434 lines
12 KiB
Go
434 lines
12 KiB
Go
package ws
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"fmt"
|
|
"log"
|
|
"net/http"
|
|
"strconv"
|
|
"strings"
|
|
"sync"
|
|
"time"
|
|
|
|
"github.com/gorilla/websocket"
|
|
"gorm.io/gorm"
|
|
|
|
"chat/src/db"
|
|
"chat/src/ratelimit"
|
|
)
|
|
|
|
// Upgrader is the WebSocket upgrader configuration.
|
|
var Upgrader = websocket.Upgrader{
|
|
CheckOrigin: func(r *http.Request) bool { return true },
|
|
}
|
|
|
|
// ClientInfo stores metadata about a connected WebSocket client.
|
|
type ClientInfo struct {
|
|
Username string
|
|
Room string
|
|
}
|
|
|
|
// Hub manages all WebSocket connections and message routing.
|
|
type Hub struct {
|
|
mu sync.Mutex
|
|
clients map[*websocket.Conn]bool
|
|
clientMeta map[*websocket.Conn]ClientInfo
|
|
DB *gorm.DB
|
|
Redis db.RedisClient
|
|
InstanceID string
|
|
RateLimitMax int
|
|
RateLimitWindow time.Duration
|
|
AdminUsername string
|
|
DMsSent int64
|
|
}
|
|
|
|
// NewHub creates a new WebSocket hub.
|
|
func NewHub(gormDB *gorm.DB, redis db.RedisClient, instanceID string, rateMax int, rateWindow time.Duration, adminUser string) *Hub {
|
|
return &Hub{
|
|
clients: make(map[*websocket.Conn]bool),
|
|
clientMeta: make(map[*websocket.Conn]ClientInfo),
|
|
DB: gormDB,
|
|
Redis: redis,
|
|
InstanceID: instanceID,
|
|
RateLimitMax: rateMax,
|
|
RateLimitWindow: rateWindow,
|
|
AdminUsername: adminUser,
|
|
}
|
|
}
|
|
|
|
// ConnectedClients returns the number of currently connected clients.
|
|
func (h *Hub) ConnectedClients() int {
|
|
h.mu.Lock()
|
|
defer h.mu.Unlock()
|
|
return len(h.clients)
|
|
}
|
|
|
|
// Broadcast sends a message to all connected clients and persists to DB.
|
|
func (h *Hub) Broadcast(msg db.Message) {
|
|
h.mu.Lock()
|
|
defer h.mu.Unlock()
|
|
|
|
// Persist to configured SQL backend.
|
|
if h.DB != nil {
|
|
h.DB.Create(&msg)
|
|
}
|
|
|
|
// Local fan-out
|
|
for client := range h.clients {
|
|
if err := client.WriteJSON(msg); err != nil {
|
|
client.Close()
|
|
delete(h.clients, client)
|
|
delete(h.clientMeta, client)
|
|
}
|
|
}
|
|
|
|
// Redis pub/sub for multi-instance fan-out
|
|
if h.Redis != nil && msg.Room != "" {
|
|
payload, _ := json.Marshal(map[string]any{
|
|
"instance": h.InstanceID,
|
|
"msg": msg,
|
|
})
|
|
_ = h.Redis.Publish(context.Background(), fmt.Sprintf("room:%s:messages", msg.Room), payload)
|
|
}
|
|
}
|
|
|
|
// PushChatHistory sends the last N messages for a room via WebSocket.
|
|
func (h *Hub) PushChatHistory(conn *websocket.Conn, room string, limit int) {
|
|
if h.DB == nil || room == "" {
|
|
return
|
|
}
|
|
var messages []db.Message
|
|
h.DB.Where("room = ?", room).Order("timestamp desc").Limit(limit).Find(&messages)
|
|
// Reverse to chronological order
|
|
for i := len(messages)/2 - 1; i >= 0; i-- {
|
|
opp := len(messages) - 1 - i
|
|
messages[i], messages[opp] = messages[opp], messages[i]
|
|
}
|
|
for _, msg := range messages {
|
|
conn.WriteJSON(msg)
|
|
}
|
|
}
|
|
|
|
// HandleConnection handles a new WebSocket connection.
|
|
func (h *Hub) HandleConnection(w http.ResponseWriter, r *http.Request) {
|
|
clientIP := r.Header.Get("X-Forwarded-For")
|
|
if clientIP == "" {
|
|
clientIP = r.RemoteAddr
|
|
}
|
|
|
|
// IP-based connection rate limiting
|
|
allowed, err := ratelimit.CheckWSConnRate(context.Background(), h.Redis, clientIP)
|
|
if err != nil {
|
|
log.Printf("WS rate limit check error: %v", err)
|
|
}
|
|
if !allowed {
|
|
http.Error(w, "too many connections from this IP", http.StatusTooManyRequests)
|
|
return
|
|
}
|
|
defer ratelimit.DecrWSConn(context.Background(), h.Redis, clientIP)
|
|
|
|
conn, err := Upgrader.Upgrade(w, r, nil)
|
|
if err != nil {
|
|
log.Println("Upgrade error:", err)
|
|
return
|
|
}
|
|
|
|
username := r.URL.Query().Get("username")
|
|
room := r.URL.Query().Get("room")
|
|
if username == "" {
|
|
username = fmt.Sprintf("anon_%d", time.Now().UnixNano()%10000)
|
|
}
|
|
|
|
h.mu.Lock()
|
|
h.clients[conn] = true
|
|
h.clientMeta[conn] = ClientInfo{Username: username, Room: room}
|
|
h.mu.Unlock()
|
|
|
|
// Announce presence (ADR-007)
|
|
if h.Redis != nil && room != "" {
|
|
_ = h.Redis.Set(context.Background(), fmt.Sprintf("presence:%s:%s", room, username), "online", 60*time.Second)
|
|
h.Broadcast(db.Message{Room: room, Username: "system", Text: username + " joined " + room, Type: "system", Timestamp: time.Now()})
|
|
}
|
|
|
|
// Push chat history on join (last 50 messages)
|
|
h.PushChatHistory(conn, room, 50)
|
|
|
|
// Heartbeat goroutine: refresh presence TTL every 30s
|
|
heartbeatStop := make(chan struct{})
|
|
go func() {
|
|
ticker := time.NewTicker(30 * time.Second)
|
|
defer ticker.Stop()
|
|
for {
|
|
select {
|
|
case <-ticker.C:
|
|
if h.Redis != nil && room != "" {
|
|
_ = h.Redis.Set(context.Background(), fmt.Sprintf("presence:%s:%s", room, username), "online", 60*time.Second)
|
|
}
|
|
// Ping to detect dead connections
|
|
if err := conn.WriteMessage(websocket.PingMessage, nil); err != nil {
|
|
return
|
|
}
|
|
case <-heartbeatStop:
|
|
return
|
|
}
|
|
}
|
|
}()
|
|
|
|
defer func() {
|
|
close(heartbeatStop)
|
|
h.mu.Lock()
|
|
delete(h.clients, conn)
|
|
delete(h.clientMeta, conn)
|
|
h.mu.Unlock()
|
|
if h.Redis != nil && room != "" {
|
|
_ = h.Redis.Del(context.Background(), fmt.Sprintf("presence:%s:%s", room, username))
|
|
h.Broadcast(db.Message{Room: room, Username: "system", Text: username + " left " + room, Type: "system", Timestamp: time.Now()})
|
|
}
|
|
conn.Close()
|
|
}()
|
|
|
|
for {
|
|
var msg db.Message
|
|
if err := conn.ReadJSON(&msg); err != nil {
|
|
break
|
|
}
|
|
|
|
msg.Timestamp = time.Now()
|
|
msg.Username = username // enforce from query param for security
|
|
|
|
// Check if banned before processing
|
|
if h.Redis != nil {
|
|
if isBanned, _ := h.Redis.Exists(context.Background(), fmt.Sprintf("banned:%s", username)); isBanned > 0 {
|
|
conn.WriteJSON(db.Message{Type: "system", Text: "You are banned from chat.", Timestamp: time.Now()})
|
|
continue
|
|
}
|
|
if isMuted, _ := h.Redis.Exists(context.Background(), fmt.Sprintf("muted:%s", username)); isMuted > 0 {
|
|
conn.WriteJSON(db.Message{Type: "system", Text: "You are muted.", Timestamp: time.Now()})
|
|
continue
|
|
}
|
|
}
|
|
|
|
// Rate limiting check
|
|
allowed, _ := ratelimit.CheckMessageRate(context.Background(), h.Redis, msg.Username, h.RateLimitMax, h.RateLimitWindow)
|
|
if !allowed {
|
|
conn.WriteJSON(db.Message{Type: "system", Text: "Rate limit exceeded. Please slow down.", Timestamp: time.Now()})
|
|
continue
|
|
}
|
|
|
|
// Handle moderation commands
|
|
if strings.HasPrefix(msg.Text, "/mod ") || strings.HasPrefix(msg.Text, "/mute ") || strings.HasPrefix(msg.Text, "/ban ") {
|
|
h.handleModerationCommand(msg, conn)
|
|
continue
|
|
}
|
|
|
|
// Handle private DMs
|
|
if strings.HasPrefix(msg.Room, "dm:") {
|
|
h.handlePrivateDM(msg)
|
|
continue
|
|
}
|
|
|
|
msg.Type = "public"
|
|
log.Printf("Chat message in room %s from %s: %s", msg.Room, msg.Username, msg.Text)
|
|
h.Broadcast(msg)
|
|
}
|
|
}
|
|
|
|
// handlePrivateDM routes private messages to the specific user pair.
|
|
func (h *Hub) handlePrivateDM(msg db.Message) {
|
|
msg.Type = "private"
|
|
h.DMsSent++
|
|
|
|
// Persist for history
|
|
if h.DB != nil {
|
|
h.DB.Create(&msg)
|
|
}
|
|
|
|
// Deliver to target if online
|
|
h.mu.Lock()
|
|
targetOnline := false
|
|
for conn, info := range h.clientMeta {
|
|
if info.Username == msg.Room || strings.Contains(msg.Room, info.Username) {
|
|
if err := conn.WriteJSON(msg); err == nil {
|
|
targetOnline = true
|
|
}
|
|
}
|
|
}
|
|
h.mu.Unlock()
|
|
|
|
if !targetOnline {
|
|
log.Printf("DM stored for offline user: %s -> room %s", msg.Username, msg.Room)
|
|
}
|
|
|
|
// Also deliver to sender for confirmation
|
|
h.mu.Lock()
|
|
for conn, info := range h.clientMeta {
|
|
if info.Username == msg.Username {
|
|
conn.WriteJSON(msg)
|
|
break
|
|
}
|
|
}
|
|
h.mu.Unlock()
|
|
|
|
log.Printf("Private DM from %s to room %s: %s", msg.Username, msg.Room, msg.Text)
|
|
}
|
|
|
|
// logModerationAction logs an action to the immutable ModerationLog table.
|
|
func (h *Hub) logModerationAction(action, target, moderator, reason string) {
|
|
logEntry := db.ModerationLog{
|
|
Action: action,
|
|
Target: target,
|
|
Moderator: moderator,
|
|
Reason: reason,
|
|
Timestamp: time.Now(),
|
|
}
|
|
if h.DB != nil {
|
|
h.DB.Create(&logEntry)
|
|
}
|
|
log.Printf("Moderation action logged: %s %s by %s - %s", action, target, moderator, reason)
|
|
}
|
|
|
|
// handleModerationCommand processes /mod, /mute, /ban commands.
|
|
func (h *Hub) handleModerationCommand(msg db.Message, conn *websocket.Conn) {
|
|
parts := strings.Fields(msg.Text)
|
|
if len(parts) < 3 {
|
|
conn.WriteJSON(db.Message{Type: "system", Text: "Usage: /mute <user> <duration_m> <reason> or /ban <user> <reason>", Timestamp: time.Now()})
|
|
return
|
|
}
|
|
|
|
if msg.Username != h.AdminUsername {
|
|
conn.WriteJSON(db.Message{Type: "system", Text: "Only moderators can use moderation commands.", Timestamp: time.Now()})
|
|
return
|
|
}
|
|
|
|
action := strings.TrimPrefix(parts[0], "/")
|
|
target := parts[1]
|
|
reason := strings.Join(parts[2:], " ")
|
|
|
|
if h.Redis != nil {
|
|
switch action {
|
|
case "mute", "mod":
|
|
durationMinutes := 10 // default
|
|
if len(parts) >= 3 {
|
|
if d, err := strconv.Atoi(parts[2]); err == nil && d > 0 {
|
|
durationMinutes = d
|
|
if len(parts) > 3 {
|
|
reason = strings.Join(parts[3:], " ")
|
|
} else {
|
|
reason = "no reason given"
|
|
}
|
|
}
|
|
}
|
|
_ = h.Redis.Set(context.Background(), fmt.Sprintf("muted:%s", target), "1", time.Duration(durationMinutes)*time.Minute)
|
|
case "ban":
|
|
_ = h.Redis.Set(context.Background(), fmt.Sprintf("banned:%s", target), "1", 0) // no expiry
|
|
}
|
|
}
|
|
|
|
h.logModerationAction(action, target, msg.Username, reason)
|
|
|
|
h.Broadcast(db.Message{
|
|
Room: msg.Room,
|
|
Username: "system",
|
|
Type: "system",
|
|
Text: fmt.Sprintf("%s was %s by %s: %s", target, action, msg.Username, reason),
|
|
Timestamp: time.Now(),
|
|
})
|
|
}
|
|
|
|
// SubscribeRoomMessages listens for Redis pub/sub messages and fans out locally.
|
|
func (h *Hub) SubscribeRoomMessages() {
|
|
if h.Redis == nil {
|
|
return
|
|
}
|
|
pubsub := h.Redis.PSubscribe(context.Background(), "room:*:messages")
|
|
defer pubsub.Close()
|
|
|
|
log.Printf("Redis pub/sub subscriber started (instance %s)", h.InstanceID)
|
|
|
|
for {
|
|
msg, err := pubsub.ReceiveMessage(context.Background())
|
|
if err != nil {
|
|
log.Printf("Redis pub/sub error: %v — retrying in 5s", err)
|
|
time.Sleep(5 * time.Second)
|
|
pubsub = h.Redis.PSubscribe(context.Background(), "room:*:messages")
|
|
continue
|
|
}
|
|
|
|
var envelope struct {
|
|
Instance string `json:"instance"`
|
|
Msg db.Message `json:"msg"`
|
|
}
|
|
if err := json.Unmarshal([]byte(msg.Payload), &envelope); err != nil {
|
|
continue
|
|
}
|
|
|
|
// Skip messages originating from this instance
|
|
if envelope.Instance == h.InstanceID {
|
|
continue
|
|
}
|
|
|
|
// Fan out to local clients
|
|
h.mu.Lock()
|
|
for client := range h.clients {
|
|
client.WriteJSON(envelope.Msg)
|
|
}
|
|
h.mu.Unlock()
|
|
}
|
|
}
|
|
|
|
// ScanPresence returns all active presence entries from Redis.
|
|
func (h *Hub) ScanPresence(ctx context.Context) ([]db.Presence, error) {
|
|
if h.Redis == nil {
|
|
return nil, nil
|
|
}
|
|
keys, _, err := h.Redis.Scan(ctx, 0, "presence:*", 200)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
var presences []db.Presence
|
|
for _, key := range keys {
|
|
parts := strings.SplitN(key, ":", 3)
|
|
if len(parts) != 3 {
|
|
continue
|
|
}
|
|
presences = append(presences, db.Presence{
|
|
Room: parts[1],
|
|
Username: parts[2],
|
|
Online: true,
|
|
})
|
|
}
|
|
return presences, nil
|
|
}
|
|
|
|
// MuteUser mutes a user for a duration.
|
|
func (h *Hub) MuteUser(ctx context.Context, username string, duration time.Duration) error {
|
|
if h.Redis == nil {
|
|
return fmt.Errorf("redis unavailable")
|
|
}
|
|
return h.Redis.Set(ctx, fmt.Sprintf("muted:%s", username), "1", duration)
|
|
}
|
|
|
|
// UnmuteUser removes a mute.
|
|
func (h *Hub) UnmuteUser(ctx context.Context, username string) error {
|
|
if h.Redis == nil {
|
|
return fmt.Errorf("redis unavailable")
|
|
}
|
|
return h.Redis.Del(ctx, fmt.Sprintf("muted:%s", username))
|
|
}
|
|
|
|
// BanUser bans a user permanently.
|
|
func (h *Hub) BanUser(ctx context.Context, username string) error {
|
|
if h.Redis == nil {
|
|
return fmt.Errorf("redis unavailable")
|
|
}
|
|
return h.Redis.Set(ctx, fmt.Sprintf("banned:%s", username), "1", 0)
|
|
}
|
|
|
|
// UnbanUser removes a ban.
|
|
func (h *Hub) UnbanUser(ctx context.Context, username string) error {
|
|
if h.Redis == nil {
|
|
return fmt.Errorf("redis unavailable")
|
|
}
|
|
return h.Redis.Del(ctx, fmt.Sprintf("banned:%s", username))
|
|
} |