From 98e46ff7cc824fac03f0a58524e2ca9c408afa42 Mon Sep 17 00:00:00 2001
From: Erik Brakkee <erik@brakkee.org>
Date: Sat, 20 Jul 2024 19:37:17 +0200
Subject: [PATCH] moved websocket utils to separate package to allow use by the
 agent.

---
 cmd/agent/agent.go            |  2 +-
 cmd/converge/server.go        | 40 ++++++-----------------------------
 pkg/converge/admin.go         |  5 +++--
 pkg/websocketutil/services.go | 34 +++++++++++++++++++++++++++++
 4 files changed, 45 insertions(+), 36 deletions(-)
 create mode 100644 pkg/websocketutil/services.go

diff --git a/cmd/agent/agent.go b/cmd/agent/agent.go
index b98637b..dbb2fa6 100755
--- a/cmd/agent/agent.go
+++ b/cmd/agent/agent.go
@@ -171,7 +171,7 @@ func main() {
 	wsConn := iowrappers.NewWebSocketConn(conn)
 	defer wsConn.Close()
 
-	listener, err := yamux.Client(wsConn, nil)
+	listener, err := yamux.Server(wsConn, nil)
 	if err != nil {
 		panic(err)
 	}
diff --git a/cmd/converge/server.go b/cmd/converge/server.go
index 5908ef7..ed0524a 100644
--- a/cmd/converge/server.go
+++ b/cmd/converge/server.go
@@ -3,39 +3,13 @@ package main
 import (
 	"cidebug/pkg/converge"
 	"cidebug/pkg/iowrappers"
+	"cidebug/pkg/websocketutil"
 	"fmt"
-	"github.com/gorilla/websocket"
 	"log"
 	"net/http"
 	"regexp"
 )
 
-var upgrader = websocket.Upgrader{
-	ReadBufferSize:  1024,
-	WriteBufferSize: 1024,
-}
-
-func handleWebSocket(w http.ResponseWriter, r *http.Request,
-	handler func(w http.ResponseWriter, r *http.Request, websockerConnection iowrappers.ReadWriteAddrCloser)) {
-	conn, err := upgrader.Upgrade(w, r, nil)
-	if err != nil {
-		log.Println("Error upgrading to WebSocket:", err)
-		return
-	}
-	wsConn := iowrappers.NewWebSocketConn(conn)
-	defer wsConn.Close()
-
-	handler(w, r, wsConn)
-}
-
-type WebSocketService struct {
-	handler func(w http.ResponseWriter, r *http.Request, conn iowrappers.ReadWriteAddrCloser)
-}
-
-func (endpoint *WebSocketService) handle(w http.ResponseWriter, r *http.Request) {
-	handleWebSocket(w, r, endpoint.handler)
-}
-
 func parsePublicId(path string) (publicId string, _ error) {
 	pattern := regexp.MustCompile("^/[^/]+/([^/]+)$")
 	matches := pattern.FindStringSubmatch(path)
@@ -48,8 +22,8 @@ func parsePublicId(path string) (publicId string, _ error) {
 func main() {
 
 	admin := converge.NewAdmin()
-	registrationService := WebSocketService{
-		handler: func(w http.ResponseWriter, r *http.Request, conn iowrappers.ReadWriteAddrCloser) {
+	registrationService := websocketutil.WebSocketService{
+		Handler: func(w http.ResponseWriter, r *http.Request, conn iowrappers.ReadWriteAddrCloser) {
 			publicId, err := parsePublicId(r.URL.Path)
 			if err != nil {
 				log.Printf("Cannot parse public id from url: '%v'\n", err)
@@ -62,8 +36,8 @@ func main() {
 			}
 		},
 	}
-	clientService := WebSocketService{
-		handler: func(w http.ResponseWriter, r *http.Request, conn iowrappers.ReadWriteAddrCloser) {
+	clientService := websocketutil.WebSocketService{
+		Handler: func(w http.ResponseWriter, r *http.Request, conn iowrappers.ReadWriteAddrCloser) {
 			publicId, err := parsePublicId(r.URL.Path)
 			if err != nil {
 				log.Printf("Cannot parse public id from url: '%v'\n", err)
@@ -77,8 +51,8 @@ func main() {
 		},
 	}
 
-	http.HandleFunc("/agent/", registrationService.handle)
-	http.HandleFunc("/client/", clientService.handle)
+	http.HandleFunc("/agent/", registrationService.Handle)
+	http.HandleFunc("/client/", clientService.Handle)
 
 	// Start HTTP server
 	fmt.Println("Rendez-vous server listening on :8000")
diff --git a/pkg/converge/admin.go b/pkg/converge/admin.go
index 2f7e058..a4241ec 100644
--- a/pkg/converge/admin.go
+++ b/pkg/converge/admin.go
@@ -85,11 +85,12 @@ func (admin *Admin) addAgent(publicId string, conn io.ReadWriteCloser) (*Agent,
 	if agent != nil {
 		return nil, fmt.Errorf("A different agent with same publicId '%s' already registered", publicId)
 	}
-	session, err := yamux.Client(conn, nil)
+
+	clientSession, err := yamux.Client(conn, nil)
 	if err != nil {
 		return nil, err
 	}
-	agent = NewAgent(publicId, session)
+	agent = NewAgent(publicId, clientSession)
 	admin.agents[publicId] = agent
 	admin.logStatus()
 	return agent, nil
diff --git a/pkg/websocketutil/services.go b/pkg/websocketutil/services.go
new file mode 100644
index 0000000..3fe5162
--- /dev/null
+++ b/pkg/websocketutil/services.go
@@ -0,0 +1,34 @@
+package websocketutil
+
+import (
+	"cidebug/pkg/iowrappers"
+	"github.com/gorilla/websocket"
+	"log"
+	"net/http"
+)
+
+var upgrader = websocket.Upgrader{
+	ReadBufferSize:  1024,
+	WriteBufferSize: 1024,
+}
+
+func handleWebSocket(w http.ResponseWriter, r *http.Request,
+	handler func(w http.ResponseWriter, r *http.Request, websockerConnection iowrappers.ReadWriteAddrCloser)) {
+	conn, err := upgrader.Upgrade(w, r, nil)
+	if err != nil {
+		log.Println("Error upgrading to WebSocket:", err)
+		return
+	}
+	wsConn := iowrappers.NewWebSocketConn(conn)
+	defer wsConn.Close()
+
+	handler(w, r, wsConn)
+}
+
+type WebSocketService struct {
+	Handler func(w http.ResponseWriter, r *http.Request, conn iowrappers.ReadWriteAddrCloser)
+}
+
+func (endpoint *WebSocketService) Handle(w http.ResponseWriter, r *http.Request) {
+	handleWebSocket(w, r, endpoint.Handler)
+}