From 10dffb3eb284ef2cae53f86d07fe083c800bb9e2 Mon Sep 17 00:00:00 2001 From: Weetile Date: Sun, 25 May 2025 15:40:18 +0100 Subject: [PATCH] refactor --- main.go | 292 ++++++++++++++++++++++++++++++-------------------------- 1 file changed, 154 insertions(+), 138 deletions(-) diff --git a/main.go b/main.go index aa0769d..4fba475 100644 --- a/main.go +++ b/main.go @@ -133,15 +133,16 @@ func (a *app) Start() { go func() { if err = a.ListenAndServe(); err != nil && !errors.Is(err, ssh.ErrServerClosed) { log.Error("Could not start server", "error", err) - done <- syscall.SIGTERM + done <- syscall.SIGTERM // Trigger shutdown if server fails to start } }() - <-done + <-done // Wait for shutdown signal log.Info("Stopping SSH server...") ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) defer func() { cancel() + // Clear program and room data after shutdown attempt a.mu.Lock() a.activePrograms = make(map[*tea.Program]programInfo) a.rooms = make(map[string][]*tea.Program) @@ -149,7 +150,7 @@ func (a *app) Start() { log.Info("Cleared app program stores.") }() - if err := a.Shutdown(ctx); err != nil { + if err := a.Shutdown(ctx); err != nil && !errors.Is(err, ssh.ErrServerClosed) { // Don't log ErrServerClosed as an error during graceful shutdown log.Error("Could not gracefully stop server", "error", err) } log.Info("SSH server stopped.") @@ -157,6 +158,9 @@ func (a *app) Start() { func (a *app) ProgramHandler(s ssh.Session) *tea.Program { userID := s.User() + if userID == "" { + userID = "guest" // Default userID if empty + } m := initialModel(a, userID) p := tea.NewProgram(m, bubbletea.MakeOptions(s)...) m.SetProgram(p) // Give the model a reference to its own program @@ -185,18 +189,20 @@ func (a *app) AssignProgramToRoom(prog *tea.Program, roomName string) { // 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] { - if pInRoom != prog { - updatedRoomProgs = append(updatedRoomProgs, pInRoom) + if progsInOldRoom, ok := a.rooms[currentRoom]; ok { + var updatedRoomProgs []*tea.Program + for _, pInRoom := range progsInOldRoom { + if pInRoom != prog { + updatedRoomProgs = append(updatedRoomProgs, pInRoom) + } } + if len(updatedRoomProgs) > 0 { + a.rooms[currentRoom] = updatedRoomProgs + } else { + delete(a.rooms, currentRoom) // Clean up empty room + } + log.Debugf("Program %p (user: %s) moved from room %s", prog, userID, currentRoom) } - if len(updatedRoomProgs) > 0 { - a.rooms[currentRoom] = updatedRoomProgs - } else { - delete(a.rooms, currentRoom) - } - log.Debugf("Program %p (user: %s) moved from room %s", prog, userID, currentRoom) } // Ensure the new room exists in the map @@ -236,34 +242,34 @@ 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)) // Allocate with capacity + recipients = make([]*tea.Program, 0, len(progs)) for _, p := range progs { - if p != prog { // Don't send to the exiting program + if p != prog { recipients = append(recipients, p) } } } - // Now, update the room structure var updatedRoomProgs []*tea.Program - for _, pInRoom := range a.rooms[info.roomName] { - if pInRoom != prog { - updatedRoomProgs = append(updatedRoomProgs, pInRoom) + if progsInRoom, ok := a.rooms[info.roomName]; ok { + for _, pInRoom := range progsInRoom { + if pInRoom != prog { + updatedRoomProgs = append(updatedRoomProgs, pInRoom) + } } } + if len(updatedRoomProgs) > 0 { a.rooms[info.roomName] = updatedRoomProgs } else { - delete(a.rooms, info.roomName) // Clean up empty room + delete(a.rooms, info.roomName) } - a.activePrograms[prog] = programInfo{roomName: "", userID: userID} // Update program's state - a.mu.Unlock() // Unlock before broadcasting + a.activePrograms[prog] = programInfo{roomName: "", userID: userID} + 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}) @@ -288,40 +294,38 @@ 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 { // Don't collect the disconnecting program + if p != prog { leaveRecipients = append(leaveRecipients, p) } } } } - // Remove from room tracking if roomName != "" { var updatedRoomProgs []*tea.Program - for _, pInRoom := range a.rooms[roomName] { - if pInRoom != prog { - updatedRoomProgs = append(updatedRoomProgs, pInRoom) + if progsInRoom, ok := a.rooms[roomName]; ok { + for _, pInRoom := range progsInRoom { + if pInRoom != prog { + updatedRoomProgs = append(updatedRoomProgs, pInRoom) + } } } if len(updatedRoomProgs) > 0 { a.rooms[roomName] = updatedRoomProgs } else { - delete(a.rooms, roomName) // Clean up room if it becomes empty + delete(a.rooms, roomName) } 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}) @@ -331,24 +335,27 @@ func (a *app) RemoveProgram(prog *tea.Program) { log.Infof("Program %p (user: %s) removed from active list (full disconnect).", prog, userID) } -// --- main, model types, and functions continue below --- func main() { // Optional: Setup file logging - // file, err := os.OpenFile("chatter_server.log", os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0666) + // logFile, err := os.OpenFile("chatter_server.log", os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0666) // if err == nil { - // log.Default().SetOutput(io.MultiWriter(os.Stderr, file)) // Log to both stderr and file - // defer file.Close() + // defer logFile.Close() + // // You can use a multiwriter to log to both stderr and file if desired + // // log.Default().SetOutput(io.MultiWriter(os.Stderr, logFile)) + // log.Default().SetOutput(logFile) // Log only to file + // log.SetTimeFormat(time.RFC3339) + // log.SetLevel(log.DebugLevel) // Set desired log level + // } else { + // log.Error("Failed to open log file", "error", err) // } - // log.SetLevel(log.DebugLevel) - // log.SetTimeFormat(time.Kitchen) app := newApp() app.Start() } type ( - errMsg error - chatMsg struct { + errMsg error // Used for tea.Cmd to send errors to Update + chatMsg struct { // Used for tea.Cmd to send chat messages to Update id string text string room string @@ -363,7 +370,7 @@ const ( phaseChat ) -var roomNameRegex = regexp.MustCompile("^[a-zA-Z0-9]{1,24}$") +var roomNameRegex = regexp.MustCompile("^[a-zA-Z0-9]{1," + fmt.Sprintf("%d", maxRoomNameLength) + "}$") func isValidRoomName(name string) bool { return roomNameRegex.MatchString(name) @@ -371,18 +378,18 @@ func isValidRoomName(name string) bool { type model struct { app *app - program *tea.Program + program *tea.Program // Reference to the Bubble Tea program managing this model phase modelPhase roomInput textinput.Model currentRoom string id string // User ID from SSH session viewport viewport.Model - messages []string + messages []string // Chat messages displayed in the viewport textarea textarea.Model senderStyle lipgloss.Style - systemMessageStyle lipgloss.Style // For join/leave/command output/errors - messageTimestamps []time.Time // For rate limiting - err error + systemMessageStyle lipgloss.Style // For join/leave/command output (not for m.err) + messageTimestamps []time.Time // For rate limiting user messages + err error // For displaying transient input errors to the user (e.g., rate limit, too long) } func initialModel(appInstance *app, userID string) *model { @@ -395,14 +402,14 @@ 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 = maxChatMessageLength // Updated to new max length - ta.SetWidth(60) - ta.SetHeight(3) - ta.FocusedStyle.CursorLine = lipgloss.NewStyle() + ta.CharLimit = maxChatMessageLength // User-facing limit in textarea + ta.SetWidth(60) // Example width, adjust as needed + ta.SetHeight(3) // Multi-line input area + ta.FocusedStyle.CursorLine = lipgloss.NewStyle() // No special styling for cursor line ta.ShowLineNumbers = false - ta.KeyMap.InsertNewline.SetEnabled(false) + ta.KeyMap.InsertNewline.SetEnabled(false) // Use Enter to send, not for newlines in message - vp := viewport.New(60, 15) + vp := viewport.New(60, 15) // Example dimensions, adjust as needed vp.SetContent("Welcome! Please enter a room name to join or create.") return &model{ @@ -414,9 +421,9 @@ func initialModel(appInstance *app, userID string) *model { 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 - err: nil, + senderStyle: lipgloss.NewStyle().Foreground(lipgloss.Color("12")), // Light Blue for sender ID + systemMessageStyle: lipgloss.NewStyle().Foreground(lipgloss.AdaptiveColor{Light: "#909090", Dark: "#6C6C6C"}), // Grey for system messages + err: nil, // No error initially } } @@ -441,35 +448,34 @@ func (m *model) Update(msg tea.Msg) (tea.Model, tea.Cmd) { switch m.phase { case phaseRoomInput: - switch msg := msg.(type) { + switch keyEvent := msg.(type) { case tea.KeyMsg: - switch msg.Type { + switch keyEvent.Type { case tea.KeyCtrlC, tea.KeyEsc: if m.program != nil { - m.app.RemoveProgram(m.program) + m.app.RemoveProgram(m.program) // Deregister before quitting } return m, tea.Quit case tea.KeyEnter: roomName := strings.TrimSpace(m.roomInput.Value()) - if roomName == "" { - roomName = generateRandomRoomName(maxRoomNameLength) - m.roomInput.SetValue(roomName) // Show generated name in input briefly + if roomName == "" { // Generate random room name if input is empty + roomName = generateRandomRoomName(8) // Shorter random name for readability + m.roomInput.SetValue(roomName) // Show generated name in input briefly log.Infof("User %s generated random room name: %s", m.id, roomName) } if isValidRoomName(roomName) { m.currentRoom = roomName m.phase = phaseChat - m.err = nil + m.err = nil // Clear any previous error from room input phase if m.program == nil { log.Error("CRITICAL: m.program is nil when trying to assign to room.") m.err = errors.New("internal setup error: program reference missing") - return m, nil + return m, nil // Stay in current phase with error displayed } m.app.AssignProgramToRoom(m.program, m.currentRoom) - // Send join message joinMsgText := fmt.Sprintf("%s has joined the room.", m.id) m.app.BroadcastToRoom(m.currentRoom, chatMsg{ id: "system", @@ -480,48 +486,62 @@ func (m *model) Update(msg tea.Msg) (tea.Model, tea.Cmd) { m.textarea.Focus() m.roomInput.Blur() + // Initial welcome message for the room, not appended to m.messages directly to avoid duplication m.viewport.SetContent(fmt.Sprintf("Welcome to room '%s'! (User: %s)\nType /help for commands.", m.currentRoom, m.id)) - m.messages = []string{} // Clear previous messages if any + m.messages = []string{} // Clear message history from potential previous room/session cmds = append(cmds, textarea.Blink) } else { - m.err = errors.New("invalid room name: 1-24 alphanumeric characters required") + m.err = fmt.Errorf("invalid room name: use 1-%d alphanumeric characters", maxRoomNameLength) m.roomInput.Focus() // Keep focus on room input if error cmds = append(cmds, textinput.Blink) } return m, tea.Batch(cmds...) } - case errMsg: - m.err = msg + case errMsg: // This would be an error from a command, not typically used in room input phase + m.err = keyEvent return m, nil } m.roomInput, cmd = m.roomInput.Update(msg) cmds = append(cmds, cmd) case phaseChat: - var tiCmd tea.Cmd - var vpCmd tea.Cmd + // Order of updates: components first, then specific message handling. + m.textarea, cmd = m.textarea.Update(msg) + cmds = append(cmds, cmd) + m.viewport, cmd = m.viewport.Update(msg) // Handles scrolling based on content changes + cmds = append(cmds, 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) // Handles scrolling, etc. - cmds = append(cmds, tiCmd, vpCmd) - - switch msg := msg.(type) { + switch event := msg.(type) { case tea.KeyMsg: - switch msg.Type { + // If a key is pressed (and it's not Enter potentially causing a new error), + // and there was a previous user input error (m.err), clear it. + if m.err != nil && event.Type != tea.KeyEnter { + isTextModifyingKey := false + switch event.Type { + case tea.KeyRunes, tea.KeySpace, tea.KeyBackspace, tea.KeyDelete: + isTextModifyingKey = true + } + if isTextModifyingKey { + m.err = nil // Clear error as user types new input + } + } + + switch event.Type { case tea.KeyCtrlC, tea.KeyEsc: if m.program != nil { m.app.RemoveProgram(m.program) } return m, tea.Quit + case tea.KeyEnter: val := strings.TrimSpace(m.textarea.Value()) - m.textarea.Reset() // Reset textarea immediately + m.textarea.Reset() // Reset textarea immediately after getting value if strings.HasPrefix(val, "/") { + m.err = nil // Clear previous input errors when a command is attempted parts := strings.Fields(val) command := parts[0] - // args := parts[1:] // For future use if commands take arguments + // args := parts[1:] // For future use switch command { case "/help": @@ -537,7 +557,7 @@ func (m *model) Update(msg tea.Msg) (tea.Model, tea.Cmd) { m.phase = phaseRoomInput m.currentRoom = "" m.messages = []string{} - m.err = nil + m.err = nil // Ensure error is clear when changing phase m.roomInput.Reset() m.roomInput.Placeholder = "Enter another room or press Enter for random..." m.roomInput.Focus() @@ -552,88 +572,83 @@ 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." + // This case (only one here) should be handled by GetUserIDsInRoom including self. + userListText = "It seems you are the only one here." } m.messages = append(m.messages, m.systemMessageStyle.Render(userListText)) } else { + // This should not happen if in phaseChat, currentRoom should be set. m.messages = append(m.messages, m.systemMessageStyle.Render("Error: Not in a room.")) } default: m.messages = append(m.messages, m.systemMessageStyle.Render("Unknown command: "+command)) } + // Update viewport with command output m.viewport.SetContent(strings.Join(m.messages, "\n")) m.viewport.GotoBottom() - } else if val != "" { // Regular chat message + } else if val != "" { // Regular chat message attempt // 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 - } + m.err = fmt.Errorf("message too long (max %d chars). Not sent.", maxChatMessageLength) + } else { + // 2. Rate Limiting Check + now := time.Now() + // Remove timestamps older than the window + var recentTimestamps []time.Time + for _, ts := range m.messageTimestamps { + if now.Sub(ts) < rateLimitWindow { + recentTimestamps = append(recentTimestamps, ts) + } + } + m.messageTimestamps = recentTimestamps - // 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) + if len(m.messageTimestamps) >= rateLimitMaxMessages { + m.err = fmt.Errorf("too many messages (max %d per %s). Wait. Not sent.", rateLimitMaxMessages, rateLimitWindow.String()) + } else { + // SUCCESSFUL SEND + m.err = nil // Clear any prior input error + m.messageTimestamps = append(m.messageTimestamps, now) // Add current message's timestamp + + m.app.BroadcastToRoom(m.currentRoom, chatMsg{ + id: m.id, + text: val, + room: m.currentRoom, + isSystemMessage: false, + }, nil) // Broadcast to all, including self (sender) } } - 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) + } else { // val == "" (empty input submitted via Enter) + m.err = nil // Clear any previous error, do nothing for empty message } - // After sending or command, ensure textarea is focused for next input + // After sending, command, or empty Enter, ensure textarea is focused for next input cmds = append(cmds, m.textarea.Focus(), textarea.Blink) - } + + } // end switch event.Type (inner for KeyMsg) case chatMsg: // Received a message (could be own, from others, or system) - if msg.room == m.currentRoom { // Process only if message is for the current room + // Process only if message is for the current room + if event.room == m.currentRoom { var styledMessage string - if msg.isSystemMessage { - styledMessage = m.systemMessageStyle.Render(msg.text) + if event.isSystemMessage { + styledMessage = m.systemMessageStyle.Render(event.text) } else { - styledId := m.senderStyle.Render(msg.id) - styledMessage = styledId + ": " + msg.text + styledId := m.senderStyle.Render(event.id) + // Handle potential multi-line messages by ensuring they are displayed correctly + messageText := strings.ReplaceAll(event.text, "\n", "\n ") // Indent subsequent lines + styledMessage = styledId + ": " + messageText } m.messages = append(m.messages, styledMessage) m.viewport.SetContent(strings.Join(m.messages, "\n")) m.viewport.GotoBottom() - } 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) + } else if event.room != "" && !event.isSystemMessage { + // Log if a non-system message for a different room is somehow received by this model's Update + log.Warnf("Client '%s' in room '%s' received non-system message for different room '%s': %s", m.id, m.currentRoom, event.room, event.text) } - case errMsg: - m.err = msg - log.Error("Error in chat model", "user", m.id, "room", m.currentRoom, "error", msg) - // Optionally display this error in the viewport too - // 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 // Or tea.Batch(cmds...) if other cmds are important + case errMsg: // This is for system/internal errors passed as errMsg type to Update + m.err = event // Display this error using the m.err mechanism + log.Error("Internal error in model", "user", m.id, "room", m.currentRoom, "error", event) } } return m, tea.Batch(cmds...) @@ -642,14 +657,14 @@ 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")) // Purple - errorStyle := lipgloss.NewStyle().Foreground(lipgloss.Color("9")).Bold(true) // Red + titleStyle := lipgloss.NewStyle().Bold(true).Foreground(lipgloss.Color("63")) // Purple for titles + errorStyle := lipgloss.NewStyle().Foreground(lipgloss.Color("9")).Bold(true) // Red for m.err switch m.phase { case phaseRoomInput: viewBuilder.WriteString(titleStyle.Render("chatter.sh") + "\n\n") viewBuilder.WriteString("Please enter a room name to join or create.\n") - viewBuilder.WriteString("Rules: 1-24 characters, letters (a-z, A-Z) and numbers (0-9) only.\n") + viewBuilder.WriteString(fmt.Sprintf("Rules: 1-%d characters, letters (a-z, A-Z) and numbers (0-9) only.\n", maxRoomNameLength)) viewBuilder.WriteString("(Leave empty and press Enter for a random room name.)\n\n") viewBuilder.WriteString(m.roomInput.View()) viewBuilder.WriteString("\n\n(Ctrl+C or Esc to quit)") @@ -664,6 +679,7 @@ func (m *model) View() string { viewBuilder.WriteString("Unknown application state. This is a bug.") } + // Display m.err if it's set (for transient input errors like rate limit, too long message) if m.err != nil { viewBuilder.WriteString("\n\n" + errorStyle.Render("Error: "+m.err.Error())) }