From 3b2b9162ba6014c7487b9fd7efbadd1aa94f2678 Mon Sep 17 00:00:00 2001
From: Erik Brakkee <erik@brakkee.org>
Date: Fri, 19 Jul 2024 21:48:54 +0200
Subject: [PATCH] cleanup now works. Clients and agents can connect and
 disconnect at will.

---
 cmd/converge/server.go | 47 ++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 47 insertions(+)

diff --git a/cmd/converge/server.go b/cmd/converge/server.go
index 0ee815f..41b2dba 100644
--- a/cmd/converge/server.go
+++ b/cmd/converge/server.go
@@ -105,12 +105,48 @@ func (admin *Admin) addClient(publicId string, conn *iowrappers.WebSocketConn) (
 	return agent, nil
 }
 
+func (admin *Admin) RemoveAgent(publicId string) error {
+	admin.mutex.Lock()
+	defer admin.mutex.Unlock()
+
+	agent := admin.agents[publicId]
+	if agent == nil {
+		return fmt.Errorf("Cannot remove agent: '%s' not found", publicId)
+	}
+	agent.agent = nil
+	if agent.client == nil {
+		log.Printf("Removing agent: '%s'", publicId)
+		delete(admin.agents, publicId)
+	}
+	return nil
+}
+
+func (admin *Admin) RemoveClient(publicId string) error {
+	admin.mutex.Lock()
+	defer admin.mutex.Unlock()
+
+	agent := admin.agents[publicId]
+	if agent == nil {
+		return fmt.Errorf("Cannot remove agent: '%s' not found", publicId)
+	}
+	agent.client = nil
+	if agent.agent == nil {
+		log.Printf("Removing client: '%s'", publicId)
+		delete(admin.agents, publicId)
+	}
+	return nil
+}
+
 func (admin *Admin) Register(publicId string, conn *iowrappers.WebSocketConn) error {
 	defer conn.Close()
 	agent, err := admin.addAgent(publicId, conn)
 	if err != nil {
 		return err
 	}
+	defer func() {
+		admin.RemoveAgent(publicId)
+	}()
+	log.Printf("After defer remove agent\n")
 	agent.agentAvailable <- true
 	log.Printf("Agent registered: '%s'\n", publicId)
 	<-agent.clientClosed
@@ -123,6 +159,9 @@ func (admin *Admin) Connect(publicId string, conn *iowrappers.WebSocketConn) err
 	if err != nil {
 		return err
 	}
+	defer func() {
+		admin.RemoveClient(publicId)
+	}()
 	<-agent.agentAvailable
 	log.Printf("Connecting client and agent: '%s'\n", publicId)
 	iowrappers.SynchronizeStreams(agent.client, agent.agent)
@@ -130,6 +169,14 @@ func (admin *Admin) Connect(publicId string, conn *iowrappers.WebSocketConn) err
 	return nil
 }
 
+func (admin *Admin) log() {
+	log.Println("CONNECTIONS")
+	for _, agent := range admin.agents {
+		log.Println(agent.publicId)
+	}
+	log.Printf("\n")
+}
+
 func parsePublicId(path string) (publicId string, _ error) {
 	pattern := regexp.MustCompile("^/[^/]+/([^/]+)$")
 	matches := pattern.FindStringSubmatch(path)