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

@@ -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)
 }

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)
+}

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