path: root/internal/handler/websocket/handler.go

package websocket

import (
	"context"
	"encoding/json"
	"fmt"
	"log"
	"net/http"

	pb "northwest.io/muxstr/api/nostr/v1"
	"northwest.io/muxstr/internal/storage"
	"northwest.io/muxstr/internal/subscription"
	"northwest.io/muxstr/internal/websocket"
	"northwest.io/nostr"
)

type EventStore interface {
	StoreEvent(context.Context, *storage.EventData) error
	QueryEvents(context.Context, []*pb.Filter, *storage.QueryOptions) ([]*pb.Event, error)
	ProcessDeletion(context.Context, *pb.Event) error
}

type AuthStore interface {
	CreateAuthChallenge(context.Context) (string, error)
	ValidateAndConsumeChallenge(context.Context, string) error
}

type MetricsRecorder interface {
	IncrementSubscriptions()
	DecrementSubscriptions()
	SetActiveSubscriptions(count int)
}

type RateLimiter interface {
	Allow(identifier, method string) bool
}

type AuthConfig struct {
	ReadEnabled         bool
	WriteEnabled        bool
	ReadAllowedPubkeys  []string
	WriteAllowedPubkeys []string
}

type connState struct {
	authenticatedPubkey string
	authChallenge       string
	clientIP            string
}

type Handler struct {
	store      EventStore
	auth       AuthStore
	subs       *subscription.Manager
	metrics    MetricsRecorder
	limiter    RateLimiter
	authConfig *AuthConfig
	indexData  IndexData
}

func NewHandler(store EventStore, subs *subscription.Manager) *Handler {
	return &Handler{
		store: store,
		subs:  subs,
	}
}

func (h *Handler) SetMetrics(m MetricsRecorder) {
	h.metrics = m
}

func (h *Handler) SetAuth(a AuthStore) {
	h.auth = a
}

func (h *Handler) SetRateLimiter(l RateLimiter) {
	h.limiter = l
}

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{
		GrpcAddr: grpcAddr,
		HttpAddr: httpAddr,
		WsAddr:   wsAddr,
	}
}

