change db structure, add goroutine connection handling

database.go

@@ -19,11 +19,26 @@ func InitDatabase(ctx context.Context) {
 
 	_, err = conn.Exec(ctx, `
 		CREATE TABLE IF NOT EXISTS users (
-		    Id SERIAL PRIMARY KEY,
+		    id SERIAL PRIMARY KEY,
-		    Name VARCHAR(20) UNIQUE NOT NULL,
+		    name VARCHAR(20) UNIQUE NOT NULL,
-		    PassHash VARCHAR(60) NOT NULL,
+		    pass_hash VARCHAR(60) NOT NULL,
-		    Color VARCHAR(3) NOT NULL
+		    color VARCHAR(3) NOT NULL
-		)
+		);
+		CREATE TABLE IF NOT EXISTS chat_groups (
+		    id SERIAL PRIMARY KEY,
+		    name VARCHAR(48) NOT NULL,
+		    createor INTEGER NOT NULL REFERANCE users(id) ON DELETE SET NULL,
+		    owner INTEGER NOT NULL REFERANCE user(id) ON DELETE SET NULL
+		    enable_user_colors BOOLEAN NOT NULL DEFAULT true,
+		    group_color VARCHAR(3),
+		    created_at TIMESTAMP NOT NULL DEFAULT NOW()
+		);
+		CREATE TABLE IF NOT EXISTS chat_group_members (
+			group_id INTEGER NOT NULL REFERENCES chat_groups(id) ON DELETE CASCADE,
+			user_id  INTEGER NOT NULL REFERENCES users(id) ON DELETE CASCADE,
+			joined_at TIMESTAMP NOT NULL DEFAULT NOW(),
+			PRIMARY KEY (group_id, user_id)
+		);
 	`)
 	if err != nil {
 		panic(err)
@@ -54,16 +69,16 @@ func AddNewUser(ctx context.Context, user User) (uint32, error) {
 		return 0, errors.New("color invalid")
 	}
 	err = dbConnection.QueryRow(ctx, `
-        INSERT INTO users (Name, PassHash, Color)
+        INSERT INTO users (name, pass_hash, color)
         VALUES ($1, $2, $3)
-        RETURNING Id
+        RETURNING id
     `, user.Name, user.Password, user.Color).Scan(&id)
 	return id, err
 }
 
 func isPassValid(ctx context.Context, id uint32, plainPassword string) bool {
 	var controlHash string
-	err := dbConnection.QueryRow(ctx, "SELECT PassHash FROM users WHERE Id = $1", id).Scan(&controlHash)
+	err := dbConnection.QueryRow(ctx, "SELECT pass_hash FROM users WHERE id = $1", id).Scan(&controlHash)
 	if err != nil {
 		return false
 	}
@@ -73,7 +88,7 @@ func isPassValid(ctx context.Context, id uint32, plainPassword string) bool {
 
 func GetUserDataById(ctx context.Context, id uint32) (*User, error) {
 	var user User
-	err := dbConnection.QueryRow(ctx, "SELECT Id, Name, PassHash, Color FROM users WHERE Id = $1", id).
+	err := dbConnection.QueryRow(ctx, "SELECT id, name, pass_hash, color FROM users WHERE id = $1", id).
 		Scan(&user.Id, &user.Name, &user.Password, &user.Color)
 	if err != nil {
 		return &User{}, err
@@ -83,7 +98,7 @@ func GetUserDataById(ctx context.Context, id uint32) (*User, error) {
 }
 func GetUserDataByName(ctx context.Context, name string) (*User, error) {
 	var user User
-	err := dbConnection.QueryRow(ctx, "SELECT Id, Name, PassHash, Color FROM users WHERE Name = $1", name).
+	err := dbConnection.QueryRow(ctx, "SELECT id, name, pass_hash, color FROM users WHERE name = $1", name).
 		Scan(&user.Id, &user.Name, &user.Password, &user.Color)
 	if err != nil {
 		return &User{}, err
@@ -91,3 +106,7 @@ func GetUserDataByName(ctx context.Context, name string) (*User, error) {
 	user.IsPasswordHashed = true
 	return &user, nil
 }
+
+func CreateGroup() {
+
+}

main.go

@@ -4,32 +4,23 @@ import (
 	"context"
 	"log"
 	"net/http"
-
-	"github.com/coder/websocket"
 )
 
 func main() {
 	InitDatabase(context.Background())
 	srv := &wsServer{
-		OnOpen: func(ctx context.Context, conn *websocket.Conn) {
+		OnOpen: func(c *Client) {
 			log.Println("client connected")
-			if getConnectionDataIfAuth(conn) != nil {
-				mu.Lock()
-				unauthenticatedConnections = append(unauthenticatedConnections, conn)
-				mu.Unlock()
-			}
 		},
-		OnClose: func(ctx context.Context, conn *websocket.Conn, err error) {
+		OnClose: func(c *Client, err error) {
 			log.Println("client disconnected:", err)
-			removeConnectionCache(conn)
 		},
-		OnMessage: func(ctx context.Context, conn *websocket.Conn, msg map[string]any) {
+		OnMessage: func(c *Client, msg map[string]any) {
 			log.Printf("received: %v\n", msg)
-			authConnOrNil := getConnectionDataIfAuth(conn)
+			if c.User == nil {
-			if authConnOrNil == nil {
+				handleUnauthenticatedMessage(c, msg)
-				handleUnauthenticatedMessage(conn, msg)
 			} else {
-				handleAuthenticatedMessage(conn, msg)
+				handleAuthenticatedMessage(c, msg)
 			}
 		},
 	}

structures.go

@@ -9,8 +9,10 @@ type User struct {
 	Color            string
 	IsPasswordHashed bool
 }
-
+type Client struct {
-type AuthConnection struct {
+	conn *websocket.Conn
-	connection *websocket.Conn
+	User *User
-	user       User
+}
+
+type ChatGroup struct {
 }

wsServer.go

@@ -12,19 +12,18 @@ import (
 )
 
 type wsServer struct {
-	OnOpen    func(ctx context.Context, conn *websocket.Conn)
+	OnOpen    func(c *Client)
-	OnClose   func(ctx context.Context, conn *websocket.Conn, err error)
+	OnClose   func(c *Client, err error)
-	OnMessage func(ctx context.Context, conn *websocket.Conn, msg map[string]any)
+	OnMessage func(c *Client, msg map[string]any)
 }
 
 var (
-	unauthenticatedConnections []*websocket.Conn
+	clients []*Client
-	authenticatedConnections   []AuthConnection
 	mu      sync.Mutex
 )
 
-func (s *wsServer) ServeHTTP(w http.ResponseWriter, r *http.Request) {
+func (s *wsServer) ServeHTTP(responseWriter http.ResponseWriter, request *http.Request) {
-	conn, err := websocket.Accept(w, r, &websocket.AcceptOptions{
+	conn, err := websocket.Accept(responseWriter, request, &websocket.AcceptOptions{
 		InsecureSkipVerify: true,
 	})
 	if err != nil {
@@ -33,11 +32,16 @@ func (s *wsServer) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 	}
 	defer conn.CloseNow()
 
+	client := &Client{conn: conn}
+	mu.Lock()
+	clients = append(clients, client)
+	mu.Unlock()
+
 	ctx, cancel := context.WithCancel(context.Background())
 	defer cancel()
 
 	if s.OnOpen != nil {
-		s.OnOpen(ctx, conn)
+		s.OnOpen(client)
 	}
 
 	var readErr error
@@ -47,52 +51,29 @@ func (s *wsServer) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 			break
 		}
 		if s.OnMessage != nil {
-			s.OnMessage(ctx, conn, msg)
+			s.OnMessage(client, msg)
 		}
 	}
 
 	cancel() // cancel before OnClose so any in-flight queries are canceled first
 
 	if s.OnClose != nil {
-		s.OnClose(ctx, conn, readErr)
+		s.OnClose(client, readErr)
 	}
 
+	mu.Lock()
+	for i, c := range clients {
+		if c == client {
+			clients[i] = clients[len(clients)-1]
+			clients = clients[:len(clients)-1]
+			break
+		}
+	}
+	mu.Unlock()
+
 	conn.Close(websocket.StatusNormalClosure, "done")
 }
 
-func removeConnectionCache(conn *websocket.Conn) {
-	mu.Lock()
-	defer mu.Unlock()
-	if getConnectionDataIfAuth(conn) != nil {
-		for i, c := range authenticatedConnections {
-			if c.connection == conn {
-				authenticatedConnections[i] = authenticatedConnections[len(authenticatedConnections)-1]
-				authenticatedConnections = authenticatedConnections[:len(authenticatedConnections)-1]
-				return
-			}
-		}
-	} else {
-		for i, c := range unauthenticatedConnections {
-			if c == conn {
-				unauthenticatedConnections[i] = unauthenticatedConnections[len(unauthenticatedConnections)-1]
-				unauthenticatedConnections = unauthenticatedConnections[:len(unauthenticatedConnections)-1]
-				return
-			}
-		}
-	}
-}
-
-func getConnectionDataIfAuth(conn *websocket.Conn) *AuthConnection {
-	mu.Lock()
-	defer mu.Unlock()
-	for _, c := range authenticatedConnections {
-		if c.connection == conn {
-			return &c
-		}
-	}
-	return nil
-}
-
 func sendAndCloseIfFails(conn *websocket.Conn, message map[string]any) {
 	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
 	defer cancel()
@@ -101,54 +82,39 @@ func sendAndCloseIfFails(conn *websocket.Conn, message map[string]any) {
 	}
 }
 
-func sendToAllExceptAndCloseIfFails(conn *websocket.Conn, message map[string]any) {
+func sendToAllExceptAndCloseIfFails(client *Client, message map[string]any) {
-	_, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+	for _, other := range clients {
-	defer cancel()
+		if other != client {
-	for _, aConn := range authenticatedConnections {
+			sendAndCloseIfFails(other.conn, message)
-		if aConn.connection != conn {
-			sendAndCloseIfFails(aConn.connection, message)
 		}
 	}
 }
 
-func handleUnauthenticatedMessage(conn *websocket.Conn, msg map[string]any) {
+func handleUnauthenticatedMessage(client *Client, msg map[string]any) {
 	token := msg["token"].(string)
 	user, err := GetUserFromToken(token)
 	if err != nil {
-		log.Println("invalid or expired token:", err)
+		client.conn.Close(websocket.StatusPolicyViolation, "invalid token")
-		err := conn.Close(websocket.StatusPolicyViolation, "invalid token")
-		if err != nil {
 		return
 	}
-		return
+	client.User = &user
-	}
+	sendAndCloseIfFails(client.conn, map[string]any{
-	mu.Lock()
-	authenticatedConnections = append(authenticatedConnections, AuthConnection{connection: conn, user: user})
-	mu.Unlock()
-	sendAndCloseIfFails(conn, map[string]any{
 		"authAs": user.Name,
 	})
+	log.Println("New User authenticated as: " + user.Name)
 }
 
-func handleAuthenticatedMessage(conn *websocket.Conn, msg map[string]any) {
+func handleAuthenticatedMessage(client *Client, msg map[string]any) {
 	message := msg["message"].(string)
 	if message == "" {
-		sendAndCloseIfFails(conn, map[string]any{
+		sendAndCloseIfFails(client.conn, map[string]any{
 			"error": "no message",
 		})
 		return
 	}
 
-	auth := getConnectionDataIfAuth(conn)
+	sendToAllExceptAndCloseIfFails(client, map[string]any{
-	if auth == nil {
+		"username": client.User.Name,
-		sendAndCloseIfFails(conn, map[string]any{
-			"error": "no auth",
-		})
-		return
-	}
-
-	sendToAllExceptAndCloseIfFails(conn, map[string]any{
-		"username": auth.user.Name,
 		"message":  message,
 	})
 }