add 1 on 1 chat

cmd/server/chat.go

@@ -2,11 +2,14 @@ package main
 
 import (
 	"context"
+	"encoding/json"
 	"errors"
 	"io"
-	"log"
+	"log/slog"
 	"net"
 	"net/http"
+	"net/url"
+	"strings"
 	"sync"
 	"time"
 
@@ -15,68 +18,75 @@ import (
 	"github.com/coder/websocket"
 )
 
-// chatServer enables broadcasting to a set of subscribers.
+// chatMessage is the payload sent from the frontend to /publish.
+type chatMessage struct {
+	To      string `json:"to"`
+	Message string `json:"message"`
+	From    string `json:"from,omitempty"` // optional, server fills this in
+}
+
+// chatServer enables 1:1 messaging between named subscribers.
 type chatServer struct {
 	// subscriberMessageBuffer controls the max number
 	// of messages that can be queued for a subscriber
 	// before it is kicked.
-	//
-	// Defaults to 16.
 	subscriberMessageBuffer int
 
 	// publishLimiter controls the rate limit applied to the publish endpoint.
-	//
-	// Defaults to one publish every 100ms with a burst of 8.
 	publishLimiter *rate.Limiter
 
 	// logf controls where logs are sent.
-	// Defaults to log.Printf.
 	logf func(f string, v ...any)
 
-	// serveMux routes the various endpoints to the appropriate handler.
+	// serveMux routes the various endpoints.
 	serveMux http.ServeMux
 
-	subscribersMu sync.Mutex
-	subscribers   map[*subscriber]struct{}
+	// subscribersMu protects the subscribers map.
+	// Key is the username (from query param).
+	subscribersMu sync.RWMutex
+	subscribers   map[string]*subscriber
 }
 
 // newChatServer constructs a chatServer with the defaults.
 func newChatServer() *chatServer {
 	cs := &chatServer{
 		subscriberMessageBuffer: 16,
-		logf:                    log.Printf,
-		subscribers:             make(map[*subscriber]struct{}),
+		logf:                    slog.Info,
+		subscribers:             make(map[string]*subscriber),
 		publishLimiter:          rate.NewLimiter(rate.Every(time.Millisecond*100), 8),
 	}
-	// Serve static files (index.html, index.js, index.css) from cwd
 	cs.serveMux.Handle("/", http.FileServer(http.Dir(".")))
-	// WebSocket endpoint: clients connect here to receive messages
 	cs.serveMux.HandleFunc("/subscribe", cs.subscribeHandler)
-	// HTTP POST endpoint: clients POST here to send a message
 	cs.serveMux.HandleFunc("/publish", cs.publishHandler)
+	cs.serveMux.HandleFunc("/users", cs.usersHandler) // list online users
 
 	return cs
 }
 
 // subscriber represents a single connected client.
-// Messages are sent on the msgs channel. If the client
-// cannot keep up (buffer full), closeSlow is called
-// to disconnect them.
 type subscriber struct {
-	msgs      chan []byte
-	closeSlow func()
+	username  string             // unique identifier from query param
+	msgs      chan []byte        // buffered channel for outbound messages
+	closeSlow func()             // kicks the subscriber if buffer fills
+	cancel    context.CancelFunc // cancels the read context on disconnect
 }
 
-// ServeHTTP makes chatServer implement http.Handler,
-// delegating to the internal serveMux.
+// ServeHTTP delegates to the internal serveMux.
 func (cs *chatServer) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 	cs.serveMux.ServeHTTP(w, r)
 }
 
-// subscribeHandler accepts the WebSocket connection and then subscribes
-// it to all future messages.
+// subscribeHandler extracts the username from query params and upgrades to WebSocket.
 func (cs *chatServer) subscribeHandler(w http.ResponseWriter, r *http.Request) {
-	err := cs.subscribe(w, r)
+	// Extract username from ?user= parameter
+	username := r.URL.Query().Get("user")
+	if username == "" {
+		http.Error(w, "missing user parameter", http.StatusBadRequest)
+		return
+	}
+	username = url.PathEscape(strings.TrimSpace(username))
+
+	err := cs.subscribe(w, r, username)
 	if errors.Is(err, context.Canceled) {
 		return
 	}
@@ -85,48 +95,20 @@ func (cs *chatServer) subscribeHandler(w http.ResponseWriter, r *http.Request) {
 		return
 	}
 	if err != nil {
-		cs.logf("%v", err)
+		cs.logf("subscribe error for user %s: %v", username, err)
 		return
 	}
 }
 
