simple server relay with clients

cmd/client/main.go

@@ -0,0 +1,183 @@
+package main
+
+import (
+	"crypto/rand"
+	"encoding/base64"
+	"encoding/hex"
+	"encoding/json"
+	"fmt"
+	"log"
+	"strings"
+
+	"gossip/pkg/protocol"
+
+	"github.com/charmbracelet/bubbles/textarea"
+	"github.com/charmbracelet/bubbles/viewport"
+	tea "github.com/charmbracelet/bubbletea"
+	"github.com/gorilla/websocket"
+	"golang.org/x/crypto/nacl/box"
+)
+
+type KeyPair struct {
+	Public  *[32]byte
+	Private *[32]byte
+	PubHex  string
+}
+
+type model struct {
+	conn      *websocket.Conn
+	keys      KeyPair
+	targetHex string
+	targetPub *[32]byte
+	viewport  viewport.Model
+	textarea  textarea.Model
+	messages  []string
+	err       error
+}
+
+type wsMsg protocol.Message
+
+func main() {
+	pub, priv, err := box.GenerateKey(rand.Reader)
+	if err != nil {
+		log.Fatal(err)
+	}
+	keys := KeyPair{Public: pub, Private: priv, PubHex: hex.EncodeToString(pub[:])}
+
+	c, _, err := websocket.DefaultDialer.Dial("ws://localhost:8080/ws", nil)
+	if err != nil {
+		log.Fatal("Could not connect to server:", err)
+	}
+	defer c.Close()
+
+	loginMsg := protocol.Message{Type: "login", Sender: keys.PubHex}
+	c.WriteJSON(loginMsg)
+
+	ta := textarea.New()
+	ta.Placeholder = "Type a message (or /connect <PUBKEY>)..."
+	ta.Focus()
+	ta.SetHeight(2)
+	ta.ShowLineNumbers = false
+
+	vp := viewport.New(80, 20)
+	vp.SetContent(fmt.Sprintf("Your ID: %s\nTo start, type: /connect <FRIEND_ID>", keys.PubHex))
+
+	m := model{
+		conn:     c,
+		keys:     keys,
+		textarea: ta,
+		viewport: vp,
+	}
+
+	p := tea.NewProgram(m)
+
+	go func() {
+		for {
+			_, data, err := c.ReadMessage()
+			if err != nil {
+				return
+			}
+			var msg protocol.Message
+			json.Unmarshal(data, &msg)
+			p.Send(wsMsg(msg))
+		}
+	}()
+
+	if _, err := p.Run(); err != nil {
+		log.Fatal(err)
+	}
+}
+
+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
+	)
+
+	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 == "" {
+				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]+"...")
+				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,
+				}
+
+				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()
+		}
+
+	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)))
+		}
+		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)
+}
+
+func (m model) View() string {
+	return fmt.Sprintf(
+		"%s\n\n%s",
+		m.viewport.View(),
+		m.textarea.View(),
+	) + "\n\nPress Esc to quit."
+}

cmd/server/main.go

@@ -0,0 +1,83 @@
+package main
+
+import (
+	"encoding/json"
+	"log"
+	"net/http"
+	"sync"
+
+	"gossip/pkg/protocol"
+
+	"github.com/gorilla/websocket"
+)
+
+type Server struct {
+	clients  map[string]*websocket.Conn
+	mu       sync.Mutex
+	upgrader websocket.Upgrader
+}
+
+func main() {
+	srv := &Server{
+		clients: make(map[string]*websocket.Conn),
+		upgrader: websocket.Upgrader{
+			ReadBufferSize:  1024,
+			WriteBufferSize: 1024,
+			CheckOrigin:     func(r *http.Request) bool { return true },
+		},
+	}
+
+	http.HandleFunc("/ws", srv.handleWS)
+	log.Println("Relay Server listening on :8080")
+	log.Fatal(http.ListenAndServe(":8080", nil))
+}
+
+func (s *Server) handleWS(w http.ResponseWriter, r *http.Request) {
+	conn, err := s.upgrader.Upgrade(w, r, nil)
+	if err != nil {
+		log.Println("Upgrade error:", err)
+		return
+	}
+	defer conn.Close()
+
+	var myPubKey string
+
+	for {
+		_, data, err := conn.ReadMessage()
+		if err != nil {
+			break
+		}
+
+		var msg protocol.Message
+		if err := json.Unmarshal(data, &msg); err != nil {
+			continue
+		}
+
+		switch msg.Type {
+		case "login":
+			s.mu.Lock()
+			s.clients[msg.Sender] = conn
+			s.mu.Unlock()
+			myPubKey = msg.Sender
+			log.Printf("Client connected: %s...", myPubKey[:8])
+		case "msg":
+			s.mu.Lock()
+			targetConn, ok := s.clients[msg.Target]
+			s.mu.Unlock()
+
+			if ok {
+				err = targetConn.WriteMessage(websocket.TextMessage, data)
+				if err != nil {
+					log.Printf("Failed to relay to %s", msg.Target[:8])
+				}
+			}
+		}
+	}
+
+	if myPubKey != "" {
+		s.mu.Lock()
+		delete(s.clients, myPubKey)
+		s.mu.Unlock()
+		log.Printf("Client disconnected: %s...", myPubKey[:8])
+	}
+}