Add project-scoped Codex execution over a private socket

internal/codex/http.go

@@ -0,0 +1,143 @@
+package codex
+
+import (
+	"bytes"
+	"context"
+	"crypto/sha256"
+	"crypto/subtle"
+	"encoding/hex"
+	"encoding/json"
+	"errors"
+	"io"
+	"net"
+	"net/http"
+	"os"
+	"strings"
+	"time"
+)
+
+const (
+	maxRequest = 32 << 10
+	maxOutput  = 512 << 10
+	timeout    = 90 * time.Second
+)
+
+type task struct {
+	Prompt string `json:"prompt"`
+}
+
+func writeJSON(w http.ResponseWriter, status int, value any) {
+	w.Header().Set("Content-Type", "application/json")
+	w.Header().Set("Cache-Control", "no-store")
+	w.WriteHeader(status)
+	_ = json.NewEncoder(w).Encode(value)
+}
+
+func fail(w http.ResponseWriter, status int, code string) {
+	writeJSON(w, status, map[string]string{"error": code})
+}
+
+func readTask(w http.ResponseWriter, r *http.Request, limit int64) (task, bool) {
+	var input task
+	decoder := json.NewDecoder(http.MaxBytesReader(w, r.Body, limit))
+	decoder.DisallowUnknownFields()
+	err := decoder.Decode(&input)
+	if err == nil {
+		err = decoder.Decode(new(any))
+		if err == io.EOF && strings.TrimSpace(input.Prompt) != "" {
+			return input, true
+		}
+	}
+	var tooLarge *http.MaxBytesError
+	if errors.As(err, &tooLarge) {
+		fail(w, 413, "request_too_large")
+	} else {
+		fail(w, 400, "invalid_task")
+	}
+	return task{}, false
+}
+
+// Register enables the public endpoint only when both configuration files are supplied.
+// Each credential digest belongs to exactly one project.
+func Register(mux *http.ServeMux, projects map[string]string, socket, grantsFile string) error {
+	if socket == "" && grantsFile == "" {
+		return nil
+	}
+	if socket == "" || grantsFile == "" {
+		return errors.New("Codex requires both a socket and a grants file")
+	}
+	transport := &http.Transport{
+		DialContext: func(ctx context.Context, _, _ string) (net.Conn, error) {
+			return (&net.Dialer{}).DialContext(ctx, "unix", socket)
+		},
+		DisableKeepAlives: true,
+	}
+	client := &http.Client{
+		Transport:     transport,
+		Timeout:       timeout + 2*time.Second,
+		CheckRedirect: func(*http.Request, []*http.Request) error { return http.ErrUseLastResponse },
+	}
+	mux.HandleFunc("POST /projects/{project}/codex/exec", func(w http.ResponseWriter, r *http.Request) {
+		scheme, token, _ := strings.Cut(r.Header.Get("Authorization"), " ")
+		key, err := hex.DecodeString(token)
+		if !strings.EqualFold(scheme, "Bearer") || err != nil || len(key) != 32 {
+			w.Header().Set("WWW-Authenticate", "Bearer")
+			fail(w, 401, "unauthorized")
+			return
+		}
+		data, err := os.ReadFile(grantsFile)
+		var grants map[string]string
+		if err != nil || json.Unmarshal(data, &grants) != nil {
+			fail(w, 503, "codex_unavailable")
+			return
+		}
+		hash := sha256.Sum256(key)
+		var scope string
+		for digest, project := range grants {
+			value, err := hex.DecodeString(digest)
+			if err == nil && subtle.ConstantTimeCompare(hash[:], value) == 1 {
+				scope = project
+			}
+		}
+		if scope == "" {
+			w.Header().Set("WWW-Authenticate", "Bearer")
+			fail(w, 401, "unauthorized")
+			return
+		}
+		project := r.PathValue("project")
+		if _, exists := projects[project]; !exists || scope != project {
+			fail(w, 403, "forbidden")
+			return
+		}
+		input, ok := readTask(w, r, maxRequest)
+		if !ok {
+			return
+		}
+		body, _ := json.Marshal(input)
+		req, err := http.NewRequestWithContext(r.Context(), http.MethodPost, "http://codex/exec", bytes.NewReader(body))
+		if err != nil {
+			fail(w, 502, "codex_unavailable")
+			return
+		}
+		resp, err := client.Do(req)
+		if err != nil {
+			if errors.Is(err, context.DeadlineExceeded) {
+				fail(w, 504, "codex_timeout")
+			} else {
+				fail(w, 502, "codex_unavailable")
+			}
+			return
+		}
+		defer resp.Body.Close()
+		data, err = io.ReadAll(io.LimitReader(resp.Body, maxOutput+1))
+		if err != nil || len(data) > maxOutput || !json.Valid(data) {
+			fail(w, 502, "invalid_codex_response")
+			return
+		}
+		if retry := resp.Header.Get("Retry-After"); retry != "" {
+			w.Header().Set("Retry-After", retry)
+		}
+		writeJSON(w, resp.StatusCode, json.RawMessage(data))
+	})
+	return nil
+}