improve server relay and client

This commit is contained in:
juancwu 2025-12-07 21:17:52 -05:00
commit fd0e2afddb
6 changed files with 461 additions and 109 deletions

View file

@ -7,66 +7,149 @@ import (
"encoding/json"
"fmt"
"log"
"os"
"sort"
"strings"
"gossip/pkg/protocol"
"github.com/charmbracelet/bubbles/textarea"
"github.com/charmbracelet/bubbles/textinput"
"github.com/charmbracelet/bubbles/viewport"
tea "github.com/charmbracelet/bubbletea"
"github.com/gorilla/websocket"
"golang.org/x/crypto/nacl/box"
)
const (
IdentityFile = "identity.json"
ContactsFile = "contacts.json"
)
type AppMode int
const (
ModeNormal AppMode = iota
ModeInsert
ModeConnect
ModeList
)
type KeyPair struct {
Public *[32]byte
Private *[32]byte
PubHex string
}
type StoredIdentity struct {
PublicKey string `json:"public_key"`
PrivateKey string `json:"private_key"`
}
func loadOrGenerateKeys() (KeyPair, error) {
fileBytes, err := os.ReadFile(IdentityFile)
if err == nil {
var stored StoredIdentity
if err := json.Unmarshal(fileBytes, &stored); err == nil {
pubBytes, _ := hex.DecodeString(stored.PublicKey)
privBytes, _ := hex.DecodeString(stored.PrivateKey)
var pubKey, privKey [32]byte
copy(pubKey[:], pubBytes)
copy(privKey[:], privBytes)
return KeyPair{Public: &pubKey, Private: &privKey, PubHex: stored.PublicKey}, nil
}
}
pub, priv, err := box.GenerateKey(rand.Reader)
if err != nil {
return KeyPair{}, err
}
kp := KeyPair{Public: pub, Private: priv, PubHex: hex.EncodeToString(pub[:])}
saveData := StoredIdentity{PublicKey: kp.PubHex, PrivateKey: hex.EncodeToString(priv[:])}
bytes, _ := json.MarshalIndent(saveData, "", " ")
_ = os.WriteFile(IdentityFile, bytes, 0600)
return kp, nil
}
type ContactBook map[string]string
func loadContacts() ContactBook {
cb := make(ContactBook)
data, err := os.ReadFile(ContactsFile)
if err == nil {
json.Unmarshal(data, &cb)
}
return cb
}
func (cb ContactBook) save() {
data, _ := json.MarshalIndent(cb, "", " ")
os.WriteFile(ContactsFile, data, 0600)
}
type model struct {
conn *websocket.Conn
keys KeyPair
targetHex string
targetPub *[32]byte
viewport viewport.Model
textarea textarea.Model
messages []string
err error
conn *websocket.Conn
keys KeyPair
contacts ContactBook
mode AppMode
myAccountNum string
targetAcc string
targetName string
targetHex string
targetPub *[32]byte
viewport viewport.Model
textarea textarea.Model
connectInput textinput.Model
contactList []string
listCursor int
messages []string
err error
}
type wsMsg protocol.Message
func main() {
pub, priv, err := box.GenerateKey(rand.Reader)
keys, err := loadOrGenerateKeys()
if err != nil {
log.Fatal(err)
}
keys := KeyPair{Public: pub, Private: priv, PubHex: hex.EncodeToString(pub[:])}
contacts := loadContacts()
c, _, err := websocket.DefaultDialer.Dial("ws://localhost:8080/ws", nil)
if err != nil {
log.Fatal("Could not connect to server:", err)
log.Fatal("Could not connect:", err)
}
defer c.Close()
loginMsg := protocol.Message{Type: "login", Sender: keys.PubHex}
c.WriteJSON(loginMsg)
c.WriteJSON(protocol.Message{Type: protocol.TypeLogin, Sender: keys.PubHex})
ta := textarea.New()
ta.Placeholder = "Type a message (or /connect <PUBKEY>)..."
ta.Focus()
ta.SetHeight(2)
ta.Placeholder = "Type message..."
ta.SetHeight(3)
ta.ShowLineNumbers = false
ta.Blur()
ti := textinput.New()
ti.Placeholder = "Enter Name or Account #"
ti.CharLimit = 20
ti.Width = 30
vp := viewport.New(80, 20)
vp.SetContent(fmt.Sprintf("Your ID: %s\nTo start, type: /connect <FRIEND_ID>", keys.PubHex))
vp.SetContent("Welcome. Press 'i' to type, 'c' to connect, 'l' for contacts.")
m := model{
conn: c,
keys: keys,
textarea: ta,
viewport: vp,
conn: c,
keys: keys,
contacts: contacts,
textarea: ta,
connectInput: ti,
viewport: vp,
mode: ModeNormal,
}
p := tea.NewProgram(m)
@ -88,96 +171,239 @@ func main() {
}
}
func (m model) Init() tea.Cmd {
return textarea.Blink
}
func (m model) Init() tea.Cmd { return textarea.Blink }
func (m model) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
var (
tiCmd tea.Cmd
vpCmd tea.Cmd
)
var cmd tea.Cmd
var cmds []tea.Cmd
switch msg := msg.(type) {
case tea.KeyMsg:
switch msg.Type {
case tea.KeyCtrlC, tea.KeyEsc:
return m, tea.Quit
case tea.KeyEnter:
input := m.textarea.Value()
if input == "" {
switch m.mode {
case ModeNormal:
switch msg.String() {
case "ctrl+c", "q":
return m, tea.Quit
case "i":
m.mode = ModeInsert
m.textarea.Focus()
m.textarea.CursorEnd()
return m, nil
}
if dest, ok := strings.CutPrefix(input, "/connect "); ok {
m.targetHex = dest
decoded, _ := hex.DecodeString(dest)
var keyArr [32]byte
copy(keyArr[:], decoded)
m.targetPub = &keyArr
m.messages = append(m.messages, "System: Target set to "+dest[:8]+"...")
case "c":
m.mode = ModeConnect
m.connectInput.Reset()
m.connectInput.Focus()
return m, nil
case "l":
m.mode = ModeList
m.contactList = []string{}
for k := range m.contacts {
m.contactList = append(m.contactList, k)
}
sort.Strings(m.contactList)
m.listCursor = 0
return m, nil
case "d":
m.targetPub = nil
m.targetAcc = ""
m.targetName = ""
m.messages = append(m.messages, "System: Disconnected.")
m.viewport.SetContent(strings.Join(m.messages, "\n"))
m.textarea.Reset()
m.viewport.GotoBottom()
return m, nil
}
if m.targetPub == nil {
m.messages = append(m.messages, "System: No target set! Use /connect first.")
} else {
var nonce [24]byte
rand.Read(nonce[:])
encrypted := box.Seal(nonce[:], []byte(input), &nonce, m.targetPub, m.keys.Private)
b64Content := base64.StdEncoding.EncodeToString(encrypted)
outMsg := protocol.Message{
Type: "msg",
Sender: m.keys.PubHex,
Target: m.targetHex,
Content: b64Content,
case "enter":
input := strings.TrimSpace(m.textarea.Value())
if input == "" {
return m, nil
}
m.conn.WriteJSON(outMsg)
m.messages = append(m.messages, "Me: "+input)
if m.targetPub == nil {
m.messages = append(m.messages, "System: ⛔ Not connected! Press 'c' to connect.")
} else {
var nonce [24]byte
rand.Read(nonce[:])
encrypted := box.Seal(nonce[:], []byte(input), &nonce, m.targetPub, m.keys.Private)
outMsg := protocol.Message{
Type: protocol.TypeMsg,
Sender: m.keys.PubHex,
Target: m.targetHex,
Content: base64.StdEncoding.EncodeToString(encrypted),
}
m.conn.WriteJSON(outMsg)
m.messages = append(m.messages, "Me: "+input)
m.textarea.Reset()
}
m.viewport.SetContent(strings.Join(m.messages, "\n"))
m.viewport.GotoBottom()
}
m.textarea.Reset()
m.viewport.SetContent(strings.Join(m.messages, "\n"))
m.viewport.GotoBottom()
case ModeInsert:
if msg.String() == "esc" {
m.mode = ModeNormal
m.textarea.Blur()
return m, nil
}
case ModeConnect:
if msg.String() == "esc" {
m.mode = ModeNormal
m.connectInput.Blur()
return m, nil
}
if msg.String() == "enter" {
target := m.connectInput.Value()
if num, ok := m.contacts[target]; ok {
m.targetName = target
target = num
} else {
m.targetName = target
}
req := protocol.Message{
Type: protocol.TypeLookup,
Sender: m.keys.PubHex,
Content: target,
}
m.conn.WriteJSON(req)
m.messages = append(m.messages, fmt.Sprintf("System: Connecting to %s...", m.targetName))
m.viewport.SetContent(strings.Join(m.messages, "\n"))
m.viewport.GotoBottom()
m.mode = ModeNormal
m.connectInput.Blur()
return m, nil
}
case ModeList:
switch msg.String() {
case "esc", "q":
m.mode = ModeNormal
return m, nil
case "k", "up":
if m.listCursor > 0 {
m.listCursor--
}
case "j", "down":
if m.listCursor < len(m.contactList)-1 {
m.listCursor++
}
case "enter":
if len(m.contactList) > 0 {
selection := m.contactList[m.listCursor]
targetNum := m.contacts[selection]
m.targetName = selection
req := protocol.Message{
Type: protocol.TypeLookup,
Sender: m.keys.PubHex,
Content: targetNum,
}
m.conn.WriteJSON(req)
m.messages = append(m.messages, fmt.Sprintf("System: Connecting to %s...", m.targetName))
m.viewport.SetContent(strings.Join(m.messages, "\n"))
m.viewport.GotoBottom()
}
m.mode = ModeNormal
return m, nil
}
}
case wsMsg:
encBytes, _ := base64.StdEncoding.DecodeString(msg.Content)
senderBytes, _ := hex.DecodeString(msg.Sender)
var senderKey [32]byte
copy(senderKey[:], senderBytes)
var nonce [24]byte
copy(nonce[:], encBytes[:24])
decrypted, ok := box.Open(nil, encBytes[24:], &nonce, &senderKey, m.keys.Private)
if !ok {
m.messages = append(m.messages, "System: Failed to decrypt message from "+msg.Sender[:8])
} else {
m.messages = append(m.messages, fmt.Sprintf("Friend (%s): %s", msg.Sender[:8], string(decrypted)))
switch msg.Type {
case protocol.TypeIdentity:
m.myAccountNum = msg.Content
m.messages = append(m.messages, fmt.Sprintf("System: 🟢 Online. Account #%s", m.myAccountNum))
case protocol.TypeLookupResponse:
if msg.Content == "" {
m.messages = append(m.messages, fmt.Sprintf("System: ❌ Account #%s not found.", msg.Target))
} else {
m.targetAcc = msg.Target
m.targetHex = msg.Content
decoded, _ := hex.DecodeString(m.targetHex)
var keyArr [32]byte
copy(keyArr[:], decoded)
m.targetPub = &keyArr
displayName := m.targetName
if displayName == "" {
displayName = m.targetAcc
}
m.messages = append(m.messages, fmt.Sprintf("System: 🔒 Secure channel established with %s", displayName))
}
case protocol.TypeMsg:
encBytes, _ := base64.StdEncoding.DecodeString(msg.Content)
senderBytes, _ := hex.DecodeString(msg.Sender)
var senderKey [32]byte
copy(senderKey[:], senderBytes)
var nonce [24]byte
copy(nonce[:], encBytes[:24])
decrypted, ok := box.Open(nil, encBytes[24:], &nonce, &senderKey, m.keys.Private)
if !ok {
m.messages = append(m.messages, "System: ⚠️ Decryption failed!")
} else {
m.messages = append(m.messages, fmt.Sprintf("%s: %s", m.targetName, string(decrypted)))
}
}
m.viewport.SetContent(strings.Join(m.messages, "\n"))
m.viewport.GotoBottom()
}
m.textarea, tiCmd = m.textarea.Update(msg)
m.viewport, vpCmd = m.viewport.Update(msg)
return m, tea.Batch(tiCmd, vpCmd)
m.viewport, cmd = m.viewport.Update(msg)
cmds = append(cmds, cmd)
if m.mode == ModeInsert {
m.textarea, cmd = m.textarea.Update(msg)
cmds = append(cmds, cmd)
}
if m.mode == ModeConnect {
m.connectInput, cmd = m.connectInput.Update(msg)
cmds = append(cmds, cmd)
}
return m, tea.Batch(cmds...)
}
func (m model) View() string {
return fmt.Sprintf(
"%s\n\n%s",
m.viewport.View(),
m.textarea.View(),
) + "\n\nPress Esc to quit."
var status string
var modeColor string
switch m.mode {
case ModeNormal:
status = "-- NORMAL -- (i: insert, c: connect, l: list, d: disconnect)"
modeColor = "\033[1;34m" // Blue
case ModeInsert:
status = "-- INSERT -- (Esc: normal, Enter: newline)"
modeColor = "\033[1;32m" // Green
case ModeConnect:
status = "-- CONNECT -- (Enter: confirm, Esc: cancel)"
modeColor = "\033[1;33m" // Yellow
case ModeList:
status = "-- CONTACTS -- (j/k: move, Enter: select)"
modeColor = "\033[1;35m" // Purple
}
var mainView string
if m.mode == ModeConnect {
dialog := fmt.Sprintf(
"Connect to User:\n\n%s",
m.connectInput.View(),
)
mainView = fmt.Sprintf("\n\n %s\n\n", dialog)
} else if m.mode == ModeList {
var items []string
items = append(items, "Select a Contact:\n")
for i, name := range m.contactList {
cursor := " "
if i == m.listCursor {
cursor = ">"
}
items = append(items, fmt.Sprintf("%s %s (#%s)", cursor, name, m.contacts[name]))
}
if len(m.contactList) == 0 {
items = append(items, " (No contacts saved. Use /add in Insert mode or edit contacts.json)")
}
mainView = strings.Join(items, "\n")
} else {
mainView = fmt.Sprintf("%s\n\n%s", m.viewport.View(), m.textarea.View())
}
return fmt.Sprintf("%s%s\033[0m\n\n%s", modeColor, status, mainView)
}