104 lines
2.3 KiB
Go
104 lines
2.3 KiB
Go
package main
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"log"
|
|
"net/http"
|
|
"sync"
|
|
"time"
|
|
|
|
"github.com/coder/websocket"
|
|
"github.com/coder/websocket/wsjson"
|
|
)
|
|
|
|
type wsServer struct {
|
|
OnOpen func(c *Client)
|
|
OnClose func(c *Client, err error)
|
|
OnMessage func(c *Client, msg map[string]any)
|
|
}
|
|
|
|
var (
|
|
mu sync.Mutex
|
|
)
|
|
|
|
func (s *wsServer) ServeHTTP(responseWriter http.ResponseWriter, request *http.Request) {
|
|
conn, err := websocket.Accept(responseWriter, request, &websocket.AcceptOptions{
|
|
InsecureSkipVerify: true,
|
|
})
|
|
if err != nil {
|
|
log.Println("accept error:", err)
|
|
return
|
|
}
|
|
defer conn.CloseNow()
|
|
|
|
client := &Client{conn: conn}
|
|
|
|
ctx, cancel := context.WithCancel(context.Background())
|
|
defer cancel()
|
|
|
|
if s.OnOpen != nil {
|
|
s.OnOpen(client)
|
|
}
|
|
|
|
var readErr error
|
|
for {
|
|
var msg map[string]any
|
|
if readErr = wsjson.Read(ctx, conn, &msg); readErr != nil {
|
|
break
|
|
}
|
|
if s.OnMessage != nil {
|
|
s.OnMessage(client, msg)
|
|
}
|
|
}
|
|
|
|
cancel() // cancel before OnClose so any in-flight queries are canceled first
|
|
|
|
if s.OnClose != nil {
|
|
s.OnClose(client, readErr)
|
|
}
|
|
|
|
conn.Close(websocket.StatusNormalClosure, "done")
|
|
}
|
|
|
|
func sendAndCloseIfFails(conn *websocket.Conn, message *map[string]any) {
|
|
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
|
defer cancel()
|
|
if err := wsjson.Write(ctx, conn, message); err != nil {
|
|
conn.Close(websocket.StatusGoingAway, "Write error")
|
|
}
|
|
}
|
|
func sendToGroup(id uint32, excludedUserId uint32, message *map[string]any) error {
|
|
if _, ok := Groups[id]; !ok {
|
|
return errors.New("Group Not Found")
|
|
}
|
|
|
|
for client := range ConnectedClients[id] {
|
|
if client.User.Id != excludedUserId {
|
|
sendAndCloseIfFails(client.conn, message)
|
|
}
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func handleUnauthenticatedMessage(client *Client, msg map[string]any) {
|
|
token := msg["token"].(string)
|
|
user, err := GetUserFromToken(token)
|
|
if err != nil {
|
|
client.conn.Close(websocket.StatusPolicyViolation, "invalid token")
|
|
return
|
|
}
|
|
client.User = &user
|
|
m := map[string]any{
|
|
"authAs": user.Name,
|
|
}
|
|
sendAndCloseIfFails(client.conn, &m)
|
|
AddOrUpdateConnectedClientToCache(&mu, client)
|
|
log.Println("New User authenticated as: " + user.Name)
|
|
}
|
|
|
|
func handleAuthenticatedMessage(client *Client, msg map[string]any) {
|
|
m := map[string]any{"temporary": "unauthorized"}
|
|
sendAndCloseIfFails(client.conn, &m)
|
|
}
|