clean up challenge if client aborts (#17)

docker-compose.yaml

@@ -15,28 +15,29 @@ services:
             max-size: 10m
         ports:
           - '3000:3000'
-        image: fd49561087a8563fece3be7eff59f6ff728c3e749fbd4ebfaee96a6ee3982b7d
-        container_name: ssh-sync-server
+        image: 856da056de3abb7b317e59481bd500c27eb0f3b2a4c2432e8db930c1d20e71bf
+        container_name: ssh-sync-server-debug
     ssh-sync-db:
         image: therealpaulgg/ssh-sync-db:latest 
-        container_name: ssh-sync-db
+        container_name: ssh-sync-db-debug
         environment:
           - POSTGRES_USER=sshsync
           - POSTGRES_PASSWORD=sshsync
           - POSTGRES_DB=sshsync
         restart: always
     ssh-sync:
-        image: 46204e8109ce
+        image: 9065faaa7a20a821f7323f42cbddac1b594f7d01d57f8b2a2837e433769b86f4
         container_name: ssh-sync
         stdin_open: true # Allows Docker container to keep STDIN open
         tty: true # Allocates a pseudo-TTY
     ssh-sync-2:
-        image: 46204e8109ce
+        image: 9065faaa7a20a821f7323f42cbddac1b594f7d01d57f8b2a2837e433769b86f4
         container_name: ssh-sync-2
         stdin_open: true # Allows Docker container to keep STDIN open
         tty: true # Allocates a pseudo-TTY
     ssh-sync-3:
-        image: 46204e8109ce
+        image: 9065faaa7a20a821f7323f42cbddac1b594f7d01d57f8b2a2837e433769b86f4
         container_name: ssh-sync-3
         stdin_open: true # Allows Docker container to keep STDIN open
-        tty: true # Allocates a pseudo-TTY
+        tty: true # Allocates a pseudo-TTY
+        #http://ssh-sync-server-debug:3000

pkg/web/live/main.go

@@ -3,6 +3,7 @@ package live
 import (
 	"database/sql"
 	"errors"
+	"io"
 	"net"
 	"net/http"
 	"strings"
@@ -78,6 +79,7 @@ func MachineChallengeResponse(i *do.Injector, r *http.Request, w http.ResponseWr
 func MachineChallengeResponseHandler(i *do.Injector, r *http.Request, w http.ResponseWriter, c *net.Conn) {
 	conn := *c
 	defer conn.Close()
+
 	user, ok := r.Context().Value(context_keys.UserContextKey).(*models.User)
 	if !ok {
 		log.Warn().Msg("Could not get user from context")
@@ -132,6 +134,24 @@ func NewMachineChallenge(i *do.Injector, r *http.Request, w http.ResponseWriter)
 func NewMachineChallengeHandler(i *do.Injector, r *http.Request, w http.ResponseWriter, c *net.Conn) {
 	conn := *c
 	defer conn.Close()
+	closeChan := make(chan struct{}) // Channel to signal connection closure
+
+	// Start a goroutine to monitor connection for closure
+	go func() {
+		buf := make([]byte, 1)
+
+		for {
+			_, err := conn.Read(buf)
+			if err != nil {
+				if err == io.EOF {
+					close(closeChan)
+				} else {
+					close(closeChan)
+				}
+				return
+			}
+		}
+	}()
 	// first message sent should be JSON payload
 	userMachine, err := utils.ReadClientMessage[dto.UserMachineDto](&conn)
 	if err != nil {
@@ -235,6 +255,15 @@ func NewMachineChallengeHandler(i *do.Injector, r *http.Request, w http.Response
 					timer.Stop()
 				}
 				return
+			case <-closeChan:
+				log.Debug().Msg("Connection closed by client")
+				ChallengeResponseDict.mux.Lock()
+				// Check if the challenge still exists before sending to the channel
+				if _, exists := ChallengeResponseDict.dict[challengePhrase]; exists {
+					ChallengeResponseDict.dict[challengePhrase].ChallengeAccepted <- false
+				}
+				ChallengeResponseDict.mux.Unlock()
+				return
 			}
 		}
 	}()