2 files changed, 214 insertions, 6 deletions

diff --git a/internal/handler/websocket/handler.go b/internal/handler/websocket/handler.go
index b7ea71d..c8fb6cc 100644
--- a/internal/handler/websocket/handler.go
+++ b/internal/handler/websocket/handler.go
@@ -26,10 +26,18 @@ type MetricsRecorder interface {
         SetActiveSubscriptions(count int)
 }
 
+type AuthConfig struct {
+        ReadEnabled      bool
+        WriteEnabled     bool
+        ReadAllowedPubkeys []string
+        WriteAllowedPubkeys []string
+}
+
 type Handler struct {
         store     EventStore
         subs      *subscription.Manager
         metrics   MetricsRecorder
+        authConfig *AuthConfig
         indexData IndexData
 }
 
@@ -44,6 +52,10 @@ func (h *Handler) SetMetrics(m MetricsRecorder) {
         h.metrics = m
 }
 
+func (h *Handler) SetAuthConfig(cfg *AuthConfig) {
+        h.authConfig = cfg
+}
+
 // SetIndexData sets the addresses for the index page
 func (h *Handler) SetIndexData(grpcAddr, httpAddr, wsAddr string) {
         h.indexData = IndexData{
@@ -79,6 +91,9 @@ func (h *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 
         ctx := r.Context()
         clientSubs := make(map[string]*subscription.Subscription)
+        var authenticatedPubkey string
+        var authChallenge string
+
         defer func() {
                 count := len(clientSubs)
                 for subID := range clientSubs {
@@ -97,14 +112,14 @@ func (h *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
                         return
                 }
 
-                if err := h.handleMessage(ctx, conn, data, clientSubs); err != nil {
+                if err := h.handleMessage(ctx, conn, data, clientSubs, &authenticatedPubkey, &authChallenge); err != nil {
                         log.Printf("Message handling error: %v", err)
                         h.sendNotice(ctx, conn, err.Error())
                 }
         }
 }
 
-func (h *Handler) handleMessage(ctx context.Context, conn *websocket.Conn, data []byte, clientSubs map[string]*subscription.Subscription) error {
+func (h *Handler) handleMessage(ctx context.Context, conn *websocket.Conn, data []byte, clientSubs map[string]*subscription.Subscription, authenticatedPubkey *string, authChallenge *string) error {
         var raw []json.RawMessage
         if err := json.Unmarshal(data, &raw); err != nil {
                 return fmt.Errorf("invalid JSON")
@@ -121,21 +136,75 @@ func (h *Handler) handleMessage(ctx context.Context, conn *websocket.Conn, data
 
         switch msgType {
         case "EVENT":
-                return h.handleEvent(ctx, conn, raw)
+                return h.handleEvent(ctx, conn, raw, authenticatedPubkey, authChallenge)
         case "REQ":
-                return h.handleReq(ctx, conn, raw, clientSubs)
+                return h.handleReq(ctx, conn, raw, clientSubs, authenticatedPubkey, authChallenge)
         case "CLOSE":
                 return h.handleClose(raw, clientSubs)
+        case "AUTH":
+                return h.handleAuth(ctx, conn, raw, authenticatedPubkey, authChallenge)
         default:
                 return fmt.Errorf("unknown message type: %s", msgType)
         }
 }
 
-func (h *Handler) handleEvent(ctx context.Context, conn *websocket.Conn, raw []json.RawMessage) error {
+func (h *Handler) requireAuth(ctx context.Context, conn *websocket.Conn, isWrite bool, authenticatedPubkey *string, authChallenge *string) error {
+        authRequired := false
+        var allowedPubkeys []string
+
+        if h.authConfig != nil {
+                if isWrite && h.authConfig.WriteEnabled {
+                        authRequired = true
+                        allowedPubkeys = h.authConfig.WriteAllowedPubkeys
+                } else if !isWrite && h.authConfig.ReadEnabled {
+                        authRequired = true
+                        allowedPubkeys = h.authConfig.ReadAllowedPubkeys
+                }
+        }
+
+        if !authRequired {
+                return nil
+        }
+
+        if *authenticatedPubkey == "" {
+                if *authChallenge == "" {
+                        challenge, err := h.store.(interface {
+                                CreateAuthChallenge(context.Context) (string, error)
+                        }).CreateAuthChallenge(ctx)
+                        if err != nil {
+                                return fmt.Errorf("failed to create auth challenge: %w", err)
+                        }
+                        *authChallenge = challenge
+                        h.sendAuthChallenge(ctx, conn, challenge)
+                }
+                return fmt.Errorf("restricted: authentication required")
+        }
+
+        if len(allowedPubkeys) > 0 {
+                allowed := false
+                for _, pk := range allowedPubkeys {
+                        if pk == *authenticatedPubkey {
+                                allowed = true
+                                break
+                        }
+                }
+                if !allowed {
+                        return fmt.Errorf("restricted: pubkey not authorized")
+                }
+        }
+
+        return nil
+}
+
+func (h *Handler) handleEvent(ctx context.Context, conn *websocket.Conn, raw []json.RawMessage, authenticatedPubkey *string, authChallenge *string) error {
         if len(raw) != 2 {
                 return fmt.Errorf("EVENT expects 2 elements")
         }
 
+        if err := h.requireAuth(ctx, conn, true, authenticatedPubkey, authChallenge); err != nil {
+                return err
+        }
+
         var event nostr.Event
         if err := json.Unmarshal(raw[1], &event); err != nil {
                 return fmt.Errorf("invalid event: %w", err)
@@ -185,11 +254,15 @@ func (h *Handler) handleEvent(ctx context.Context, conn *websocket.Conn, raw []j
         return nil
 }
 
-func (h *Handler) handleReq(ctx context.Context, conn *websocket.Conn, raw []json.RawMessage, clientSubs map[string]*subscription.Subscription) error {
+func (h *Handler) handleReq(ctx context.Context, conn *websocket.Conn, raw []json.RawMessage, clientSubs map[string]*subscription.Subscription, authenticatedPubkey *string, authChallenge *string) error {
         if len(raw) < 3 {
                 return fmt.Errorf("REQ expects at least 3 elements")
         }
 
+        if err := h.requireAuth(ctx, conn, false, authenticatedPubkey, authChallenge); err != nil {
+                return err
+        }
+
         var subID string
         if err := json.Unmarshal(raw[1], &subID); err != nil {
                 return fmt.Errorf("invalid subscription ID")
@@ -308,3 +381,50 @@ func (h *Handler) sendNotice(ctx context.Context, conn *websocket.Conn, notice s
         data, _ := json.Marshal(msg)
         return conn.Write(ctx, websocket.MessageText, data)
 }
+
+func (h *Handler) sendAuthChallenge(ctx context.Context, conn *websocket.Conn, challenge string) error {
+        msg := []interface{}{"AUTH", challenge}
+        data, _ := json.Marshal(msg)
+        return conn.Write(ctx, websocket.MessageText, data)
+}
+
+func (h *Handler) handleAuth(ctx context.Context, conn *websocket.Conn, raw []json.RawMessage, authenticatedPubkey *string, authChallenge *string) error {
+        if len(raw) != 2 {
+                return fmt.Errorf("AUTH expects 2 elements")
+        }
+
+        var authEvent nostr.Event
+        if err := json.Unmarshal(raw[1], &authEvent); err != nil {
+                return fmt.Errorf("invalid auth event: %w", err)
+        }
+
+        if authEvent.Kind != 22242 {
+                return fmt.Errorf("invalid auth event kind: expected 22242, got %d", authEvent.Kind)
+        }
+
+        if !authEvent.Verify() {
+                return fmt.Errorf("invalid auth event signature")
+        }
+
+        challengeTag := authEvent.Tags.Find("challenge")
+        if challengeTag == nil {
+                return fmt.Errorf("missing challenge tag in auth event")
+        }
+
+        eventChallenge := challengeTag.Value()
+        if eventChallenge != *authChallenge {
+                return fmt.Errorf("challenge mismatch")
+        }
+
+        if err := h.store.(interface {
+                ValidateAndConsumeChallenge(context.Context, string) error
+        }).ValidateAndConsumeChallenge(ctx, eventChallenge); err != nil {
+                return fmt.Errorf("invalid challenge: %w", err)
+        }
+
+        *authenticatedPubkey = authEvent.PubKey
+        log.Printf("WebSocket client authenticated: %s", authEvent.PubKey[:16])
+
+        h.sendOK(ctx, conn, authEvent.ID, true, "")
+        return nil
+}
diff --git a/internal/storage/auth.go b/internal/storage/auth.go
new file mode 100644
index 0000000..6eefa41
--- /dev/null
+++ b/internal/storage/auth.go
@@ -0,0 +1,88 @@
+package storage
+
+import (
+        "context"
+        "crypto/rand"
+        "encoding/hex"
+        "fmt"
+        "time"
+)
+
+const (
+        ChallengeLength = 32 // bytes
+        ChallengeTTL    = 10 * time.Minute
+)
+
+func generateChallenge() (string, error) {
+        bytes := make([]byte, ChallengeLength)
+        if _, err := rand.Read(bytes); err != nil {
+                return "", fmt.Errorf("failed to generate challenge: %w", err)
+        }
+        return hex.EncodeToString(bytes), nil
+}
+
+func (s *Storage) CreateAuthChallenge(ctx context.Context) (string, error) {
+        challenge, err := generateChallenge()
+        if err != nil {
+                return "", err
+        }
+
+        now := time.Now().Unix()
+        expiresAt := time.Now().Add(ChallengeTTL).Unix()
+
+        _, err = s.db.ExecContext(ctx,
+                "INSERT INTO auth_challenges (challenge, created_at, expires_at, used) VALUES (?, ?, ?, 0)",
+                challenge, now, expiresAt,
+        )
+        if err != nil {
+                return "", fmt.Errorf("failed to store challenge: %w", err)
+        }
+
+        return challenge, nil
+}
+
+func (s *Storage) ValidateAndConsumeChallenge(ctx context.Context, challenge string) error {
+        tx, err := s.db.BeginTx(ctx, nil)
+        if err != nil {
+                return fmt.Errorf("failed to begin transaction: %w", err)
+        }
+        defer tx.Rollback()
+
+        var expiresAt int64
+        var used int
+        err = tx.QueryRowContext(ctx,
+                "SELECT expires_at, used FROM auth_challenges WHERE challenge = ?",
+                challenge,
+        ).Scan(&expiresAt, &used)
+
+        if err != nil {
+                return fmt.Errorf("challenge not found or invalid")
+        }
+
+        if used != 0 {
+                return fmt.Errorf("challenge already used")
+        }
+
+        if time.Now().Unix() > expiresAt {
+                return fmt.Errorf("challenge expired")
+        }
+
+        _, err = tx.ExecContext(ctx,
+                "UPDATE auth_challenges SET used = 1 WHERE challenge = ?",
+                challenge,
+        )
+        if err != nil {
+                return fmt.Errorf("failed to mark challenge as used: %w", err)
+        }
+
+        return tx.Commit()
+}
+
+func (s *Storage) CleanupExpiredChallenges(ctx context.Context) error {
+        now := time.Now().Unix()
+        _, err := s.db.ExecContext(ctx,
+                "DELETE FROM auth_challenges WHERE expires_at < ?",
+                now,
+        )
+        return err
+}