so-many-crows/server.go

package main

import (
	"context"
	"crypto/rand"
	"encoding/hex"
	"encoding/json"
	"errors"
	"flag"
	"fmt"
	"io"
	"math"
	"net"
	"net/http"
	"os"
	"os/signal"
	"path/filepath"
	"strconv"
	"time"
)

const (
	maxBody   = 64 << 10 // POST bodies
	heartbeat = 15 * time.Second
	webDir    = "web"
)

// Server is the web front door: the same pipeline as the command line,
// with web approvals and an SSE hub instead of stdin and stdout.
type Server struct {
	p        *Pipeline
	approver *WebApprover
	hub      *Hub
	ctx      context.Context // cancelled on shutdown; lists and streams end with it
	beat     time.Duration
}

func NewServer(ctx context.Context, p *Pipeline, approver *WebApprover) *Server {
	return &Server{p: p, approver: approver, hub: NewHub(), ctx: ctx, beat: heartbeat}
}

func (s *Server) Handler() http.Handler {
	mux := http.NewServeMux()
	mux.HandleFunc("POST /api/lists", s.postList)
	mux.HandleFunc("GET /api/events", s.events)
	mux.HandleFunc("POST /api/verdicts", s.postVerdict)
	mux.HandleFunc("GET /api/trust", s.trust)
	if st, err := os.Stat(webDir); err == nil && st.IsDir() {
		mux.Handle("GET /", http.FileServer(http.Dir(webDir)))
	} else {
		mux.HandleFunc("GET /{$}", func(w http.ResponseWriter, r *http.Request) {
			fmt.Fprintln(w, "So Many Crows is listening. The web front end lands in ./web later.")
		})
	}
	return mux
}

func cmdServe(args []string) {
	fs := flag.NewFlagSet("serve", flag.ExitOnError)
	addr := fs.String("addr", "127.0.0.1:8080", "address to listen on (localhost only by default: there is no login yet)")
	timeout := fs.Duration("approval-timeout", 10*time.Minute, "unanswered approvals count as rejected after this long")
	fs.Parse(args)

	g := DefaultGuardrails()
	ledger, err := LoadLedger(filepath.Join(dataDir, "ledger.json"))
	must(err)
	flog, err := OpenFlightLog(filepath.Join(dataDir, "flight.jsonl"))
	must(err)
	defer flog.Close()

	ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt)
	defer stop()
	brain := NewBrain()
	approver := NewWebApprover(*timeout)
	p := &Pipeline{Flock: NewFlock(brain, g, flog), Ledger: ledger, Log: flog, G: g, Approver: approver}
	s := NewServer(ctx, p, approver)

	srv := &http.Server{Addr: *addr, Handler: s.Handler(), ReadHeaderTimeout: 10 * time.Second}
	go func() {
		<-ctx.Done()
		shut, cancel := context.WithTimeout(context.Background(), 5*time.Second)
		defer cancel()
		srv.Shutdown(shut)
	}()
	fmt.Printf("Brain: %s · listening on %s\n", brain.Name(), *addr)
	if err := srv.ListenAndServe(); err != nil && !errors.Is(err, http.ErrServerClosed) {
		must(err)
	}
	if err := ledger.Save(); err != nil {
		fmt.Fprintln(os.Stderr, "could not save ledger:", err)
	}
}

func writeJSON(w http.ResponseWriter, code int, v any) {
	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(code)
	json.NewEncoder(w).Encode(v)
}

func httpError(w http.ResponseWriter, code int, msg string) {
	writeJSON(w, code, map[string]string{"message": msg})
}

// decode reads a size-capped JSON body.
func decode(w http.ResponseWriter, r *http.Request, v any) bool {
	r.Body = http.MaxBytesReader(w, r.Body, maxBody)
	if err := json.NewDecoder(r.Body).Decode(v); err != nil {
		var tooBig *http.MaxBytesError
		if errors.As(err, &tooBig) {
			httpError(w, http.StatusRequestEntityTooLarge, fmt.Sprintf("body over %d bytes", maxBody))
		} else {
			httpError(w, http.StatusBadRequest, "body must be JSON: "+err.Error())
		}
		return false
	}
	return true
}

