package main import ( "context" "errors" "log" "net/http" "time" "go-socket/Enums/WsEventType" "github.com/coder/websocket" "github.com/coder/websocket/wsjson" ) func ServeWsConnection(responseWriter http.ResponseWriter, request *http.Request) { connection, err := websocket.Accept(responseWriter, request, &websocket.AcceptOptions{ InsecureSkipVerify: true, }) if err != nil { log.Printf("websocket accept error: %v", err) return } ctx, cancel := context.WithCancel(context.Background()) defer cancel() var user = User{WsConn: connection} var isAuthenticated bool var ignoreCache bool defer func() { closeConnection(&user, ignoreCache) }() for { var userMessage map[string]any err = wsjson.Read(ctx, connection, &userMessage) if err != nil { log.Printf("read error: %v", err) return } if len(userMessage) > 0 { if isAuthenticated { if !handleAuthenticatedMessage(&user, &userMessage) { return } } else { if !handleUnauthenticatedMessage(ctx, &user, &userMessage) { ignoreCache = true return } isAuthenticated = true } } } } func WsSendMessageCloseIfTimeout(user *User, message any) { if user.WsConn == nil { return } ctx, cancel := context.WithTimeout(context.Background(), time.Second*10) defer cancel() err := wsjson.Write(ctx, user.WsConn, message) if err != nil { if errors.Is(err, context.DeadlineExceeded) { closeConnection(user, false) } log.Printf("write error: %v", err) } } func sendToAllMessageCloseIfTimeout(message *map[string]any) { mu.RLock() users := make([]*User, 0, len(CacheUsers)) for _, user := range CacheUsers { users = append(users, user) } mu.RUnlock() for _, user := range users { WsSendMessageCloseIfTimeout(user, message) } } func WsSendToGroupAsUser(group *Group, sender *User, message string) error { for groupUserId := range group.Users { groupUser, err := CacheGetUserById(groupUserId) if err != nil || groupUser.Id == sender.Id { continue } // TODO update on groups rework var msg = map[string]any{ // "type": WsEventType.Group, "from": group.Id, "sender": sender.Id, "content": message, } WsSendMessageCloseIfTimeout(groupUser, &msg) } return nil } func handleAuthenticatedMessage(user *User, userMessage *map[string]any) bool { WsSendMessageCloseIfTimeout(user, userMessage) return true } func handleUnauthenticatedMessage(ctx context.Context, user *User, userMessage *map[string]any) bool { token, ok := (*userMessage)["token"].(string) response := WsEventMessage{Type: WsEventType.Authentication} if !ok { response.Event = WsAuthMessage{ Success: false, Error: "no token in message", } WsSendMessageCloseIfTimeout(user, response) return false } userId, err := TokenValidateGetId(token) if err != nil { response.Event = WsAuthMessage{ Success: false, Error: "invalid token", } WsSendMessageCloseIfTimeout(user, response) return false } userFromCache, err := CacheGetUserById(userId) if err != nil { response.Event = WsAuthMessage{ Success: false, Error: "user not found", } WsSendMessageCloseIfTimeout(user, response) return false } userFromCache.WsConn = user.WsConn *user = *userFromCache for groupId, _ := range userFromCache.Groups { _, err = CacheGetGroup(groupId) if err != nil { dbGroup := &Group{Id: groupId} err = DbGroupGetById(ctx, dbGroup) if err != nil { response.Event = WsAuthMessage{ Success: false, Error: "invalid user data", } WsSendMessageCloseIfTimeout(user, response) return false } err = DbGroupGetMembers(ctx, dbGroup) if err != nil { response.Event = WsAuthMessage{ Success: false, Error: "invalid user data", } WsSendMessageCloseIfTimeout(user, response) return false } CacheSaveGroup(dbGroup) } } response.Event = WsAuthMessage{ Success: true, Error: "", } WsSendMessageCloseIfTimeout(user, response) return true } func closeConnection(user *User, ignoreCache bool) { if !ignoreCache { CacheDeleteUser(user.Id) } if user.WsConn != nil { user.WsConn.CloseNow() } }