diff --git a/internal/database.go b/internal/database.go index 966493a..e764152 100644 --- a/internal/database.go +++ b/internal/database.go @@ -48,7 +48,6 @@ func InitDB() { FOREIGN KEY (room_id) REFERENCES rooms(id), FOREIGN KEY (username) REFERENCES users(username) ); - -- Stores the encrypted room key for a specific user for a specific epoch CREATE TABLE IF NOT EXISTS user_room_keys ( username TEXT, room_id TEXT, @@ -132,6 +131,12 @@ func getRoomIDByName(name string) (string, error) { return id, err } +func getRoomCreator(roomID string) (string, error) { + var creator string + err := db.QueryRow("SELECT creator FROM rooms WHERE id = ?", roomID).Scan(&creator) + return creator, err +} + func getRoomCurrentEpoch(roomID string) (int, error) { var epoch int err := db.QueryRow("SELECT current_epoch FROM rooms WHERE id = ?", roomID).Scan(&epoch) @@ -140,7 +145,6 @@ func getRoomCurrentEpoch(roomID string) (int, error) { func incrementRoomEpoch(roomID string) (int, error) { var newEpoch int - // Atomic increment err := db.QueryRow("UPDATE rooms SET current_epoch = current_epoch + 1 WHERE id = ? RETURNING current_epoch", roomID).Scan(&newEpoch) return newEpoch, err } @@ -175,7 +179,6 @@ func leaveRoom(roomID, username string) error { } func joinRoomMember(roomID, username string) error { - // Check if already joined var count int db.QueryRow("SELECT COUNT(*) FROM room_members WHERE room_id = ? AND username = ?", roomID, username).Scan(&count) if count > 0 { diff --git a/internal/model.go b/internal/model.go index af1d85f..18dae1e 100644 --- a/internal/model.go +++ b/internal/model.go @@ -17,13 +17,13 @@ import ( ) var ( - msgStyle = lipgloss.NewStyle().Foreground(lipgloss.Color("252")).PaddingLeft(1) - senderStyle = lipgloss.NewStyle().Foreground(lipgloss.Color("205")).Bold(true) - errStyle = lipgloss.NewStyle().Foreground(lipgloss.Color("196")).Bold(true) - sysStyle = lipgloss.NewStyle().Foreground(lipgloss.Color("117")).Italic(true) - roomStyle = lipgloss.NewStyle().Foreground(lipgloss.Color("86")).Bold(true) - dmStyle = lipgloss.NewStyle().Foreground(lipgloss.Color("141")).Bold(true) - noteStyle = lipgloss.NewStyle().Foreground(lipgloss.Color("228")).Bold(true) + msgStyle = lipgloss.NewStyle().Foreground(lipgloss.Color("252")).PaddingLeft(1) + senderStyle = lipgloss.NewStyle().Foreground(lipgloss.Color("205")).Bold(true) + errStyle = lipgloss.NewStyle().Foreground(lipgloss.Color("196")).Bold(true) + sysStyle = lipgloss.NewStyle().Foreground(lipgloss.Color("117")).Italic(true) + roomStyle = lipgloss.NewStyle().Foreground(lipgloss.Color("86")).Bold(true) + dmStyle = lipgloss.NewStyle().Foreground(lipgloss.Color("141")).Bold(true) + noteStyle = lipgloss.NewStyle().Foreground(lipgloss.Color("228")).Bold(true) commandStyle = lipgloss.NewStyle().Foreground(lipgloss.Color("39")).Italic(true) timeStyle = lipgloss.NewStyle().Foreground(lipgloss.Color("246")) lockedStyle = lipgloss.NewStyle().Foreground(lipgloss.Color("240")).Italic(true) @@ -51,8 +51,7 @@ type model struct { currentRoomID string currentRoomEpoch int - // We cache keys: Epoch -> Plaintext Key - roomKeyCache map[int][]byte + roomKeyCache map[int][]byte rooms []struct { ID string @@ -143,6 +142,11 @@ func (m model) Update(msg tea.Msg) (tea.Model, tea.Cmd) { return m, tea.ClearScreen case tea.KeyMsg: + // Clear any displayed errors on user interaction + if m.err != nil { + m.err = nil + } + switch msg.Type { case tea.KeyCtrlC, tea.KeyEsc: return m, tea.Quit @@ -226,23 +230,21 @@ func (m model) Update(msg tea.Msg) (tea.Model, tea.Cmd) { } if val == "/leave" { + creator, _ := getRoomCreator(m.currentRoomID) + if creator == m.username { + m.saveMessage("System", "Error: Owners cannot leave. Use /delete to destroy room.") + m.input.Reset() + return m, nil + } leaveRoom(m.currentRoomID, m.username) - m.rotateRoomKey(m.currentRoomID) // Rotate key on leave + m.rotateRoomKey(m.currentRoomID) m.exitChat() return m, tea.ClearScreen } if val == "/delete" { - isCreator := false - for _, r := range m.rooms { - if r.ID == m.currentRoomID { - if r.Creator == m.username { - isCreator = true - } - break - } - } - if isCreator { + creator, _ := getRoomCreator(m.currentRoomID) + if creator == m.username { deleteRoom(m.currentRoomID) m.exitChat() return m, tea.ClearScreen @@ -289,15 +291,10 @@ func (m *model) loadRooms() { m.rooms = rooms m.viewport.SetContent("") m.input.SetValue("") - m.input.Placeholder = "Enter room # or: /join , /dm , /list" + m.input.Placeholder = "Enter room # or: /join , /dm " } func (m *model) handleRoomListInput(text string) { - if text == "/list" { - m.loadRooms() - return - } - if strings.HasPrefix(text, "/join ") { roomName := strings.TrimPrefix(text, "/join ") m.handleJoinRoom(roomName) @@ -324,18 +321,15 @@ func (m *model) handleRoomListInput(text string) { func (m *model) handleJoinRoom(roomName string) { existingID, err := getRoomIDByName(roomName) if err == nil && existingID != "" { - // Join existing if err := joinRoomMember(existingID, m.username); err != nil { m.err = err return } - // ROTATE KEY so new user gets a key, but doesn't get old keys m.rotateRoomKey(existingID) m.enterRoom(existingID, roomName) return } - // Create new roomID := generateRoomID() if err := createRoom(roomID, roomName, m.username, false); err != nil { m.err = err @@ -345,7 +339,6 @@ func (m *model) handleJoinRoom(roomName string) { m.err = err return } - // Initial Key m.rotateRoomKey(roomID) m.enterRoom(roomID, roomName) } @@ -382,7 +375,7 @@ func (m *model) rotateRoomKey(roomID string) { func (m *model) handleSelectUserForDM(username string) { if username == "/cancel" { m.state = 1 - m.input.Placeholder = "Enter room # or: /join , /dm , /list" + m.input.Placeholder = "Enter room # or: /join , /dm " return } if username == m.username { @@ -390,6 +383,40 @@ func (m *model) handleSelectUserForDM(username string) { return } + // CHECK IF USER EXISTS + identityKey, prekey, prekeySignature, err := getUserKeys(username) + if err != nil { + if err == sql.ErrNoRows { + m.err = fmt.Errorf("user '%s' not found", username) + } else { + m.err = err + } + return + } + + if len(identityKey) == 0 { + m.err = fmt.Errorf("user '%s' has no keys set up", username) + return + } + + // Pre-calc shared secret for DM + var theirIdentityKey, theirPrekey [32]byte + copy(theirIdentityKey[:], identityKey) + copy(theirPrekey[:], prekey) + + ephemeralKey, err := generateIdentityKeyPair() + if err != nil { + m.err = err + return + } + + dh1, _ := performDH(m.identityKey.PrivateKey, theirPrekey) + dh2, _ := performDH(ephemeralKey.PrivateKey, theirIdentityKey) + dh3, _ := performDH(ephemeralKey.PrivateKey, theirPrekey) + + _ = deriveSharedSecret(dh1, dh2, dh3) + _ = prekeySignature + roomID := generateRoomID() roomName := fmt.Sprintf("DM: %s <-> %s", m.username, username) @@ -418,10 +445,17 @@ func (m *model) createNoteToSelf() { func (m *model) enterRoom(roomID, roomName string) { epoch, _ := getRoomCurrentEpoch(roomID) + creator, _ := getRoomCreator(roomID) + m.currentRoomID = roomID m.currentRoomEpoch = epoch m.state = 2 - m.input.Placeholder = fmt.Sprintf("[%s] /back, /leave, /delete", roomName) + + if creator == m.username { + m.input.Placeholder = fmt.Sprintf("[%s] /back to menu, /delete to destroy", roomName) + } else { + m.input.Placeholder = fmt.Sprintf("[%s] /back to menu, /leave to quit", roomName) + } m.loadMessages() } @@ -442,28 +476,23 @@ func generateRoomID() string { } func (m *model) getEpochKey(epoch int) []byte { - // Check cache if key, ok := m.roomKeyCache[epoch]; ok { return key } - encKey, err := getUserRoomKey(m.username, m.currentRoomID, epoch) if err != nil || encKey == nil { return nil } - key, err := DecryptKeyForUser(m.identityKey.PrivateKey, encKey) if err != nil { return nil } - m.roomKeyCache[epoch] = key return key } func (m *model) saveMessage(sender, text string) { m.currentRoomEpoch, _ = getRoomCurrentEpoch(m.currentRoomID) - key := m.getEpochKey(m.currentRoomEpoch) if key == nil { m.err = fmt.Errorf("no key for current epoch") @@ -555,29 +584,38 @@ No public key authentication provided. } if m.needsRegistration { + errStr := "" + if m.err != nil { + errStr = "\n" + errStyle.Render(m.err.Error()) + "\n" + } return fmt.Sprintf(` %s Welcome! You're a new user. Please choose a username to register. - %s -`, sysStyle.Render("SECURE TUI CHAT - REGISTRATION"), m.input.View()) +%s +`, sysStyle.Render("SECURE TUI CHAT - REGISTRATION"), errStr, m.input.View()) } if m.state == 0 { + errStr := "" + if m.err != nil { + errStr = "\n" + errStyle.Render(m.err.Error()) + "\n" + } return fmt.Sprintf(` %s Welcome back, %s. This environment is encrypted at rest. Please enter your passphrase to unlock your keys. - %s -`, sysStyle.Render("SECURE TUI CHAT"), senderStyle.Render(m.username), m.input.View()) +%s +`, sysStyle.Render("SECURE TUI CHAT"), senderStyle.Render(m.username), errStr, m.input.View()) } if m.state == 1 { var b strings.Builder b.WriteString(sysStyle.Render("=== YOUR ROOMS ===") + "\n\n") + if len(m.rooms) == 0 { b.WriteString(sysStyle.Render("No rooms yet.") + "\n") b.WriteString(commandStyle.Render(" /join ") + " - Join/Create room\n") @@ -601,7 +639,11 @@ Please enter your passphrase to unlock your keys. prefix, style.Render(room.Name))) } - b.WriteString("\n" + sysStyle.Render("Commands: ") + commandStyle.Render("/join /dm /list") + "\n") + b.WriteString("\n" + sysStyle.Render("Commands: ") + commandStyle.Render("/join /dm") + "\n") + } + + if m.err != nil { + b.WriteString("\n" + errStyle.Render(m.err.Error())) } b.WriteString("\n" + m.input.View()) return b.String()