-// publishHandler reads the request body with a limit of 8192 bytes
-// and then publishes the received message.
-func (cs *chatServer) publishHandler(w http.ResponseWriter, r *http.Request) {
-	if r.Method != "POST" {
-		http.Error(w, http.StatusText(http.StatusMethodNotAllowed), http.StatusMethodNotAllowed)
-		return
-	}
-	body := http.MaxBytesReader(w, r.Body, 8192)
-	msg, err := io.ReadAll(body)
-	if err != nil {
-		http.Error(w, http.StatusText(http.StatusRequestEntityTooLarge), http.StatusRequestEntityTooLarge)
-		return
-	}
-
-	cs.publish(msg)
-
-	w.WriteHeader(http.StatusAccepted)
-}
-
-// subscribe upgrades the HTTP connection to a WebSocket,
-// registers the subscriber, and then loops forever writing
-// messages from the subscriber's channel to the WebSocket.
-//
-// It uses CloseRead to keep reading control frames (close, ping, pong)
-// and cancel the context if the connection drops. This means we
-// never actually read message data from the WebSocket —
-// messages come in via /publish instead.
-func (cs *chatServer) subscribe(w http.ResponseWriter, r *http.Request) error {
-	// We need a mutex here because the WebSocket connection might
-	// be closed by closeSlow (from another goroutine) before we've
-	// assigned it to c. The mutex ensures we don't race.
+// subscribe upgrades the HTTP connection to WebSocket and registers the subscriber.
+func (cs *chatServer) subscribe(w http.ResponseWriter, r *http.Request, username string) error {
 	var mu sync.Mutex
 	var c *websocket.Conn
 	var closed bool
 
 	s := &subscriber{
-		msgs: make(chan []byte, cs.subscriberMessageBuffer),
+		username: username,
+		msgs:     make(chan []byte, cs.subscriberMessageBuffer),
 		closeSlow: func() {
 			mu.Lock()
 			defer mu.Unlock()
@@ -137,19 +119,17 @@ func (cs *chatServer) subscribe(w http.ResponseWriter, r *http.Request) error {
 		},
 	}
 
-	// Register the subscriber BEFORE accepting the WebSocket.
-	// This means we won't miss any messages published between
-	// the accept and the registration.
+	// Register BEFORE accepting WebSocket
 	cs.addSubscriber(s)
 	defer cs.deleteSubscriber(s)
 
-	// Upgrade HTTP to WebSocket
+	// Upgrade to WebSocket
 	c2, err := websocket.Accept(w, r, nil)
 	if err != nil {
 		return err
 	}
 
-	// Check if closeSlow was already called before we got the connection
+	// Check if already closed
 	mu.Lock()
 	if closed {
 		mu.Unlock()
@@ -160,68 +140,116 @@ func (cs *chatServer) subscribe(w http.ResponseWriter, r *http.Request) error {
 
 	defer c.CloseNow()
 
-	// CloseRead returns a context that is cancelled when:
-	// - The client sends a close frame
-	// - The connection drops
-	// This frees us from manually handling ping/pong and deadlines.
-	ctx := c.CloseRead(context.Background())
+	// Create context that cancels on close/read error
+	ctx, cancel := context.WithCancel(context.Background())
+	s.cancel = cancel
+
+	// CloseRead handles ping/pong/close frames automatically
+	connCtx := c.CloseRead(ctx)
+
+	cs.logf("user %s connected", username)
 
-	// Main loop: write messages to the WebSocket as they arrive
 	for {
 		select {
 		case msg := <-s.msgs:
-			err := writeTimeout(ctx, time.Second*5, c, msg)
+			err := writeTimeout(connCtx, time.Second*5, c, msg)
 			if err != nil {
 				return err
 			}
-		case <-ctx.Done():
-			// Context cancelled = client disconnected
-			return ctx.Err()
+		case <-connCtx.Done():
+			cs.logf("user %s disconnected", username)
+			cancel() // trigger cleanup
+			return connCtx.Err()
 		}
 	}
 }
 
-// publish sends a message to ALL subscribers.
-// It never blocks — if a subscriber's buffer is full,
-// they get kicked via closeSlow.
-func (cs *chatServer) publish(msg []byte) {
-	cs.subscribersMu.Lock()
-	defer cs.subscribersMu.Unlock()
+// publishHandler reads the JSON message and routes it to the recipient.
+func (cs *chatServer) publishHandler(w http.ResponseWriter, r *http.Request) {
+	if r.Method != http.MethodPost {
+		http.Error(w, http.StatusText(http.StatusMethodNotAllowed), http.StatusMethodNotAllowed)
+		return
+	}
+
+	body := http.MaxBytesReader(w, r.Body, 8192)
+	data, err := io.ReadAll(body)
+	if err != nil {
+		http.Error(w, http.StatusText(http.StatusRequestEntityTooLarge), http.StatusRequestEntityTooLarge)
+		return
+	}
+
+	var msg chatMessage
+	if err := json.Unmarshal(data, &msg); err != nil {
+		http.Error(w, "invalid JSON", http.StatusBadRequest)
+		return
+	}
+
+	if msg.To == "" || msg.Message == "" {
+		http.Error(w, "missing 'to' or 'message' field", http.StatusBadRequest)
+		return
+	}
+
+	// Deliver the message to the recipient
+	delivered := cs.deliverMessage(msg.To, msg.Message)
+
+	if delivered {
+		w.WriteHeader(http.StatusAccepted)
+	} else {
+		http.Error(w, "recipient not found", http.StatusNotFound)
+	}
+}
+
+// deliverMessage sends a message to a specific recipient.
+func (cs *chatServer) deliverMessage(toUsername, message string) bool {
+	cs.subscribersMu.RLock()
+	recipient, ok := cs.subscribers[toUsername]
+	cs.subscribersMu.RUnlock()
+
+	if !ok {
+		return false
+	}
 
-	// Rate limit: wait for a token before broadcasting
 	cs.publishLimiter.Wait(context.Background())
 
-	for s := range cs.subscribers {
-		select {
-		case s.msgs <- msg:
-			// Message delivered to subscriber's buffer
-		default:
-			// Buffer is full — kick the slow subscriber
-			go s.closeSlow()
-		}
+	select {
+	case recipient.msgs <- []byte(message):
+		return true
+	default:
+		// Buffer full — kick slow client
+		go recipient.closeSlow()
+		return false
 	}
 }
 
+// usersHandler returns list of currently connected usernames.
+func (cs *chatServer) usersHandler(w http.ResponseWriter, r *http.Request) {
+	cs.subscribersMu.RLock()
+	users := make([]string, 0, len(cs.subscribers))
+	for username := range cs.subscribers {
+		users = append(users, username)
+	}
+	cs.subscribersMu.RUnlock()
+
+	w.Header().Set("Content-Type", "application/json")
+	json.NewEncoder(w).Encode(users)
+}
+
 // addSubscriber registers a subscriber.
 func (cs *chatServer) addSubscriber(s *subscriber) {
 	cs.subscribersMu.Lock()
-	cs.subscribers[s] = struct{}{}
+	cs.subscribers[s.username] = s
 	cs.subscribersMu.Unlock()
 }
 
 // deleteSubscriber removes a subscriber.
 func (cs *chatServer) deleteSubscriber(s *subscriber) {
 	cs.subscribersMu.Lock()
-	delete(cs.subscribers, s)
+	delete(cs.subscribers, s.username)
 	cs.subscribersMu.Unlock()
 }
 
-// writeTimeout writes a message to the WebSocket with a deadline.
-// If the write doesn't complete within the timeout, the context
-// is cancelled and the operation aborts.
 func writeTimeout(ctx context.Context, timeout time.Duration, c *websocket.Conn, msg []byte) error {
 	ctx, cancel := context.WithTimeout(ctx, timeout)
 	defer cancel()
-
 	return c.Write(ctx, websocket.MessageText, msg)
 }