diff --git a/main.go b/main.go index 0aeddd0..aa0769d 100644 --- a/main.go +++ b/main.go @@ -31,10 +31,13 @@ import ( ) const ( - host = "0.0.0.0" - port = "23234" - maxRoomNameLength = 24 - roomNameCharset = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789" + host = "0.0.0.0" + port = "23234" + maxRoomNameLength = 24 + roomNameCharset = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789" + rateLimitMaxMessages = 15 // Max messages per user within the rateLimitWindow + rateLimitWindow = 15 * time.Second // The time window for rate limiting + maxChatMessageLength = 2000 // Max characters per chat message ) // generateRandomRoomName creates a random alphanumeric string of a given length. @@ -156,7 +159,7 @@ func (a *app) ProgramHandler(s ssh.Session) *tea.Program { userID := s.User() m := initialModel(a, userID) p := tea.NewProgram(m, bubbletea.MakeOptions(s)...) - m.SetProgram(p) + m.SetProgram(p) // Give the model a reference to its own program a.addProgram(p, userID) return p } @@ -180,6 +183,7 @@ func (a *app) AssignProgramToRoom(prog *tea.Program, roomName string) { userID := currentInfo.userID currentRoom := currentInfo.roomName + // If program was in a different room, remove it from the old room's list if currentRoom != "" && currentRoom != roomName { var updatedRoomProgs []*tea.Program for _, pInRoom := range a.rooms[currentRoom] { @@ -195,10 +199,12 @@ func (a *app) AssignProgramToRoom(prog *tea.Program, roomName string) { log.Debugf("Program %p (user: %s) moved from room %s", prog, userID, currentRoom) } + // Ensure the new room exists in the map if _, ok := a.rooms[roomName]; !ok { a.rooms[roomName] = []*tea.Program{} } + // Add program to the new room's list if not already there foundInNewRoomList := false for _, pInList := range a.rooms[roomName] { if pInList == prog { @@ -230,17 +236,19 @@ func (a *app) HandleUserExitFromRoom(prog *tea.Program, oldRoomName string) { } userID := info.userID + // Prepare list of recipients for the leave message BEFORE modifying room structure leaveMsgText := fmt.Sprintf("%s has left the room.", userID) var recipients []*tea.Program if progs, ok := a.rooms[info.roomName]; ok { - recipients = make([]*tea.Program, 0, len(progs)) + recipients = make([]*tea.Program, 0, len(progs)) // Allocate with capacity for _, p := range progs { - if p != prog { + if p != prog { // Don't send to the exiting program recipients = append(recipients, p) } } } + // Now, update the room structure var updatedRoomProgs []*tea.Program for _, pInRoom := range a.rooms[info.roomName] { if pInRoom != prog { @@ -250,11 +258,12 @@ func (a *app) HandleUserExitFromRoom(prog *tea.Program, oldRoomName string) { if len(updatedRoomProgs) > 0 { a.rooms[info.roomName] = updatedRoomProgs } else { - delete(a.rooms, info.roomName) + delete(a.rooms, info.roomName) // Clean up empty room } - a.activePrograms[prog] = programInfo{roomName: "", userID: userID} - a.mu.Unlock() // Unlock before broadcasting + a.activePrograms[prog] = programInfo{roomName: "", userID: userID} // Update program's state + a.mu.Unlock() // Unlock before broadcasting + // Broadcast leave message if len(recipients) > 0 { for _, p := range recipients { go p.Send(chatMsg{id: "system", text: leaveMsgText, room: oldRoomName, isSystemMessage: true}) @@ -278,18 +287,21 @@ func (a *app) RemoveProgram(prog *tea.Program) { var leaveRecipients []*tea.Program var leaveMsgText string + + // If the user was in a room, prepare to notify others if roomName != "" { leaveMsgText = fmt.Sprintf("%s has disconnected.", userID) if progs, ok := a.rooms[roomName]; ok { leaveRecipients = make([]*tea.Program, 0, len(progs)) for _, p := range progs { - if p != prog { + if p != prog { // Don't collect the disconnecting program leaveRecipients = append(leaveRecipients, p) } } } } + // Remove from room tracking if roomName != "" { var updatedRoomProgs []*tea.Program for _, pInRoom := range a.rooms[roomName] { @@ -300,13 +312,16 @@ func (a *app) RemoveProgram(prog *tea.Program) { if len(updatedRoomProgs) > 0 { a.rooms[roomName] = updatedRoomProgs } else { - delete(a.rooms, roomName) + delete(a.rooms, roomName) // Clean up room if it becomes empty } log.Infof("Program %p (user: %s) removed from room %s tracking (full disconnect).", prog, userID, roomName) } + + // Remove from active programs list delete(a.activePrograms, prog) a.mu.Unlock() // Unlock before broadcasting + // Broadcast disconnect message to former room members if roomName != "" && len(leaveRecipients) > 0 { for _, p := range leaveRecipients { go p.Send(chatMsg{id: "system", text: leaveMsgText, room: roomName, isSystemMessage: true}) @@ -321,8 +336,8 @@ func main() { // Optional: Setup file logging // file, err := os.OpenFile("chatter_server.log", os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0666) // if err == nil { - // defer file.Close() - // log.Default().SetOutput(io.MultiWriter(os.Stderr, file)) // Log to both stderr and file + // log.Default().SetOutput(io.MultiWriter(os.Stderr, file)) // Log to both stderr and file + // defer file.Close() // } // log.SetLevel(log.DebugLevel) // log.SetTimeFormat(time.Kitchen) @@ -365,7 +380,8 @@ type model struct { messages []string textarea textarea.Model senderStyle lipgloss.Style - systemMessageStyle lipgloss.Style // For join/leave/command output + systemMessageStyle lipgloss.Style // For join/leave/command output/errors + messageTimestamps []time.Time // For rate limiting err error } @@ -379,7 +395,7 @@ func initialModel(appInstance *app, userID string) *model { ta := textarea.New() ta.Placeholder = "Type your message or a command (e.g. /help)..." ta.Prompt = "┃ " - ta.CharLimit = 280 + ta.CharLimit = maxChatMessageLength // Updated to new max length ta.SetWidth(60) ta.SetHeight(3) ta.FocusedStyle.CursorLine = lipgloss.NewStyle() @@ -396,6 +412,7 @@ func initialModel(appInstance *app, userID string) *model { roomInput: roomTi, textarea: ta, messages: []string{}, + messageTimestamps: []time.Time{}, // Initialize for rate limiting viewport: vp, senderStyle: lipgloss.NewStyle().Foreground(lipgloss.Color("12")), // Blue for sender systemMessageStyle: lipgloss.NewStyle().Foreground(lipgloss.AdaptiveColor{Light: "#909090", Dark: "#6C6C6C"}), // Grey for system @@ -468,7 +485,7 @@ func (m *model) Update(msg tea.Msg) (tea.Model, tea.Cmd) { cmds = append(cmds, textarea.Blink) } else { m.err = errors.New("invalid room name: 1-24 alphanumeric characters required") - m.roomInput.Focus() + m.roomInput.Focus() // Keep focus on room input if error cmds = append(cmds, textinput.Blink) } return m, tea.Batch(cmds...) @@ -484,8 +501,9 @@ func (m *model) Update(msg tea.Msg) (tea.Model, tea.Cmd) { var tiCmd tea.Cmd var vpCmd tea.Cmd + // Order matters: update inputs first, then handle specific key presses or messages m.textarea, tiCmd = m.textarea.Update(msg) - m.viewport, vpCmd = m.viewport.Update(msg) + m.viewport, vpCmd = m.viewport.Update(msg) // Handles scrolling, etc. cmds = append(cmds, tiCmd, vpCmd) switch msg := msg.(type) { @@ -508,13 +526,12 @@ func (m *model) Update(msg tea.Msg) (tea.Model, tea.Cmd) { switch command { case "/help": helpText := m.systemMessageStyle.Render(`Available commands: - /help Show this help message. - /exit Leave the current room. - /users List users in this room.`) + /help Show this help message. + /exit Leave the current room. + /users List users in this room.`) m.messages = append(m.messages, helpText) case "/exit", "/leave": if m.program != nil && m.currentRoom != "" { - // App handles broadcasting leave message m.app.HandleUserExitFromRoom(m.program, m.currentRoom) } m.phase = phaseRoomInput @@ -522,13 +539,12 @@ func (m *model) Update(msg tea.Msg) (tea.Model, tea.Cmd) { m.messages = []string{} m.err = nil m.roomInput.Reset() - m.roomInput.Placeholder = "Enter another room or press Enter for random..." // Update placeholder + m.roomInput.Placeholder = "Enter another room or press Enter for random..." m.roomInput.Focus() m.textarea.Blur() m.viewport.SetContent("You have left the room. Enter a new room name.") cmds = append(cmds, textinput.Blink) - // No further processing for this Update cycle after state change. - return m, tea.Batch(cmds...) + return m, tea.Batch(cmds...) // Return early as state has significantly changed case "/users": if m.currentRoom != "" { userIDs := m.app.GetUserIDsInRoom(m.currentRoom) @@ -536,7 +552,7 @@ func (m *model) Update(msg tea.Msg) (tea.Model, tea.Cmd) { if len(userIDs) > 0 { userListText = "Users in '" + m.currentRoom + "':\n " + strings.Join(userIDs, "\n ") } else { - userListText = "You are the only one here." // Should not happen if self is listed + userListText = "You are the only one here." } m.messages = append(m.messages, m.systemMessageStyle.Render(userListText)) } else { @@ -549,36 +565,65 @@ func (m *model) Update(msg tea.Msg) (tea.Model, tea.Cmd) { m.viewport.GotoBottom() } else if val != "" { // Regular chat message + // 1. Message Length Check + if len(val) > maxChatMessageLength { + errMsgText := fmt.Sprintf("Your message is too long (max %d chars). It was not sent.", maxChatMessageLength) + m.messages = append(m.messages, m.systemMessageStyle.Render(errMsgText)) + m.viewport.SetContent(strings.Join(m.messages, "\n")) + m.viewport.GotoBottom() + // Do not broadcast, keep textarea focused for correction or new message + cmds = append(cmds, m.textarea.Focus(), textarea.Blink) + return m, tea.Batch(cmds...) // Return early, message not sent + } + + // 2. Rate Limiting Check + now := time.Now() + var recentTimestamps []time.Time + for _, ts := range m.messageTimestamps { + if now.Sub(ts) < rateLimitWindow { + recentTimestamps = append(recentTimestamps, ts) + } + } + m.messageTimestamps = recentTimestamps + + if len(m.messageTimestamps) >= rateLimitMaxMessages { + errMsgText := fmt.Sprintf("You're sending messages too quickly (max %d per %s). Please wait. Your message was not sent.", rateLimitMaxMessages, rateLimitWindow.String()) + m.messages = append(m.messages, m.systemMessageStyle.Render(errMsgText)) + m.viewport.SetContent(strings.Join(m.messages, "\n")) + m.viewport.GotoBottom() + // Do not broadcast, keep textarea focused + cmds = append(cmds, m.textarea.Focus(), textarea.Blink) + return m, tea.Batch(cmds...) // Return early, message not sent + } + + // If all checks pass, add current timestamp and broadcast + m.messageTimestamps = append(m.messageTimestamps, now) + m.app.BroadcastToRoom(m.currentRoom, chatMsg{ id: m.id, text: val, room: m.currentRoom, isSystemMessage: false, }, nil) // Broadcast to all, including self (sender) - // Message will be added to m.messages when this client receives its own broadcast. } // After sending or command, ensure textarea is focused for next input cmds = append(cmds, m.textarea.Focus(), textarea.Blink) } case chatMsg: // Received a message (could be own, from others, or system) - if msg.room == m.currentRoom || (msg.isSystemMessage && msg.room == m.currentRoom) { // Ensure for current room + if msg.room == m.currentRoom { // Process only if message is for the current room var styledMessage string if msg.isSystemMessage { styledMessage = m.systemMessageStyle.Render(msg.text) } else { styledId := m.senderStyle.Render(msg.id) - if msg.id == m.id { // Own regular message - // Optionally style self messages differently, e.g., if not relying on system message for own display - // For now, senderStyle applies to all non-system messages. - } styledMessage = styledId + ": " + msg.text } m.messages = append(m.messages, styledMessage) m.viewport.SetContent(strings.Join(m.messages, "\n")) m.viewport.GotoBottom() - } else if msg.room != "" { // Message for a different room, should ideally not happen if app logic is correct - log.Warnf("Client '%s' in room '%s' received message for different room '%s': %s", m.id, m.currentRoom, msg.room, msg.text) + } else if msg.room != "" && !msg.isSystemMessage { // Log if a non-system message for a different room is somehow received by this model + log.Warnf("Client '%s' in room '%s' received non-system message for different room '%s': %s", m.id, m.currentRoom, msg.room, msg.text) } case errMsg: @@ -588,7 +633,7 @@ func (m *model) Update(msg tea.Msg) (tea.Model, tea.Cmd) { // m.messages = append(m.messages, m.systemMessageStyle.Render("An error occurred: "+msg.Error())) // m.viewport.SetContent(strings.Join(m.messages, "\n")) // m.viewport.GotoBottom() - return m, nil + return m, nil // Or tea.Batch(cmds...) if other cmds are important } } return m, tea.Batch(cmds...) @@ -597,8 +642,8 @@ func (m *model) Update(msg tea.Msg) (tea.Model, tea.Cmd) { func (m *model) View() string { var viewBuilder strings.Builder - titleStyle := lipgloss.NewStyle().Bold(true).Foreground(lipgloss.Color("63")) - errorStyle := lipgloss.NewStyle().Foreground(lipgloss.Color("9")).Bold(true) // Red + titleStyle := lipgloss.NewStyle().Bold(true).Foreground(lipgloss.Color("63")) // Purple + errorStyle := lipgloss.NewStyle().Foreground(lipgloss.Color("9")).Bold(true) // Red switch m.phase { case phaseRoomInput: