diff --git a/internal/common/recovery/recovery.go b/internal/common/recovery/recovery.go new file mode 100644 index 0000000..000bd90 --- /dev/null +++ b/internal/common/recovery/recovery.go @@ -0,0 +1,20 @@ +package recovery + +import ( + "github.com/twofas/2fas-server/internal/common/logging" +) + +func DoNotPanic(fn func()) { + defer func() { + if err := recover(); err != nil { + stack := stack(3) + + logging.WithFields(logging.Fields{ + "stack": string(stack), + "error": err, + }).Error("Panic") + } + }() + + fn() +} diff --git a/internal/websocket/common/handler.go b/internal/websocket/common/handler.go index dcf5d20..a3178a3 100644 --- a/internal/websocket/common/handler.go +++ b/internal/websocket/common/handler.go @@ -4,6 +4,7 @@ import ( "github.com/gin-gonic/gin" "github.com/gorilla/websocket" "github.com/twofas/2fas-server/internal/common/logging" + "github.com/twofas/2fas-server/internal/common/recovery" "net/http" "os" "time" @@ -72,8 +73,13 @@ func (h *ConnectionHandler) serveWs(hub *Hub, w http.ResponseWriter, r *http.Req client := &Client{hub: hub, conn: conn, send: make(chan []byte, 256)} client.hub.register <- client - go client.writePump() - go client.readPump() + go recovery.DoNotPanic(func() { + client.writePump() + }) + + go recovery.DoNotPanic(func() { + client.readPump() + }) go func() { disconnectAfter := 3 * time.Minute