From ff9adfeb246aec8dbc214530837a2b266d8888cb Mon Sep 17 00:00:00 2001
From: Erik Brakkee <erik@brakkee.org>
Date: Wed, 24 Jul 2024 20:04:35 +0200
Subject: [PATCH] Extraction of communication setup in separate entity with
 client and related server code close together to make the setup easier to
 understand.

---
 cmd/agent/agent.go       |  7 +++----
 cmd/comms/agentserver.go | 44 ++++++++++++++++++++++++++++++++++++++++
 pkg/converge/admin.go    |  6 ++++--
 3 files changed, 51 insertions(+), 6 deletions(-)
 create mode 100644 cmd/comms/agentserver.go

diff --git a/cmd/agent/agent.go b/cmd/agent/agent.go
index 3c582e4..7dfd95b 100755
--- a/cmd/agent/agent.go
+++ b/cmd/agent/agent.go
@@ -2,6 +2,7 @@ package main
 
 import (
 	"bufio"
+	"converge/cmd/comms"
 	"converge/pkg/agent"
 	"converge/pkg/iowrappers"
 	"converge/pkg/terminal"
@@ -11,7 +12,6 @@ import (
 	"fmt"
 	"github.com/gliderlabs/ssh"
 	"github.com/gorilla/websocket"
-	"github.com/hashicorp/yamux"
 	"github.com/pkg/sftp"
 	"io"
 	"log"
@@ -178,12 +178,11 @@ func main() {
 	wsConn := websocketutil.NewWebSocketConn(conn)
 	defer wsConn.Close()
 
-	listener, err := yamux.Server(wsConn, nil)
+	commChannel, err := comms.NewCommChannel(comms.Agent, wsConn)
 	if err != nil {
 		panic(err)
 	}
 
-	// Need to create listener implementation that aactually listens for websocket connections.
 	var service AgentService
 	shells := []string{"bash", "sh", "ash", "ksh", "zsh", "fish", "tcsh", "csh"}
 	if runtime.GOOS == "windows" {
@@ -227,5 +226,5 @@ func main() {
 		strings.ReplaceAll(urlObject.Scheme, "ws", "http")+
 			"://"+urlObject.Host+"/docs/wsproxy")
 	log.Println()
-	service.Run(listener)
+	service.Run(commChannel.Session)
 }
diff --git a/cmd/comms/agentserver.go b/cmd/comms/agentserver.go
new file mode 100644
index 0000000..f3401d2
--- /dev/null
+++ b/cmd/comms/agentserver.go
@@ -0,0 +1,44 @@
+package comms
+
+import (
+	"fmt"
+	"github.com/hashicorp/yamux"
+	"io"
+	"net"
+)
+
+type CommChannel struct {
+	Peer    net.Conn
+	Session *yamux.Session
+}
+
+type Role int
+
+const (
+	Agent Role = iota
+	ConvergeServer
+)
+
+func NewCommChannel(role Role, wsConn io.ReadWriteCloser) (CommChannel, error) {
+	switch role {
+	case Agent:
+		listener, err := yamux.Server(wsConn, nil)
+		if err != nil {
+			return CommChannel{}, err
+		}
+		return CommChannel{
+			Peer:    nil,
+			Session: listener,
+		}, nil
+	case ConvergeServer:
+		clientSession, err := yamux.Client(wsConn, nil)
+		if err != nil {
+			return CommChannel{}, err
+		}
+		return CommChannel{
+			Peer:    nil,
+			Session: clientSession,
+		}, nil
+	}
+	return CommChannel{}, fmt.Errorf("Undefined role %d", role)
+}
diff --git a/pkg/converge/admin.go b/pkg/converge/admin.go
index a657e68..5ed282b 100644
--- a/pkg/converge/admin.go
+++ b/pkg/converge/admin.go
@@ -1,6 +1,7 @@
 package converge
 
 import (
+	"converge/cmd/comms"
 	"converge/pkg/iowrappers"
 	"fmt"
 	"github.com/hashicorp/yamux"
@@ -87,11 +88,12 @@ func (admin *Admin) addAgent(publicId string, conn io.ReadWriteCloser) (*Agent,
 		return nil, fmt.Errorf("A different agent with same publicId '%s' already registered", publicId)
 	}
 
-	clientSession, err := yamux.Client(conn, nil)
+	clientSession, err := comms.NewCommChannel(comms.ConvergeServer, conn)
+
 	if err != nil {
 		return nil, err
 	}
-	agent = NewAgent(publicId, clientSession)
+	agent = NewAgent(publicId, clientSession.Session)
 	admin.agents[publicId] = agent
 	admin.logStatus()
 	return agent, nil