go-socket/wsServer.go

package main

import (
	"context"
	"encoding/json"
	"errors"
	"log"
	"net/http"
	"time"

	"go-socket/Enums/WsMessageFrom"

	"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 sendMessageStructCloseIfTimeout(user *User, message *Message) {
	if user.WsConn == nil {
		return
	}

	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()

	jsonMsg, err := json.Marshal(message)
	if err != nil {
		log.Printf("json marshal error: %v", err)
		return
	}
	err = wsjson.Write(ctx, user.WsConn, jsonMsg)
	if err != nil {
		log.Printf("json write error: %v", err)
		return
	}
}

func sendMessageCloseIfTimeout(user *User, message *map[string]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 {
		sendMessageCloseIfTimeout(user, message)
	}
}

func WsMessageSendToUser(to *User, message *Message) {
	sendMessageStructCloseIfTimeout(to, message)
}

func WsSendToGroup(group *Group, sender *User, message string) error {
	for groupUserId := range group.Users {
		groupUser, err := CacheGetUserById(groupUserId)
		if err != nil || groupUser.Id == sender.Id {
			continue
		}

		var msg = map[string]any{
			"type":    WsMessageFrom.Group,
			"from":    group.Id,
			"sender":  sender.Id,
			"content": message,
		}
		sendMessageCloseIfTimeout(groupUser, &msg)
	}
	return nil
}

func handleAuthenticatedMessage(user *User, userMessage *map[string]any) bool {
	sendMessageCloseIfTimeout(user, userMessage)
	return true
}

func handleUnauthenticatedMessage(ctx context.Context, user *User, userMessage *map[string]any) bool {
	token, ok := (*userMessage)["token"].(string)
	if !ok {
		var msg = map[string]any{
			"type":  WsMessageFrom.Server,
			"error": "no token in message",
		}
		sendMessageCloseIfTimeout(user, &msg)
		return false
	}

	userId, err := TokenValidateGetId(token)
	if err != nil {
		var msg = map[string]any{
			"type":  WsMessageFrom.Server,
			"error": "invalid token",
		}
		sendMessageCloseIfTimeout(user, &msg)
		return false
	}

	userFromCache, err := CacheGetUserById(userId)
	if err != nil {
		var msg = map[string]any{
			"type":  WsMessageFrom.Server,
			"error": "user not found",
		}
		sendMessageCloseIfTimeout(user, &msg)
		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 {
				var msg = map[string]any{
					"type":  "server",
					"error": "invalid user data",
				}
				sendMessageCloseIfTimeout(user, &msg)
				return false
			}
			err = DbGroupGetMembers(ctx, dbGroup)
			if err != nil {
				var msg = map[string]any{
					"type":  "server",
					"error": "invalid user data",
				}
				sendMessageCloseIfTimeout(user, &msg)
				return false
			}
			CacheSaveGroup(dbGroup)
		}
	}
	return true
}

func closeConnection(user *User, ignoreCache bool) {
	if !ignoreCache {
		CacheDeleteUser(user.Id)
	}
	if user.WsConn != nil {
		user.WsConn.CloseNow()
	}
}