func (h *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	// Handle GET requests (but not WebSocket upgrades)
	if r.Method == "GET" && r.Header.Get("Upgrade") != "websocket" {
		accept := r.Header.Get("Accept")

		// NIP-11: Relay information document
		if accept == "application/nostr+json" {
			h.ServeNIP11(w, r)
			return
		}

		// Serve HTML index for browsers
		h.ServeIndex(w, r, h.indexData)
		return
	}

	// Handle WebSocket upgrade
	conn, err := websocket.Accept(w, r)
	if err != nil {
		log.Printf("WebSocket accept failed: %v", err)
		return
	}
	defer conn.Close(websocket.StatusNormalClosure, "")

	ctx := r.Context()
	clientSubs := make(map[string]*subscription.Subscription)
	state := &connState{
		clientIP: getClientIP(r),
	}

	defer func() {
		count := len(clientSubs)
		for subID := range clientSubs {
			h.subs.Remove(subID)
		}
		if h.metrics != nil && count > 0 {
			for i := 0; i < count; i++ {
				h.metrics.DecrementSubscriptions()
			}
		}
	}()

	for {
		_, data, err := conn.Read(ctx)
		if err != nil {
			return
		}

		if err := h.handleMessage(ctx, conn, data, clientSubs, state); 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, state *connState) error {
	var raw []json.RawMessage
	if err := json.Unmarshal(data, &raw); err != nil {
		return fmt.Errorf("invalid JSON")
	}

	if len(raw) == 0 {
		return fmt.Errorf("empty message")
	}

	var msgType string
	if err := json.Unmarshal(raw[0], &msgType); err != nil {
		return fmt.Errorf("invalid message type")
	}

	switch msgType {
	case "EVENT":
		return h.handleEvent(ctx, conn, raw, state)
	case "REQ":
		return h.handleReq(ctx, conn, raw, clientSubs, state)
	case "CLOSE":
		return h.handleClose(raw, clientSubs)
	case "AUTH":
		return h.handleAuth(ctx, conn, raw, state)
	default:
		return fmt.Errorf("unknown message type: %s", msgType)
	}
}

func (h *Handler) requireAuth(ctx context.Context, conn *websocket.Conn, isWrite bool, state *connState) 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 state.authenticatedPubkey == "" {
		if state.authChallenge == "" {
			if h.auth == nil {
				return fmt.Errorf("auth required but no auth store configured")
			}
			challenge, err := h.auth.CreateAuthChallenge(ctx)
			if err != nil {
				return fmt.Errorf("failed to create auth challenge: %w", err)
			}
			state.authChallenge = challenge
			h.sendAuthChallenge(ctx, conn, challenge)
		}
		return nil
	}

	if len(allowedPubkeys) > 0 {
		allowed := false
		for _, pk := range allowedPubkeys {
			if pk == state.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, state *connState) error {
	if len(raw) != 2 {
		return fmt.Errorf("EVENT expects 2 elements")
	}

	if err := h.requireAuth(ctx, conn, true, state); err != nil {
		return err
	}

	if state.authenticatedPubkey == "" && h.authConfig != nil && h.authConfig.WriteEnabled {
		return nil
	}

	// Rate limiting - use pubkey if authenticated, otherwise IP
	if h.limiter != nil {
		identifier := state.authenticatedPubkey
		if identifier == "" {
			identifier = state.clientIP
		}
		if !h.limiter.Allow(identifier, "EVENT") {
			return fmt.Errorf("rate limit exceeded")
		}
	}

	var event nostr.Event
	if err := json.Unmarshal(raw[1], &event); err != nil {
		return fmt.Errorf("invalid event: %w", err)
	}

	if !event.CheckID() {
		h.sendOK(ctx, conn, event.ID, false, "invalid: event ID mismatch")
		return nil
	}

	if !event.Verify() {
		h.sendOK(ctx, conn, event.ID, false, "invalid: signature verification failed")
		return nil
	}

	pbEvent := NostrToPB(&event)
	canonicalJSON := event.Serialize()

	// Handle deletion events (kind 5) - process but don't store
	if pbEvent.Kind == 5 {
		if err := h.store.ProcessDeletion(ctx, pbEvent); err != nil {
			h.sendOK(ctx, conn, event.ID, false, fmt.Sprintf("deletion failed: %v", err))
			return nil
		}
		h.sendOK(ctx, conn, event.ID, true, "deleted")
		return nil
	}

	eventData := &storage.EventData{
		Event:         pbEvent,
		CanonicalJSON: canonicalJSON,
	}

	err := h.store.StoreEvent(ctx, eventData)
	if err == storage.ErrEventExists {
		h.sendOK(ctx, conn, event.ID, true, "duplicate: already have this event")
		return nil
	}
	if err != nil {
		h.sendOK(ctx, conn, event.ID, false, fmt.Sprintf("error: %v", err))
		return nil
	}

	h.subs.MatchAndFan(pbEvent)

	h.sendOK(ctx, conn, event.ID, true, "")
	return nil
}

func (h *Handler) handleReq(ctx context.Context, conn *websocket.Conn, raw []json.RawMessage, clientSubs map[string]*subscription.Subscription, state *connState) error {
	if len(raw) < 3 {
		return fmt.Errorf("REQ expects at least 3 elements")
	}

	if err := h.requireAuth(ctx, conn, false, state); err != nil {
		return err
	}

	if state.authenticatedPubkey == "" && h.authConfig != nil && h.authConfig.ReadEnabled {
		return nil
	}

	// Rate limiting - use pubkey if authenticated, otherwise IP
	if h.limiter != nil {
		identifier := state.authenticatedPubkey
		if identifier == "" {
			identifier = state.clientIP
		}
		if !h.limiter.Allow(identifier, "REQ") {
			return fmt.Errorf("rate limit exceeded")
		}
	}

	var subID string
	if err := json.Unmarshal(raw[1], &subID); err != nil {
		return fmt.Errorf("invalid subscription ID")
	}

	var filters []*pb.Filter
	for i := 2; i < len(raw); i++ {
		var nostrFilter nostr.Filter
		if err := json.Unmarshal(raw[i], &nostrFilter); err != nil {
			return fmt.Errorf("invalid filter: %w", err)
		}

		pbFilter := NostrFilterToPB(&nostrFilter)
		filters = append(filters, pbFilter)
	}

	if existing, ok := clientSubs[subID]; ok {
		h.subs.Remove(existing.ID)
		delete(clientSubs, subID)
		if h.metrics != nil {
			h.metrics.DecrementSubscriptions()
		}
	}

	storedEvents, err := h.store.QueryEvents(ctx, filters, &storage.QueryOptions{Limit: 0})
	if err != nil {
		return fmt.Errorf("query failed: %w", err)
	}

	for _, pbEvent := range storedEvents {
		event := PBToNostr(pbEvent)
		h.sendEvent(ctx, conn, subID, event)
	}

	h.sendEOSE(ctx, conn, subID)

	sub := &subscription.Subscription{
		ID:      subID,
		Filters: filters,
		Events:  make(chan *pb.Event, 100),
	}
	sub.InitDone()

	h.subs.Add(sub)
	clientSubs[subID] = sub

	if h.metrics != nil {
		h.metrics.IncrementSubscriptions()
	}

	go h.streamEvents(ctx, conn, sub)

	return nil
}

func (h *Handler) handleClose(raw []json.RawMessage, clientSubs map[string]*subscription.Subscription) error {
	if len(raw) != 2 {
		return fmt.Errorf("CLOSE expects 2 elements")
	}

	var subID string
	if err := json.Unmarshal(raw[1], &subID); err != nil {
		return fmt.Errorf("invalid subscription ID")
	}

	if sub, ok := clientSubs[subID]; ok {
		h.subs.Remove(sub.ID)
		delete(clientSubs, subID)
		if h.metrics != nil {
			h.metrics.DecrementSubscriptions()
		}
	}

	return nil
}

func (h *Handler) streamEvents(ctx context.Context, conn *websocket.Conn, sub *subscription.Subscription) {
	for {
		select {
		case pbEvent, ok := <-sub.Events:
			if !ok {
				return
			}
			event := PBToNostr(pbEvent)
			h.sendEvent(ctx, conn, sub.ID, event)

		case <-ctx.Done():
			return

		case <-sub.Done():
			return
		}
	}
}

func (h *Handler) sendEvent(ctx context.Context, conn *websocket.Conn, subID string, event *nostr.Event) error {
	msg := []interface{}{"EVENT", subID, event}
	data, _ := json.Marshal(msg)
	return conn.Write(ctx, websocket.MessageText, data)
}

func (h *Handler) sendOK(ctx context.Context, conn *websocket.Conn, eventID string, accepted bool, message string) error {
	msg := []interface{}{"OK", eventID, accepted, message}
	data, _ := json.Marshal(msg)
	return conn.Write(ctx, websocket.MessageText, data)
}

func (h *Handler) sendEOSE(ctx context.Context, conn *websocket.Conn, subID string) error {
	msg := []interface{}{"EOSE", subID}
	data, _ := json.Marshal(msg)
	return conn.Write(ctx, websocket.MessageText, data)
}

func (h *Handler) sendNotice(ctx context.Context, conn *websocket.Conn, notice string) error {
	msg := []interface{}{"NOTICE", notice}
	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, state *connState) 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 != state.authChallenge {
		return fmt.Errorf("challenge mismatch")
	}

	if h.auth == nil {
		return fmt.Errorf("auth required but no auth store configured")
	}

	if err := h.auth.ValidateAndConsumeChallenge(ctx, eventChallenge); err != nil {
		return fmt.Errorf("invalid challenge: %w", err)
	}

	state.authenticatedPubkey = authEvent.PubKey
	log.Printf("WebSocket client authenticated: %s", authEvent.PubKey[:16])

	h.sendOK(ctx, conn, authEvent.ID, true, "")
	return nil
}

// getClientIP extracts the real client IP from the HTTP request.
// Checks X-Forwarded-For and X-Real-IP headers first (for reverse proxy setups),
// then falls back to RemoteAddr.
func getClientIP(r *http.Request) string {
	// Check X-Forwarded-For header (set by reverse proxies like Caddy/nginx)
	if xff := r.Header.Get("X-Forwarded-For"); xff != "" {
		return xff
	}

	// Check X-Real-IP header
	if xri := r.Header.Get("X-Real-IP"); xri != "" {
		return xri
	}

	// Fall back to RemoteAddr (direct connection or proxy IP)
	return r.RemoteAddr
}