From c7ba566c4839ee3b008d7b923ba6c5d312cfaf66 Mon Sep 17 00:00:00 2001
From: LittleSheep <littlesheep.code@hotmail.com>
Date: Mon, 27 May 2024 23:00:49 +0800
Subject: [PATCH] :ambulance: Fix write map panic with mutex

---
 pkg/server/ws.go            |  4 ++++
 pkg/services/auth.go        | 17 +++++++++++++++--
 pkg/services/connections.go | 11 ++++++++++-
 settings.toml               |  2 +-
 4 files changed, 30 insertions(+), 4 deletions(-)

diff --git a/pkg/server/ws.go b/pkg/server/ws.go
index 0c18af9..5199809 100644
--- a/pkg/server/ws.go
+++ b/pkg/server/ws.go
@@ -2,10 +2,12 @@ package server
 
 import (
 	"fmt"
+
 	"git.solsynth.dev/hydrogen/passport/pkg/models"
 	"git.solsynth.dev/hydrogen/passport/pkg/services"
 	"github.com/gofiber/contrib/websocket"
 	jsoniter "github.com/json-iterator/go"
+	"github.com/rs/zerolog/log"
 	"github.com/samber/lo"
 )
 
@@ -14,6 +16,7 @@ func listenWebsocket(c *websocket.Conn) {
 
 	// Push connection
 	services.ClientRegister(user, c)
+	log.Debug().Uint("user", user.ID).Msg("New websocket connection established...")
 
 	// Event loop
 	var task models.UnifiedCommand
@@ -76,4 +79,5 @@ func listenWebsocket(c *websocket.Conn) {
 
 	// Pop connection
 	services.ClientUnregister(user, c)
+	log.Debug().Uint("user", user.ID).Msg("A websocket connection disconnected...")
 }
diff --git a/pkg/services/auth.go b/pkg/services/auth.go
index f85967b..42b2fb1 100644
--- a/pkg/services/auth.go
+++ b/pkg/services/auth.go
@@ -2,15 +2,20 @@ package services
 
 import (
 	"fmt"
-	jsoniter "github.com/json-iterator/go"
+	"sync"
 	"time"
 
+	jsoniter "github.com/json-iterator/go"
+
 	"git.solsynth.dev/hydrogen/passport/pkg/models"
 	"github.com/gofiber/fiber/v2"
 	"github.com/rs/zerolog/log"
 )
 
-var authContextCache = make(map[string]models.AuthContext)
+var (
+	authContextMutex sync.Mutex
+	authContextCache = make(map[string]models.AuthContext)
+)
 
 func Authenticate(access, refresh string, depth int) (ctx models.AuthContext, perms map[string]any, newAccess, newRefresh string, err error) {
 	var claims PayloadClaims
@@ -50,7 +55,9 @@ func GetAuthContext(jti string) (models.AuthContext, error) {
 	if val, ok := authContextCache[jti]; ok {
 		ctx = val
 		ctx.LastUsedAt = time.Now()
+		authContextMutex.Lock()
 		authContextCache[jti] = ctx
+		authContextMutex.Unlock()
 		log.Debug().Str("jti", jti).Msg("Used an auth context cache")
 	} else {
 		ctx, err = CacheAuthContext(jti)
@@ -83,7 +90,9 @@ func CacheAuthContext(jti string) (models.AuthContext, error) {
 	}
 
 	// Put the data into memory for cache
+	authContextMutex.Lock()
 	authContextCache[jti] = ctx
+	authContextMutex.Unlock()
 
 	return ctx, nil
 }
@@ -97,7 +106,9 @@ func RecycleAuthContext() {
 	for key, val := range authContextCache {
 		if val.LastUsedAt.Add(60*time.Second).Unix() < time.Now().Unix() {
 			affected++
+			authContextMutex.Lock()
 			delete(authContextCache, key)
+			authContextMutex.Unlock()
 		}
 	}
 	log.Debug().Int("affected", affected).Msg("Recycled auth context...")
@@ -106,7 +117,9 @@ func RecycleAuthContext() {
 func InvalidAuthCacheWithUser(userId uint) {
 	for key, val := range authContextCache {
 		if val.Account.ID == userId {
+			authContextMutex.Lock()
 			delete(authContextCache, key)
+			authContextMutex.Unlock()
 		}
 	}
 }
diff --git a/pkg/services/connections.go b/pkg/services/connections.go
index 8be108e..02e84c5 100644
--- a/pkg/services/connections.go
+++ b/pkg/services/connections.go
@@ -1,22 +1,31 @@
 package services
 
 import (
+	"sync"
+
 	"git.solsynth.dev/hydrogen/passport/pkg/models"
 	"github.com/gofiber/contrib/websocket"
 )
 
-var wsConn = make(map[uint]map[*websocket.Conn]bool)
+var (
+	wsMutex sync.Mutex
+	wsConn  = make(map[uint]map[*websocket.Conn]bool)
+)
 
 func ClientRegister(user models.Account, conn *websocket.Conn) {
+	wsMutex.Lock()
 	if wsConn[user.ID] == nil {
 		wsConn[user.ID] = make(map[*websocket.Conn]bool)
 	}
 	wsConn[user.ID][conn] = true
+	wsMutex.Unlock()
 }
 
 func ClientUnregister(user models.Account, conn *websocket.Conn) {
+	wsMutex.Lock()
 	if wsConn[user.ID] == nil {
 		wsConn[user.ID] = make(map[*websocket.Conn]bool)
 	}
 	delete(wsConn[user.ID], conn)
+	wsMutex.Unlock()
 }
diff --git a/settings.toml b/settings.toml
index 9ba3c8e..cbddeb0 100644
--- a/settings.toml
+++ b/settings.toml
@@ -34,7 +34,7 @@ access_token_duration = 300
 refresh_token_duration = 2592000
 
 [database]
-dsn = "host=localhost dbname=hy_passport port=5432 sslmode=disable"
+dsn = "host=localhost user=postgres password=password dbname=hy_passport port=5432 sslmode=disable"
 prefix = "passport_"
 
 [permissions.default]