451 lines
12 KiB
Go
451 lines
12 KiB
Go
package main
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"fmt"
|
|
"net"
|
|
"os"
|
|
"os/signal"
|
|
"regexp"
|
|
"strings"
|
|
"sync"
|
|
"syscall"
|
|
"time"
|
|
|
|
"github.com/charmbracelet/bubbles/textarea"
|
|
"github.com/charmbracelet/bubbles/textinput"
|
|
"github.com/charmbracelet/bubbles/viewport"
|
|
tea "github.com/charmbracelet/bubbletea"
|
|
"github.com/charmbracelet/lipgloss"
|
|
"github.com/charmbracelet/log"
|
|
"github.com/charmbracelet/ssh"
|
|
"github.com/charmbracelet/wish"
|
|
"github.com/charmbracelet/wish/activeterm"
|
|
"github.com/charmbracelet/wish/bubbletea"
|
|
"github.com/charmbracelet/wish/logging"
|
|
"github.com/muesli/termenv"
|
|
)
|
|
|
|
const (
|
|
host = "0.0.0.0"
|
|
port = "23234"
|
|
)
|
|
|
|
type programInfo struct {
|
|
roomName string
|
|
userID string
|
|
}
|
|
|
|
// app contains a wish server and manages chat rooms and programs.
|
|
type app struct {
|
|
*ssh.Server
|
|
activePrograms map[*tea.Program]programInfo // program -> programInfo
|
|
rooms map[string][]*tea.Program // roomName -> list of programs in that room
|
|
mu sync.Mutex
|
|
}
|
|
|
|
// BroadcastToRoom sends a message to all running programs in a specific room.
|
|
func (a *app) BroadcastToRoom(roomName string, msg tea.Msg) {
|
|
a.mu.Lock()
|
|
defer a.mu.Unlock()
|
|
|
|
if progsInRoom, ok := a.rooms[roomName]; ok {
|
|
log.Debugf("Broadcasting to room %s (%d programs): %+v", roomName, len(progsInRoom), msg)
|
|
for _, p := range progsInRoom {
|
|
go p.Send(msg)
|
|
}
|
|
} else {
|
|
log.Warnf("Attempted to broadcast to non-existent or empty room: %s", roomName)
|
|
}
|
|
}
|
|
|
|
func newApp() *app {
|
|
a := &app{
|
|
activePrograms: make(map[*tea.Program]programInfo), // Initialize correctly
|
|
rooms: make(map[string][]*tea.Program),
|
|
}
|
|
s, err := wish.NewServer(
|
|
wish.WithAddress(net.JoinHostPort(host, port)),
|
|
wish.WithHostKeyPath(".ssh/id_ed25519"), // Ensure this key exists (e.g., run ssh-keygen -t ed25519 -f .ssh/id_ed25519 -N "")
|
|
wish.WithMiddleware(
|
|
bubbletea.MiddlewareWithProgramHandler(a.ProgramHandler, termenv.ANSI256),
|
|
activeterm.Middleware(),
|
|
logging.Middleware(),
|
|
),
|
|
)
|
|
if err != nil {
|
|
log.Fatal("Could not initialize server", "error", err)
|
|
}
|
|
|
|
a.Server = s
|
|
return a
|
|
}
|
|
|
|
func (a *app) Start() {
|
|
var err error
|
|
done := make(chan os.Signal, 1)
|
|
signal.Notify(done, os.Interrupt, syscall.SIGINT, syscall.SIGTERM)
|
|
log.Info("Starting SSH server", "host", host, "port", port)
|
|
go func() {
|
|
if err = a.ListenAndServe(); err != nil && !errors.Is(err, ssh.ErrServerClosed) {
|
|
log.Error("Could not start server", "error", err)
|
|
done <- syscall.SIGTERM // Signal main to exit if server fails
|
|
}
|
|
}()
|
|
|
|
<-done
|
|
log.Info("Stopping SSH server...")
|
|
ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
|
|
defer func() {
|
|
cancel()
|
|
// Clean up internal state after server shutdown
|
|
a.mu.Lock()
|
|
a.activePrograms = make(map[*tea.Program]programInfo) // Correct type
|
|
a.rooms = make(map[string][]*tea.Program)
|
|
a.mu.Unlock()
|
|
log.Info("Cleared app program stores.")
|
|
}()
|
|
|
|
if err := a.Shutdown(ctx); err != nil {
|
|
log.Error("Could not gracefully stop server", "error", err)
|
|
}
|
|
log.Info("SSH server stopped.")
|
|
}
|
|
|
|
func (a *app) ProgramHandler(s ssh.Session) *tea.Program {
|
|
userID := s.User()
|
|
m := initialModel(a, userID) // initialModel returns *model
|
|
|
|
p := tea.NewProgram(m, bubbletea.MakeOptions(s)...)
|
|
|
|
m.SetProgram(p) // Give model a reference to its own program
|
|
a.addProgram(p, userID) // Register program with the app, passing userID
|
|
|
|
return p
|
|
}
|
|
|
|
func (a *app) addProgram(p *tea.Program, userID string) {
|
|
a.mu.Lock()
|
|
defer a.mu.Unlock()
|
|
a.activePrograms[p] = programInfo{roomName: "", userID: userID} // Store userID
|
|
log.Debugf("Program %p added for user %s (unassigned)", p, userID)
|
|
}
|
|
|
|
func (a *app) AssignProgramToRoom(prog *tea.Program, roomName string) {
|
|
a.mu.Lock()
|
|
defer a.mu.Unlock()
|
|
|
|
currentInfo, exists := a.activePrograms[prog]
|
|
if !exists {
|
|
// userID is NOT available here if the program isn't in activePrograms.
|
|
// NO prog.Model() call should be here.
|
|
log.Warnf("Program %p not found in activePrograms, cannot assign to room %s", prog, roomName)
|
|
return
|
|
}
|
|
// Correct: UserID is obtained from the stored currentInfo.
|
|
userID := currentInfo.userID
|
|
currentRoom := currentInfo.roomName
|
|
|
|
if currentRoom != "" && currentRoom != roomName {
|
|
var updatedRoomProgs []*tea.Program
|
|
for _, pInRoom := range a.rooms[currentRoom] {
|
|
if pInRoom != prog {
|
|
updatedRoomProgs = append(updatedRoomProgs, pInRoom)
|
|
}
|
|
}
|
|
if len(updatedRoomProgs) > 0 {
|
|
a.rooms[currentRoom] = updatedRoomProgs
|
|
} else {
|
|
delete(a.rooms, currentRoom)
|
|
}
|
|
// Correct: userID is used from the variable. NO prog.Model() here.
|
|
log.Debugf("Program %p (user: %s) moved from room %s", prog, userID, currentRoom)
|
|
}
|
|
|
|
if _, ok := a.rooms[roomName]; !ok {
|
|
a.rooms[roomName] = []*tea.Program{}
|
|
}
|
|
|
|
foundInNewRoomList := false
|
|
for _, pInList := range a.rooms[roomName] {
|
|
if pInList == prog {
|
|
foundInNewRoomList = true
|
|
break
|
|
}
|
|
}
|
|
if !foundInNewRoomList {
|
|
a.rooms[roomName] = append(a.rooms[roomName], prog)
|
|
}
|
|
|
|
a.activePrograms[prog] = programInfo{roomName: roomName, userID: userID}
|
|
// Correct: userID is used from the variable. NO prog.Model() here.
|
|
// If your error is on this line (or similar), ensure 'userID' is not derived from prog.Model().
|
|
log.Infof("Program %p assigned/confirmed in room %s (user: %s)", prog, roomName, userID)
|
|
}
|
|
|
|
func (a *app) RemoveProgram(prog *tea.Program) {
|
|
a.mu.Lock()
|
|
defer a.mu.Unlock()
|
|
|
|
info, exists := a.activePrograms[prog]
|
|
if !exists {
|
|
// Correct: NO prog.Model() call here.
|
|
log.Warnf("Attempted to remove unknown or already removed program: %p", prog)
|
|
return
|
|
}
|
|
// Correct: UserID and roomName are obtained from the stored info.
|
|
userID := info.userID
|
|
roomName := info.roomName
|
|
|
|
if roomName != "" {
|
|
var updatedRoomProgs []*tea.Program
|
|
for _, pInRoom := range a.rooms[roomName] {
|
|
if pInRoom != prog {
|
|
updatedRoomProgs = append(updatedRoomProgs, pInRoom)
|
|
}
|
|
}
|
|
if len(updatedRoomProgs) > 0 {
|
|
a.rooms[roomName] = updatedRoomProgs
|
|
} else {
|
|
delete(a.rooms, roomName)
|
|
}
|
|
// Correct: userID is used from the variable. NO prog.Model() here.
|
|
// If your error is on this line, ensure 'userID' is not derived from prog.Model().
|
|
log.Infof("Program %p removed from room %s (user: %s)", prog, roomName, userID)
|
|
} else {
|
|
// Correct: userID is used from the variable. NO prog.Model() here.
|
|
log.Debugf("Program %p was not in any room, removing from active list (user: %s).", prog, userID)
|
|
}
|
|
|
|
delete(a.activePrograms, prog)
|
|
// Correct: userID is used from the variable. NO prog.Model() here.
|
|
log.Infof("Program %p removed from active list (user: %s).", prog, userID)
|
|
}
|
|
|
|
func main() {
|
|
// For more detailed logs:
|
|
// file, err := os.OpenFile("chat_server.log", os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0666)
|
|
// if err == nil {
|
|
// log.Default().SetOutput(file)
|
|
// }
|
|
// log.SetLevel(log.DebugLevel) // Show debug messages
|
|
|
|
app := newApp()
|
|
app.Start()
|
|
}
|
|
|
|
type (
|
|
errMsg error
|
|
chatMsg struct {
|
|
id string
|
|
text string
|
|
room string // So message carries its intended room
|
|
}
|
|
)
|
|
|
|
type modelPhase int
|
|
|
|
const (
|
|
phaseRoomInput modelPhase = iota
|
|
phaseChat
|
|
)
|
|
|
|
var roomNameRegex = regexp.MustCompile("^[a-zA-Z0-9]{1,24}$")
|
|
|
|
func isValidRoomName(name string) bool {
|
|
return roomNameRegex.MatchString(name)
|
|
}
|
|
|
|
type model struct {
|
|
app *app // Pointer to the main application state
|
|
program *tea.Program // Pointer to this model's own tea.Program
|
|
phase modelPhase
|
|
roomInput textinput.Model
|
|
currentRoom string
|
|
id string // User ID from SSH session
|
|
|
|
// Chat fields
|
|
viewport viewport.Model
|
|
messages []string
|
|
textarea textarea.Model
|
|
senderStyle lipgloss.Style
|
|
err error
|
|
}
|
|
|
|
func initialModel(appInstance *app, userID string) *model {
|
|
roomTi := textinput.New()
|
|
roomTi.Placeholder = "e.g., 'dev-chat' or 'gaming'"
|
|
roomTi.Focus()
|
|
roomTi.CharLimit = 24
|
|
roomTi.Width = 40
|
|
|
|
ta := textarea.New()
|
|
ta.Placeholder = "Type your message here and press Enter..."
|
|
ta.Prompt = "┃ "
|
|
ta.CharLimit = 280
|
|
ta.SetWidth(60)
|
|
ta.SetHeight(3)
|
|
ta.FocusedStyle.CursorLine = lipgloss.NewStyle()
|
|
ta.ShowLineNumbers = false
|
|
ta.KeyMap.InsertNewline.SetEnabled(false)
|
|
|
|
vp := viewport.New(60, 15)
|
|
vp.SetContent("Welcome! Please enter a room name to join or create.")
|
|
|
|
return &model{
|
|
app: appInstance,
|
|
id: userID,
|
|
phase: phaseRoomInput,
|
|
roomInput: roomTi,
|
|
textarea: ta,
|
|
messages: []string{},
|
|
viewport: vp,
|
|
senderStyle: lipgloss.NewStyle().Foreground(lipgloss.Color("12")),
|
|
err: nil,
|
|
}
|
|
}
|
|
|
|
func (m *model) SetProgram(p *tea.Program) {
|
|
m.program = p
|
|
}
|
|
|
|
func (m *model) Init() tea.Cmd {
|
|
switch m.phase {
|
|
case phaseRoomInput:
|
|
return textinput.Blink
|
|
case phaseChat:
|
|
return textarea.Blink
|
|
default:
|
|
return nil
|
|
}
|
|
}
|
|
|
|
func (m *model) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
|
|
var cmds []tea.Cmd
|
|
var cmd tea.Cmd
|
|
|
|
switch m.phase {
|
|
case phaseRoomInput:
|
|
switch msg := msg.(type) {
|
|
case tea.KeyMsg:
|
|
switch msg.Type {
|
|
case tea.KeyCtrlC, tea.KeyEsc:
|
|
if m.program != nil {
|
|
m.app.RemoveProgram(m.program)
|
|
}
|
|
return m, tea.Quit
|
|
case tea.KeyEnter:
|
|
roomName := strings.TrimSpace(m.roomInput.Value())
|
|
if isValidRoomName(roomName) {
|
|
m.currentRoom = roomName
|
|
m.phase = phaseChat
|
|
m.err = nil
|
|
|
|
if m.program == nil {
|
|
log.Error("CRITICAL: m.program is nil when trying to assign to room. This should not happen.")
|
|
m.err = errors.New("internal setup error: program reference missing")
|
|
return m, nil
|
|
}
|
|
m.app.AssignProgramToRoom(m.program, m.currentRoom)
|
|
|
|
m.textarea.Focus()
|
|
m.roomInput.Blur()
|
|
|
|
m.viewport.SetContent(fmt.Sprintf("Welcome to room '%s'! (User: %s)\nMessages will appear below.", m.currentRoom, m.id))
|
|
m.messages = []string{}
|
|
cmds = append(cmds, textarea.Blink)
|
|
} else {
|
|
m.err = errors.New("invalid room name: 1-24 alphanumeric characters required")
|
|
m.roomInput.Focus()
|
|
cmds = append(cmds, textinput.Blink)
|
|
}
|
|
return m, tea.Batch(cmds...)
|
|
}
|
|
case errMsg:
|
|
m.err = msg
|
|
return m, nil
|
|
}
|
|
|
|
m.roomInput, cmd = m.roomInput.Update(msg)
|
|
cmds = append(cmds, cmd)
|
|
|
|
case phaseChat:
|
|
var tiCmd tea.Cmd
|
|
var vpCmd tea.Cmd
|
|
|
|
m.textarea, tiCmd = m.textarea.Update(msg)
|
|
m.viewport, vpCmd = m.viewport.Update(msg)
|
|
cmds = append(cmds, tiCmd, vpCmd)
|
|
|
|
switch msg := msg.(type) {
|
|
case tea.KeyMsg:
|
|
switch msg.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())
|
|
if val != "" {
|
|
m.app.BroadcastToRoom(m.currentRoom, chatMsg{
|
|
id: m.id,
|
|
text: val,
|
|
room: m.currentRoom,
|
|
})
|
|
m.textarea.Reset()
|
|
m.viewport.GotoBottom()
|
|
}
|
|
}
|
|
case chatMsg:
|
|
if msg.room == m.currentRoom {
|
|
styledId := m.senderStyle.Render(msg.id)
|
|
if msg.id == m.id {
|
|
styledId = lipgloss.NewStyle().Foreground(lipgloss.Color("10")).Render(msg.id + " (you)")
|
|
}
|
|
m.messages = append(m.messages, styledId+": "+msg.text)
|
|
m.viewport.SetContent(strings.Join(m.messages, "\n"))
|
|
m.viewport.GotoBottom()
|
|
} else {
|
|
log.Warnf("Client in room '%s' (user '%s') incorrectly received msg for room '%s'. Ignored.", m.currentRoom, m.id, msg.room)
|
|
}
|
|
case errMsg:
|
|
m.err = msg
|
|
log.Error("Error in chat model", "user", m.id, "room", m.currentRoom, "error", msg)
|
|
return m, nil
|
|
}
|
|
}
|
|
return m, tea.Batch(cmds...)
|
|
}
|
|
|
|
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)
|
|
|
|
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\n")
|
|
viewBuilder.WriteString(m.roomInput.View())
|
|
viewBuilder.WriteString("\n\n(Ctrl+C or Esc to quit)")
|
|
case phaseChat:
|
|
header := fmt.Sprintf("Room: %s / User: %s", m.currentRoom, m.id)
|
|
viewBuilder.WriteString(titleStyle.Render(header) + "\n")
|
|
viewBuilder.WriteString(lipgloss.NewStyle().BorderStyle(lipgloss.NormalBorder()).BorderBottom(true).Width(m.viewport.Width).Render(m.viewport.View()) + "\n\n")
|
|
viewBuilder.WriteString(m.textarea.View())
|
|
viewBuilder.WriteString("\n\n(Ctrl+C or Esc to quit)")
|
|
default:
|
|
viewBuilder.WriteString("Unknown application state. This is a bug.")
|
|
}
|
|
|
|
if m.err != nil {
|
|
viewBuilder.WriteString("\n\n" + errorStyle.Render("Error: "+m.err.Error()))
|
|
}
|
|
|
|
return viewBuilder.String() + "\n"
|
|
}
|