144 lines
3.8 KiB
Go
144 lines
3.8 KiB
Go
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
|
|
}
|