package main

import (
	"crypto/sha256"
	"crypto/subtle"
	"fmt"
	"net"
	"net/http"
	"os"
	"strings"
)

// The boundary in front of the card.
//
// mavgpud has to listen on the LAN, because homesrv is the client and a
// loopback default takes the model arm down. That makes this the one hop on the
// workstation anything on the network could reach, and until 2026-08-11 it
// reverse-proxied every path to llama-server unauthenticated: any client could
// spend the card, hold the model resident by touching the idle clock, and read
// /slots, which carries the prompts of whoever else was using it.
//
// So: a bearer token every request must carry, read from a file, and a path
// allowlist so a token that leaks buys the model API and not the admin one. The
// CW2 transcriber beside this daemon has worked this way since it shipped; this
// is the same arrangement, not a new one.

// readToken loads the bearer token. The file holds the token and nothing else,
// trailing newline allowed. A path that is set and unreadable is fatal to the
// caller: a supervisor that silently ran without its boundary is the failure
// this exists to prevent.
func readToken(path string) (string, error) {
	b, err := os.ReadFile(path)
	if err != nil {
		return "", fmt.Errorf("token_file: %w", err)
	}
	tok := strings.TrimSpace(string(b))
	if tok == "" {
		return "", fmt.Errorf("token_file %s is empty", path)
	}
	return tok, nil
}

// loopbackListen reports whether addr can only be reached from this machine.
// An empty or wildcard host is not loopback, which is the case that matters:
// ":8080" is the shipped default and it answers the whole LAN.
func loopbackListen(addr string) bool {
	host, _, err := net.SplitHostPort(addr)
	if err != nil {
		host = addr
	}
	host = strings.Trim(host, "[]")
	if host == "" {
		return false
	}
	if host == "localhost" {
		return true
	}
	ip := net.ParseIP(host)
	return ip != nil && ip.IsLoopback()
}

// allowed is what a token buys. Everything llama-server exposes beyond this is
// refused, because the endpoints Maven does not call are the expensive ones to
// hand out: /slots returns other callers' prompts, and its save/restore actions
// write files chosen by the request.
//
// Adding a caller means adding its path here. That is deliberate — the list is
// short because Maven's use of the workstation is.
var allowed = map[string]string{
	"/v1/chat/completions": http.MethodPost,
	"/v1/completions":      http.MethodPost,
	"/v1/embeddings":       http.MethodPost,
	"/v1/models":           http.MethodGet,
	"/props":               http.MethodGet,
}

// requireToken authenticates, then bounds. Order matters: an unauthenticated
// client must not be able to make this daemon allocate a body buffer.
//
// /health is not exempt. It reports whether the card is loaded and free, which
// is exactly what someone deciding whether to take it from him would ask.
func requireToken(token string, maxBody int64, next http.Handler) http.Handler {
	want := sha256.Sum256([]byte(token))
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		got := sha256.Sum256([]byte(bearer(r)))
		if subtle.ConstantTimeCompare(got[:], want[:]) != 1 {
			w.Header().Set("WWW-Authenticate", "Bearer")
			http.Error(w, "unauthorized", http.StatusUnauthorized)
			return
		}
		r.Body = http.MaxBytesReader(w, r.Body, maxBody)
		next.ServeHTTP(w, r)
	})
}

// bearer pulls the credential out of the header. A malformed header yields the
// empty string, which fails the comparison like any other wrong token — there
// is no separate error for it, because telling a caller *how* it was wrong is
// the only thing a probe learns from a 401.
func bearer(r *http.Request) string {
	h := r.Header.Get("Authorization")
	const prefix = "Bearer "
	if len(h) <= len(prefix) || !strings.EqualFold(h[:len(prefix)], prefix) {
		return ""
	}
	return strings.TrimSpace(h[len(prefix):])
}

// allowlist refuses a path the model arm does not use. It answers 404 rather
// than 403 so a scan cannot map llama-server's surface through this hop.
func allowlist(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		method, ok := allowed[r.URL.Path]
		if !ok {
			http.NotFound(w, r)
			return
		}
		if r.Method != method {
			w.Header().Set("Allow", method)
			http.Error(w, "method not allowed", http.StatusMethodNotAllowed)
			return
		}
		next.ServeHTTP(w, r)
	})
}

// limitInflight caps concurrent proxied requests. A waiter leaves when its own
// context ends, so a client that gave up does not keep a slot: llama-server
// runs with -np 1 and queueing here is cheaper than queueing inside the child
// with a body held in memory on both sides.
func limitInflight(n int, next http.Handler) http.Handler {
	if n <= 0 {
		return next
	}
	slots := make(chan struct{}, n)
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		select {
		case slots <- struct{}{}:
			defer func() { <-slots }()
			next.ServeHTTP(w, r)
		case <-r.Context().Done():
			http.Error(w, "client went away", http.StatusServiceUnavailable)
		}
	})
}