rate limiting

main.go

@@ -31,10 +31,13 @@ import (
 )
 
 const (
-	host              = "0.0.0.0"
+	host                 = "0.0.0.0"
-	port              = "23234"
+	port                 = "23234"
-	maxRoomNameLength = 24
+	maxRoomNameLength    = 24
-	roomNameCharset   = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789"
+	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.activePrograms[prog] = programInfo{roomName: "", userID: userID} // Update program's state
-	a.mu.Unlock() // Unlock before broadcasting
+	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.
+  /help         Show this help message.
-  /exit          Leave the current room.
+  /exit         Leave the current room.
-  /users         List users in this 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 early as state has significantly changed
-						return m, tea.Batch(cmds...)
 					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
+			} 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 message for different room '%s': %s", m.id, m.currentRoom, msg.room, msg.text)
+				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"))
+	titleStyle := lipgloss.NewStyle().Bold(true).Foreground(lipgloss.Color("63")) // Purple
-	errorStyle := lipgloss.NewStyle().Foreground(lipgloss.Color("9")).Bold(true) // Red
+	errorStyle := lipgloss.NewStyle().Foreground(lipgloss.Color("9")).Bold(true)  // Red
 
 	switch m.phase {
 	case phaseRoomInput: