From e31ab7176b03f1cd8e4653b959a48f284624c344 Mon Sep 17 00:00:00 2001
From: Erik Brakkee <erik@brakkee.org>
Date: Sat, 24 Aug 2024 23:06:02 +0200
Subject: [PATCH] wsproxy mode tested.

---
 pkg/server/matchmaker/matchmaker_test.go | 80 +++++++++++++++++++++---
 1 file changed, 72 insertions(+), 8 deletions(-)

diff --git a/pkg/server/matchmaker/matchmaker_test.go b/pkg/server/matchmaker/matchmaker_test.go
index da3261e..4ba3f42 100644
--- a/pkg/server/matchmaker/matchmaker_test.go
+++ b/pkg/server/matchmaker/matchmaker_test.go
@@ -11,6 +11,8 @@ import (
 	"io"
 	"log"
 	"net/http"
+	"os"
+	"strings"
 	"sync"
 	"testing"
 	"time"
@@ -111,6 +113,10 @@ type TestClient struct {
 	serverSIdeConn iowrappers.ReadWriteAddrCloser
 	publicId       models.RendezVousId
 	clientId       models.ClientId
+
+	// for wsproxy mode
+	serverProtocol       comms.ProtocolVersion
+	clientConnectionInfo comms.ClientConnectionInfo
 }
 
 func NewTestClient(ctx context.Context) *TestClient {
@@ -127,6 +133,26 @@ func (c *TestClient) Disconnect() {
 	c.clientSideConn.Close()
 }
 
+func (c *TestClient) WsproxyInit() error {
+	channel := comms.NewGOBChannel(c.clientSideConn)
+	serverProtocol, err := comms.ReceiveWithTimeout[comms.ProtocolVersion](channel)
+	if err != nil {
+		return err
+	}
+	c.serverProtocol = serverProtocol
+
+	clientConnectionInfo, err := comms.ReceiveWithTimeout[comms.ClientConnectionInfo](channel)
+	if err != nil {
+		return err
+	}
+	c.clientConnectionInfo = clientConnectionInfo
+	err = comms.SendWithTimeout(channel, comms.NewEnvironmentInfo(os.Getenv("SHELL")))
+	if err != nil {
+		return err
+	}
+	return nil
+}
+
 func (s *MatchMakerTestSuite) Test_newMatchMaker() {
 	s.checkState(0, 0)
 }
@@ -152,7 +178,7 @@ func (s *MatchMakerTestSuite) Test_singleAgentAndClient() {
 	waitForAgentFunc := s.registerAgent(publicId, agent)
 	go waitForAgentFunc()
 
-	client, err := s.connectClient(publicId)
+	client, err := s.connectClient(publicId, false)
 	s.Nil(err)
 
 	s.checkState(1, 1)
@@ -174,9 +200,8 @@ func (s *MatchMakerTestSuite) Test_singleAgentAndClient() {
 func (s *MatchMakerTestSuite) Test_ConnectCLientToUnknownAgent() {
 	publicId := models.RendezVousId("abc")
 
-	client, err := s.connectClient(publicId)
+	_, err := s.connectClient(publicId, false)
 	s.NotNil(err)
-	s.Nil(client)
 }
 
 func (s *MatchMakerTestSuite) Test_multipleAgentsAndClients() {
@@ -194,7 +219,7 @@ func (s *MatchMakerTestSuite) Test_multipleAgentsAndClients() {
 	testClients := make([]*TestClient, 0)
 	for publicId, nclients := range clients {
 		for range nclients {
-			client, err := s.connectClient(models.RendezVousId(publicId))
+			client, err := s.connectClient(models.RendezVousId(publicId), false)
 			s.Require().Nil(err)
 			testClients = append(testClients, client)
 		}
@@ -220,14 +245,14 @@ func (s *MatchMakerTestSuite) Test_multipleAgentsAndClients() {
 	wg.Wait()
 }
 
-func (s *MatchMakerTestSuite) connectClient(publicId models.RendezVousId) (*TestClient, error) {
+func (s *MatchMakerTestSuite) connectClient(publicId models.RendezVousId, wsproxyMode bool) (*TestClient, error) {
 	client := NewTestClient(s.ctx)
 	var clientId models.ClientId
 	res := testsupport.RunAndWait(
 		&s.Suite,
 		func() any {
 			//server
-			clientIdCreated, synchronizer, err := s.matchMaker.Connect(false, publicId, client.serverSIdeConn)
+			clientIdCreated, synchronizer, err := s.matchMaker.Connect(wsproxyMode, publicId, client.serverSIdeConn)
 			clientId = clientIdCreated
 			if err == nil {
 				log.Println("test: synchronizing streams.")
@@ -236,12 +261,14 @@ func (s *MatchMakerTestSuite) connectClient(publicId models.RendezVousId) (*Test
 			return err
 		},
 		func() any {
-			// client, nothing to do with wsproxy mode off.
+			if wsproxyMode {
+				client.WsproxyInit()
+			}
 			return nil
 		})
 
 	if res[0] != nil {
-		return nil, res[0].(error)
+		return client, res[0].(error)
 	}
 	client.publicId = publicId
 	client.clientId = clientId
@@ -287,3 +314,40 @@ func (s *MatchMakerTestSuite) registerAgent(publicId models.RendezVousId, agent
 	}
 	return res[0].(WaitForAgentFunc)
 }
+
+func (s *MatchMakerTestSuite) Test_ConnectWsproxyMode() {
+	publicId := models.RendezVousId("abc")
+	agent := NewTestAgent(s.ctx)
+
+	waitForAgentFunc := s.registerAgent(publicId, agent)
+	go waitForAgentFunc()
+
+	client, err := s.connectClient(publicId, true)
+	s.Nil(err)
+
+	s.checkState(1, 1)
+
+	agentClientSideConn, err := agent.listener.GetConnection(string(client.clientId))
+	log.Printf("Agent side conn %v", agentClientSideConn)
+	s.Nil(err)
+	testsupport.BidirectionalConnectionCheck(
+		&s.Suite, "testmsg",
+		client.clientSideConn,
+		agentClientSideConn)
+
+	s.True(client.clientConnectionInfo.Ok)
+
+	client.Disconnect()
+	// It is the agents choice to exit> The test agent does not exit by default when
+	// there are no more connections.
+	s.checkState(1, 0)
+}
+
+func (s *MatchMakerTestSuite) Test_ConnectWsproxyModeAgentNotFound() {
+	publicId := models.RendezVousId("abc")
+
+	client, err := s.connectClient(publicId, true)
+	s.NotNil(err)
+	s.True(strings.Contains(err.Error(), "No agent found for rendez-vous id"))
+	s.False(client.clientConnectionInfo.Ok)
+}