Changeset 0.5.0 (#35)
This commit is contained in:
279
server/events/ws.go
Normal file
279
server/events/ws.go
Normal file
@@ -0,0 +1,279 @@
|
||||
package events
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"log"
|
||||
"net/http"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/gorilla/websocket"
|
||||
)
|
||||
|
||||
var upgrader = websocket.Upgrader{
|
||||
ReadBufferSize: 1024,
|
||||
WriteBufferSize: 1024,
|
||||
CheckOrigin: func(r *http.Request) bool { return true },
|
||||
}
|
||||
|
||||
const (
|
||||
writeWait = 10 * time.Second
|
||||
pongWait = 60 * time.Second
|
||||
pingPeriod = (pongWait * 9) / 10
|
||||
maxMsgSize = 4096
|
||||
)
|
||||
|
||||
// Hub manages WebSocket connections and bridges them to the Bus.
|
||||
type Hub struct {
|
||||
bus *Bus
|
||||
mu sync.RWMutex
|
||||
conns map[string]*Conn // connID → Conn
|
||||
}
|
||||
|
||||
// Conn represents a single WebSocket connection.
|
||||
type Conn struct {
|
||||
id string
|
||||
userID string
|
||||
rooms map[string]bool // rooms this connection is subscribed to
|
||||
ws *websocket.Conn
|
||||
send chan []byte
|
||||
hub *Hub
|
||||
unsubs []func() // bus unsubscribe functions
|
||||
}
|
||||
|
||||
// NewHub creates a Hub bound to an event bus.
|
||||
func NewHub(bus *Bus) *Hub {
|
||||
return &Hub{
|
||||
bus: bus,
|
||||
conns: make(map[string]*Conn),
|
||||
}
|
||||
}
|
||||
|
||||
// HandleWebSocket is a Gin handler for WebSocket upgrade.
|
||||
// Expects userID set in context (by auth middleware or query param).
|
||||
func (h *Hub) HandleWebSocket(c *gin.Context) {
|
||||
// Auth: check token from query param
|
||||
userID := c.GetString("user_id")
|
||||
if userID == "" {
|
||||
c.JSON(http.StatusUnauthorized, gin.H{"error": "authentication required"})
|
||||
return
|
||||
}
|
||||
|
||||
ws, err := upgrader.Upgrade(c.Writer, c.Request, nil)
|
||||
if err != nil {
|
||||
log.Printf("[ws] upgrade failed: %v", err)
|
||||
return
|
||||
}
|
||||
|
||||
connID := userID + "-" + time.Now().Format("150405.000")
|
||||
conn := &Conn{
|
||||
id: connID,
|
||||
userID: userID,
|
||||
rooms: make(map[string]bool),
|
||||
ws: ws,
|
||||
send: make(chan []byte, 64),
|
||||
hub: h,
|
||||
}
|
||||
|
||||
h.mu.Lock()
|
||||
h.conns[connID] = conn
|
||||
h.mu.Unlock()
|
||||
|
||||
log.Printf("[ws] connected: %s (user=%s, total=%d)", connID, userID, h.ConnCount())
|
||||
|
||||
// Subscribe this connection to bus events destined for clients
|
||||
conn.subscribeToBus()
|
||||
|
||||
// Publish presence
|
||||
h.bus.Publish(Event{
|
||||
Label: "user.presence",
|
||||
Payload: mustJSON(map[string]any{"user_id": userID, "status": "online"}),
|
||||
Ts: time.Now().UnixMilli(),
|
||||
SenderID: userID,
|
||||
})
|
||||
|
||||
// Start read/write pumps
|
||||
go conn.writePump()
|
||||
go conn.readPump()
|
||||
}
|
||||
|
||||
// ConnCount returns the number of active connections.
|
||||
func (h *Hub) ConnCount() int {
|
||||
h.mu.RLock()
|
||||
defer h.mu.RUnlock()
|
||||
return len(h.conns)
|
||||
}
|
||||
|
||||
// ConnsByUser returns all connections for a given user.
|
||||
func (h *Hub) ConnsByUser(userID string) []*Conn {
|
||||
h.mu.RLock()
|
||||
defer h.mu.RUnlock()
|
||||
var result []*Conn
|
||||
for _, c := range h.conns {
|
||||
if c.userID == userID {
|
||||
result = append(result, c)
|
||||
}
|
||||
}
|
||||
return result
|
||||
}
|
||||
|
||||
// removeConn cleans up a connection.
|
||||
func (h *Hub) removeConn(conn *Conn) {
|
||||
h.mu.Lock()
|
||||
delete(h.conns, conn.id)
|
||||
h.mu.Unlock()
|
||||
|
||||
// Unsubscribe from bus
|
||||
for _, unsub := range conn.unsubs {
|
||||
unsub()
|
||||
}
|
||||
|
||||
log.Printf("[ws] disconnected: %s (total=%d)", conn.id, h.ConnCount())
|
||||
|
||||
// Publish presence offline (only if no other conns for this user)
|
||||
if len(h.ConnsByUser(conn.userID)) == 0 {
|
||||
h.bus.Publish(Event{
|
||||
Label: "user.presence",
|
||||
Payload: mustJSON(map[string]any{"user_id": conn.userID, "status": "offline"}),
|
||||
Ts: time.Now().UnixMilli(),
|
||||
SenderID: conn.userID,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// subscribeToBus registers this conn as a bus subscriber for client-facing events.
|
||||
func (c *Conn) subscribeToBus() {
|
||||
unsub := c.hub.bus.Subscribe("*", func(e Event) {
|
||||
// Only forward events destined for clients
|
||||
if !ShouldSendToClient(e.Label) {
|
||||
return
|
||||
}
|
||||
|
||||
// Don't echo typing events back to the sender
|
||||
if e.Label == "chat.typing" || e.ConnID == c.id {
|
||||
if e.SenderID == c.userID && e.ConnID == c.id {
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
// Room filtering: if event has a room, only send if conn is in that room
|
||||
if e.Room != "" && !c.rooms[e.Room] {
|
||||
return
|
||||
}
|
||||
|
||||
data, err := json.Marshal(e)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
|
||||
select {
|
||||
case c.send <- data:
|
||||
default:
|
||||
// send buffer full, drop event
|
||||
}
|
||||
})
|
||||
|
||||
c.unsubs = append(c.unsubs, unsub)
|
||||
}
|
||||
|
||||
// JoinRoom adds this connection to a named room.
|
||||
func (c *Conn) JoinRoom(room string) {
|
||||
c.rooms[room] = true
|
||||
}
|
||||
|
||||
// LeaveRoom removes this connection from a named room.
|
||||
func (c *Conn) LeaveRoom(room string) {
|
||||
delete(c.rooms, room)
|
||||
}
|
||||
|
||||
// ── Read/Write Pumps ─────────────────────────
|
||||
|
||||
func (c *Conn) readPump() {
|
||||
defer func() {
|
||||
c.hub.removeConn(c)
|
||||
c.ws.Close()
|
||||
}()
|
||||
|
||||
c.ws.SetReadLimit(maxMsgSize)
|
||||
c.ws.SetReadDeadline(time.Now().Add(pongWait))
|
||||
c.ws.SetPongHandler(func(string) error {
|
||||
c.ws.SetReadDeadline(time.Now().Add(pongWait))
|
||||
return nil
|
||||
})
|
||||
|
||||
for {
|
||||
_, message, err := c.ws.ReadMessage()
|
||||
if err != nil {
|
||||
break
|
||||
}
|
||||
|
||||
var event Event
|
||||
if err := json.Unmarshal(message, &event); err != nil {
|
||||
continue
|
||||
}
|
||||
|
||||
// Handle ping
|
||||
if event.Label == "ping" {
|
||||
pong, _ := json.Marshal(Event{Label: "pong", Ts: time.Now().UnixMilli()})
|
||||
select {
|
||||
case c.send <- pong:
|
||||
default:
|
||||
}
|
||||
continue
|
||||
}
|
||||
|
||||
// Validate: only accept events allowed from clients
|
||||
if !ShouldAcceptFromClient(event.Label) {
|
||||
continue
|
||||
}
|
||||
|
||||
// Tag with sender info
|
||||
event.SenderID = c.userID
|
||||
event.ConnID = c.id
|
||||
if event.Ts == 0 {
|
||||
event.Ts = time.Now().UnixMilli()
|
||||
}
|
||||
|
||||
// Publish into bus
|
||||
c.hub.bus.Publish(event)
|
||||
}
|
||||
}
|
||||
|
||||
func (c *Conn) writePump() {
|
||||
ticker := time.NewTicker(pingPeriod)
|
||||
defer func() {
|
||||
ticker.Stop()
|
||||
c.ws.Close()
|
||||
}()
|
||||
|
||||
for {
|
||||
select {
|
||||
case msg, ok := <-c.send:
|
||||
c.ws.SetWriteDeadline(time.Now().Add(writeWait))
|
||||
if !ok {
|
||||
c.ws.WriteMessage(websocket.CloseMessage, []byte{})
|
||||
return
|
||||
}
|
||||
if err := c.ws.WriteMessage(websocket.TextMessage, msg); err != nil {
|
||||
return
|
||||
}
|
||||
|
||||
case <-ticker.C:
|
||||
c.ws.SetWriteDeadline(time.Now().Add(writeWait))
|
||||
if err := c.ws.WriteMessage(websocket.PingMessage, nil); err != nil {
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ── Helpers ──────────────────────────────────
|
||||
|
||||
func mustJSON(v any) json.RawMessage {
|
||||
data, err := json.Marshal(v)
|
||||
if err != nil {
|
||||
return json.RawMessage(`{}`)
|
||||
}
|
||||
return data
|
||||
}
|
||||
Reference in New Issue
Block a user