otto/chatter
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

rate limiting

Browse Source
This commit is contained in:
Otto Helen-Goldring
2025-05-25 15:00:23 +01:00
parent cce9b73a41
commit 2fa54c05ce
1 changed files with 82 additions and 37 deletions
Download Patch File Download Diff File Expand all files Collapse all files

119
main.go
View File

@ -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: