package main import ( "context" "errors" "log" "net/http" "time" "github.com/coder/websocket" "github.com/coder/websocket/wsjson" ) func ServeConnection(responseWriter http.ResponseWriter, request *http.Request) { connection, err := websocket.Accept(responseWriter, request, nil) if err != nil { log.Printf("websocket accept error: %v", err) return } ctx, cancel := context.WithTimeout(context.Background(), time.Second*30) defer cancel() client := Client{} for { var clientMessage map[string]any err := wsjson.Read(ctx, connection, &clientMessage) if err != nil { log.Printf("read error: %clientMessage", err) return } if len(clientMessage) > 0 { if client.IsAuthenticated { handleAuthenticatedMessage() } else { if !handleUnauthenticatedMessage(connection, &client, &clientMessage) { closeConnection(connection) return } } } } } func sendMessageCloseIfTimeout(conn *websocket.Conn, message *map[string]any) { ctx, cancel := context.WithTimeout(context.Background(), time.Second*10) defer cancel() err := wsjson.Write(ctx, conn, message) if err != nil { if errors.Is(err, context.DeadlineExceeded) { closeConnection(conn) } log.Printf("write error: %v", err) } } func handleUnauthenticatedMessage(conn *websocket.Conn, client *Client, message *map[string]any) bool { token, ok := (*message)["token"].(string) if !ok { var errmsg = map[string]any{ "type": WSServerResponse(BadMessage), "message": "token required", } sendMessageCloseIfTimeout(conn, &errmsg) return false } err := SetClientFromToken(client, token) if err != nil { var errmsg = map[string]any{ "type": WSServerResponse(InvalidCredentials), "message": "bad token", } sendMessageCloseIfTimeout(conn, &errmsg) return false } client.IsAuthenticated = true return true } func handleAuthenticatedMessage() { } func closeConnection(conn *websocket.Conn) { conn.Close(websocket.StatusNormalClosure, "closing connection") }