package handlers import ( "encoding/json" "log" "net/http" "sync" "threadr/models" "html/template" "regexp" "time" "github.com/gorilla/sessions" "github.com/gorilla/websocket" ) var upgrader = websocket.Upgrader{ ReadBufferSize: 1024, WriteBufferSize: 1024, CheckOrigin: func(r *http.Request) bool { return true // Allow all origins for now; restrict in production }, } // ChatHub manages WebSocket connections and broadcasts messages type ChatHub struct { clients map[*websocket.Conn]int // Map of connections to user IDs broadcast chan []byte register chan *websocket.Conn unregister chan *websocket.Conn mutex sync.Mutex } func NewChatHub() *ChatHub { return &ChatHub{ clients: make(map[*websocket.Conn]int), broadcast: make(chan []byte), register: make(chan *websocket.Conn), unregister: make(chan *websocket.Conn), } } func (h *ChatHub) Run() { for { select { case client := <-h.register: h.mutex.Lock() h.clients[client] = 0 // UserID set later h.mutex.Unlock() case client := <-h.unregister: h.mutex.Lock() delete(h.clients, client) h.mutex.Unlock() client.Close() case message := <-h.broadcast: h.mutex.Lock() for client := range h.clients { err := client.WriteMessage(websocket.TextMessage, message) if err != nil { log.Printf("Error broadcasting message: %v", err) client.Close() delete(h.clients, client) } } h.mutex.Unlock() } } } var hub = NewChatHub() func init() { go hub.Run() } // ChatDisplayMessage is used to pass chat message data to the template, // with content pre-formatted for HTML rendering. type ChatDisplayMessage struct { ID int UserID int Content string // Raw content for JS processing if needed, maybe not FormattedContent template.HTML // Content with mentions as HTML spans ReplyTo int Timestamp time.Time Username string PfpURL string } // formatChatMessageContent takes raw message content and wraps @mentions in HTML spans. func formatChatMessageContent(content string) template.HTML { // Regex to find @username. $0 will be replaced with the entire match (e.g., @username) r := regexp.MustCompile(`@([a-zA-Z0-9_]+)`) formatted := r.ReplaceAllString(content, `$0`) return template.HTML(formatted) } func ChatHandler(app *App) http.HandlerFunc { return func(w http.ResponseWriter, r *http.Request) { session := r.Context().Value("session").(*sessions.Session) userID, ok := session.Values["user_id"].(int) if !ok { http.Redirect(w, r, app.Config.ThreadrDir+"/login/", http.StatusFound) return } cookie, _ := r.Cookie("threadr_cookie_banner") if r.URL.Query().Get("ws") == "true" { // Handle WebSocket connection ws, err := upgrader.Upgrade(w, r, nil) if err != nil { log.Printf("Error upgrading to WebSocket: %v", err) return } hub.register <- ws hub.mutex.Lock() hub.clients[ws] = userID hub.mutex.Unlock() defer func() { hub.unregister <- ws }() for { _, msg, err := ws.ReadMessage() if err != nil { log.Printf("Error reading WebSocket message: %v", err) break } var chatMsg struct { Type string `json:"type"` Content string `json:"content"` ReplyTo int `json:"replyTo"` } if err := json.Unmarshal(msg, &chatMsg); err != nil { log.Printf("Error unmarshaling message: %v", err) continue } if chatMsg.Type == "message" { msgObj := models.ChatMessage{ UserID: userID, Content: chatMsg.Content, ReplyTo: chatMsg.ReplyTo, } // CreateChatMessage now also handles notifications and returns the new message ID newMsgID, err := models.CreateChatMessage(app.DB, msgObj) if err != nil { log.Printf("Error saving chat message: %v", err) continue } // Fetch the saved message with timestamp and user details savedMsg, err := models.GetChatMessageByID(app.DB, newMsgID) if err != nil { log.Printf("Error fetching saved message: %v", err) continue } if savedMsg == nil { log.Printf("Error: saved message with ID %d not found after creation", newMsgID) continue } // Prepare ChatDisplayMessage for broadcast displayMsg := ChatDisplayMessage{ ID: savedMsg.ID, UserID: savedMsg.UserID, Content: savedMsg.Content, // Keep raw content if needed by JS, or remove FormattedContent: formatChatMessageContent(savedMsg.Content), ReplyTo: savedMsg.ReplyTo, Timestamp: savedMsg.Timestamp, Username: savedMsg.Username, PfpURL: savedMsg.PfpURL, } response, _ := json.Marshal(displayMsg) hub.broadcast <- response } } return } if r.URL.Query().Get("autocomplete") == "true" { // Handle autocomplete for mentions prefix := r.URL.Query().Get("prefix") usernames, err := models.GetUsernamesMatching(app.DB, prefix) if err != nil { log.Printf("Error fetching usernames for autocomplete: %v", err) http.Error(w, "Internal Server Error", http.StatusInternalServerError) return } response, _ := json.Marshal(usernames) w.Header().Set("Content-Type", "application/json") w.Write(response) return } // Render chat page rawMessages, err := models.GetRecentChatMessages(app.DB, 50) if err != nil { log.Printf("Error fetching chat messages: %v", err) http.Error(w, "Internal Server Error", http.StatusInternalServerError) return } var displayMessages []ChatDisplayMessage for _, msg := range rawMessages { displayMessages = append(displayMessages, ChatDisplayMessage{ ID: msg.ID, UserID: msg.UserID, Content: msg.Content, // Raw content FormattedContent: formatChatMessageContent(msg.Content), ReplyTo: msg.ReplyTo, Timestamp: msg.Timestamp, Username: msg.Username, PfpURL: msg.PfpURL, }) } // Reverse messages to show oldest first for i, j := 0, len(displayMessages)-1; i < j; i, j = i+1, j-1 { displayMessages[i], displayMessages[j] = displayMessages[j], displayMessages[i] } data := struct { PageData Messages []ChatDisplayMessage }{ PageData: PageData{ Title: "ThreadR - Chat", Navbar: "chat", LoggedIn: true, ShowCookieBanner: cookie == nil || cookie.Value != "accepted", BasePath: app.Config.ThreadrDir, StaticPath: app.Config.ThreadrDir + "/static", CurrentURL: r.URL.Path, }, Messages: displayMessages, } if err := app.Tmpl.ExecuteTemplate(w, "chat", data); err != nil { log.Printf("Error executing template in ChatHandler: %v", err) http.Error(w, "Internal Server Error", http.StatusInternalServerError) return } } }