func (s *Server) postList(w http.ResponseWriter, r *http.Request) {
	var body struct {
		Text      string `json:"text"`
		Requester string `json:"requester"`
	}
	if !decode(w, r, &body) {
		return
	}
	if body.Requester == "" {
		body.Requester, _, _ = net.SplitHostPort(r.RemoteAddr)
	}

	ok, msg, wait := s.p.Ledger.Admit(body.Requester, time.Now(), s.p.G)
	if !ok {
		cooldown := int(math.Ceil(wait.Seconds()))
		s.p.Log.Write(Event{Kind: "denied", Requester: body.Requester, Note: msg})
		s.p.save()
		// Only a list that is still flying can hear about it; otherwise the 429 says it all.
		for _, id := range s.hub.ActiveLists(body.Requester) {
			s.hub.Publish(id, "requester.pranked", map[string]any{"message": msg, "cooldown_s": cooldown})
		}
		writeJSON(w, http.StatusTooManyRequests, map[string]any{"message": msg, "cooldown_s": cooldown})
		return
	}
	s.p.save()

	items, err := parseTodo(body.Text, s.p.G)
	switch {
	case err != nil:
		httpError(w, http.StatusBadRequest, err.Error())
		return
	case len(items) == 0:
		httpError(w, http.StatusBadRequest, "That list is empty.")
		return
	case len(items) > s.p.G.MaxItemsPerList:
		httpError(w, http.StatusBadRequest, fmt.Sprintf("That list has %d items. The flock takes at most %d at a time.", len(items), s.p.G.MaxItemsPerList))
		return
	}

	id := newListID()
	s.hub.Open(id, body.Requester)
	l := List{ID: id, Requester: body.Requester, Source: "web", Items: items, Parallel: true}
	go s.p.Run(s.ctx, l, s.hub.Emitter(id))
	writeJSON(w, http.StatusAccepted, map[string]string{"list_id": id})
}

func newListID() string {
	b := make([]byte, 6)
	rand.Read(b)
	return hex.EncodeToString(b)
}

func (s *Server) events(w http.ResponseWriter, r *http.Request) {
	listID := r.URL.Query().Get("list_id")
	lastID, _ := strconv.Atoi(r.Header.Get("Last-Event-ID"))
	past, ch, ok := s.hub.Subscribe(listID, lastID)
	if !ok {
		httpError(w, http.StatusNotFound, "no such list_id")
		return
	}
	defer s.hub.Unsubscribe(listID, ch)
	flusher, canFlush := w.(http.Flusher)
	if !canFlush {
		httpError(w, http.StatusInternalServerError, "streaming unsupported")
		return
	}

	h := w.Header()
	h.Set("Content-Type", "text/event-stream")
	h.Set("Cache-Control", "no-cache")
	h.Set("Connection", "keep-alive")
	h.Set("X-Accel-Buffering", "no")
	w.WriteHeader(http.StatusOK)

	send := func(m sseMsg) bool {
		_, err := fmt.Fprintf(w, "id: %d\nevent: %s\ndata: %s\n\n", m.ID, m.Name, m.Data)
		flusher.Flush()
		return err == nil
	}
	io.WriteString(w, ": caw\n\n")
	flusher.Flush()
	for _, m := range past {
		if !send(m) {
			return
		}
	}

	tick := time.NewTicker(s.beat)
	defer tick.Stop()
	for {
		select {
		case m, open := <-ch:
			if !open || !send(m) {
				return
			}
		case <-tick.C:
			if _, err := io.WriteString(w, ": heartbeat\n\n"); err != nil {
				return
			}
			flusher.Flush()
		case <-r.Context().Done():
			return
		case <-s.ctx.Done():
			return
		}
	}
}

func (s *Server) postVerdict(w http.ResponseWriter, r *http.Request) {
	var body struct {
		ItemID  string `json:"item_id"`
		Target  string `json:"target"`
		Approve *bool  `json:"approve"`
	}
	if !decode(w, r, &body) {
		return
	}
	if body.ItemID == "" || (body.Target != "triage" && body.Target != "output") || body.Approve == nil {
		httpError(w, http.StatusBadRequest, `need item_id, target ("triage" or "output") and approve (bool)`)
		return
	}
	if !s.approver.Resolve(body.ItemID, body.Target, *body.Approve) {
		httpError(w, http.StatusNotFound, "nothing is waiting on that verdict")
		return
	}
	w.WriteHeader(http.StatusNoContent)
}

func (s *Server) trust(w http.ResponseWriter, r *http.Request) {
	data, err := s.p.Ledger.MarshalSnapshot()
	if err != nil {
		httpError(w, http.StatusInternalServerError, err.Error())
		return
	}
	w.Header().Set("Content-Type", "application/json")
	w.Write(data)
}