70 lines
2.1 KiB
Go
70 lines
2.1 KiB
Go
package codex
|
|
|
|
import (
|
|
"crypto/sha256"
|
|
"encoding/hex"
|
|
"fmt"
|
|
"net"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"os"
|
|
"path/filepath"
|
|
"strings"
|
|
"testing"
|
|
)
|
|
|
|
func TestProjectAuthorizationAndForwarding(t *testing.T) {
|
|
dir := t.TempDir()
|
|
socket := filepath.Join(dir, "worker.sock")
|
|
listener, err := net.Listen("unix", socket)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
worker := &http.Server{Handler: http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
if r.URL.Path != "/exec" || r.Header.Get("Authorization") != "" {
|
|
fail(w, 500, "unexpected_worker_request")
|
|
return
|
|
}
|
|
input, ok := readTask(w, r, 6*maxRequest)
|
|
if ok {
|
|
writeJSON(w, 200, map[string]string{"result": input.Prompt})
|
|
}
|
|
})}
|
|
go worker.Serve(listener)
|
|
t.Cleanup(func() { _ = worker.Close() })
|
|
|
|
key := make([]byte, 32)
|
|
token := hex.EncodeToString(key)
|
|
grants := filepath.Join(dir, "grants.json")
|
|
writeGrants := func(data string) {
|
|
t.Helper()
|
|
if err := os.WriteFile(grants, []byte(data), 0600); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
}
|
|
writeGrants(fmt.Sprintf(`{"%x":"notes"}`, sha256.Sum256(key)))
|
|
mux := http.NewServeMux()
|
|
if err := Register(mux, map[string]string{"notes": "", "board": ""}, socket, grants); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
check := func(project, credential, body string, status int, response string) {
|
|
t.Helper()
|
|
r := httptest.NewRequest("POST", "/projects/"+project+"/codex/exec", strings.NewReader(body))
|
|
if credential != "" {
|
|
r.Header.Set("Authorization", "Bearer "+credential)
|
|
}
|
|
w := httptest.NewRecorder()
|
|
mux.ServeHTTP(w, r)
|
|
if w.Code != status || !strings.Contains(w.Body.String(), response) {
|
|
t.Fatalf("got %d %s, want %d containing %q", w.Code, w.Body.String(), status, response)
|
|
}
|
|
}
|
|
check("notes", "", `{"prompt":"hello"}`, 401, "unauthorized")
|
|
check("board", token, `{"prompt":"hello"}`, 403, "forbidden")
|
|
check("notes", token, `{"prompt":"hello","sandbox":"danger"}`, 400, "invalid_task")
|
|
check("notes", token, `{"prompt":"`+strings.Repeat("x", maxRequest)+`"}`, 413, "request_too_large")
|
|
check("notes", token, `{"prompt":"hello"}`, 200, `"result":"hello"`)
|
|
writeGrants("{}")
|
|
check("notes", token, `{"prompt":"hello"}`, 401, "unauthorized")
|
|
}
|