Fix memory leak

internal/api/routes.go

@@ -6,22 +6,25 @@ import (
 
 	"github.com/go-chi/chi/v5"
 	"github.com/gorilla/websocket"
+	"github.com/rs/zerolog/hlog"
 
 	"github.com/beeper/registration-relay/internal/provider"
 )
 
 var upgrader = websocket.Upgrader{}
 
 func (a *api) bridgeExecuteCommand(w http.ResponseWriter, r *http.Request) {
-	providerCode, ok := strings.CutPrefix(r.Header.Get("Authorization"), "Bearer ")
+	code, ok := strings.CutPrefix(r.Header.Get("Authorization"), "Bearer ")
 	if !ok {
 		w.WriteHeader(http.StatusBadRequest)
 		return
 	}
 
-	provider, exists := provider.GetProvider(providerCode)
+	log := hlog.FromRequest(r).With().Str("code", code).Logger()
+
+	provider, exists := provider.GetProvider(code)
 	if !exists {
-		a.log.Warn().Str("provider_code", providerCode).Msg("No provider found for code")
+		log.Warn().Msg("No provider found for code")
 		w.WriteHeader(http.StatusNotFound)
 		return
 	}
@@ -30,6 +33,7 @@ func (a *api) bridgeExecuteCommand(w http.ResponseWriter, r *http.Request) {
 
 	resp, err := provider.ExecuteCommand(command)
 	if err != nil {
+		log.Err(err).Msg("Failed to execute command")
 		w.WriteHeader(http.StatusInternalServerError)
 		return
 	}
@@ -39,9 +43,11 @@ func (a *api) bridgeExecuteCommand(w http.ResponseWriter, r *http.Request) {
 }
 
 func (a *api) providerWebsocket(w http.ResponseWriter, r *http.Request) {
+	log := hlog.FromRequest(r)
+
 	conn, err := upgrader.Upgrade(w, r, nil)
 	if err != nil {
-		a.log.Err(err).Msg("Failed to upgrade websocket connection")
+		log.Err(err).Msg("Failed to upgrade websocket connection")
 		w.WriteHeader(http.StatusInternalServerError)
 		return
 	}
@@ -50,5 +56,5 @@ func (a *api) providerWebsocket(w http.ResponseWriter, r *http.Request) {
 	provider := provider.NewProvider(conn, a.secret)
 	provider.WebsocketLoop()
 
-	a.log.Info().Msg("Websocket connection closed")
+	log.Info().Msg("Websocket connection closed")
 }

internal/provider/provider.go

@@ -5,8 +5,10 @@ import (
 	"crypto/sha256"
 	"encoding/base64"
 	"encoding/json"
+	"errors"
 	"fmt"
 	"sync"
+	"time"
 
 	"github.com/gorilla/websocket"
 	"github.com/rs/zerolog"
@@ -19,6 +21,8 @@ import (
 // Map codes -> providers, management
 //
 
+const providerRequestTimeout = 60 * time.Second
+
 var (
 	codeToProvider     map[string]*provider
 	codeToProviderLock sync.Mutex
@@ -72,10 +76,15 @@ func RegisterProvider(data registerCommandData, provider *provider) (*registerCo
 	return &data, nil
 }
 
-func UnregisterProvider(key string) {
+func UnregisterProvider(key string, provider *provider) {
 	codeToProviderLock.Lock()
 	defer codeToProviderLock.Unlock()
-	delete(codeToProvider, key)
+
+	// Only delete the provider if it's the one we're trying to unregister. If we force kick a provider
+	// due to code conflict the map key has already been overwritten.
+	if existing, exists := codeToProvider[key]; exists && existing == provider {
+		delete(codeToProvider, key)
+	}
 }
 
 // Actual provider implementation
@@ -143,6 +152,10 @@ Loop:
 				}
 				break Loop
 			}
+			// Make sure we unregister at the end of the loop
+			registerCode = response.Code
+
+			p.log = p.log.With().Str("code", request.Code).Logger()
 			p.log.Debug().Msg("Registered provider")
 
 			// Send back register response before setting the flag (ws is single writer)
@@ -176,7 +189,7 @@ Loop:
 
 	p.log.Info().Msg("Exit provider websocket loop")
 	if registerCode != "" {
-		UnregisterProvider(registerCode)
+		UnregisterProvider(registerCode, p)
 		p.log.Debug().Msg("Unregistered provider")
 	}
 }
@@ -198,10 +211,21 @@ func (p *provider) ExecuteCommand(command string) (json.RawMessage, error) {
 	}
 
 	// Send over our command and listen for the result
+	p.log.Debug().
+		Int("req_id", p.reqID).
+		Str("command", command).
+		Msg("Send request")
 	go p.ws.WriteMessage(websocket.TextMessage, buf)
-	result := <-p.resultsCh
-
-	p.log.Info().Bytes("TEST", result).Msg("")
 
-	return result, nil
+	select {
+	case result := <-p.resultsCh:
+		p.log.Debug().
+			Int("req_id", p.reqID).
+			Msg("Received response")
+		return result, nil
+	case <-time.After(providerRequestTimeout):
+		// If the request times out shoot the underlying websocket connection to force a reconnect
+		p.ws.Close()
+		return nil, errors.New("request timed out")
+	}
 }