diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index a1ce7d7..4053a72 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -53,16 +53,25 @@ jobs: include: - variant: gemma3-4b latest: true + target: "" model_url: https://huggingface.co/lmstudio-community/gemma-3-4b-it-GGUF/resolve/c536c4707e747055eecad7da65d46b6fb0ebaa79/gemma-3-4b-it-Q4_K_M.gguf model_sha256: be49949e48422e4547b00af14179a193d3777eea7fbbd7d6e1b0861304628a01 mmproj_url: https://huggingface.co/lmstudio-community/gemma-3-4b-it-GGUF/resolve/d400f8ba80bfa661d94a756ea3b663db8b00da85/mmproj-model-f16.gguf mmproj_sha256: 8c0fb064b019a6972856aaae2c7e4792858af3ca4561be2dbf649123ba6c40cb - variant: gemma3-12b latest: false + target: "" model_url: https://huggingface.co/lmstudio-community/gemma-3-12b-it-GGUF/resolve/ed6e7d1e8c65a0181cd45e3a194722d1c651f06e/gemma-3-12b-it-Q4_K_M.gguf model_sha256: 9610e3e07375303f6cd89086b496bcc1ab581177f52042eff536475a29283ba2 mmproj_url: https://huggingface.co/lmstudio-community/gemma-3-12b-it-GGUF/resolve/785c151b0dabf00cfaff1421239fafdbb8e5995d/mmproj-model-f16.gguf mmproj_sha256: 30c02d056410848227001830866e0a269fcc28aaf8ca971bded494003de9f5a5 + - variant: proxy + latest: false + target: proxy + model_url: "" + model_sha256: "" + mmproj_url: "" + mmproj_sha256: "" permissions: contents: read steps: @@ -96,6 +105,7 @@ jobs: with: context: . file: ./Dockerfile + target: ${{ matrix.target }} platforms: linux/amd64,linux/arm64 push: true tags: ${{ steps.meta.outputs.tags }} diff --git a/.gitignore b/.gitignore index adb7fe6..8259863 100644 --- a/.gitignore +++ b/.gitignore @@ -1,3 +1,4 @@ /localaik .cache/ -specs/ +/specs/ +/docs/ diff --git a/Dockerfile b/Dockerfile index 49446fb..7484a3d 100644 --- a/Dockerfile +++ b/Dockerfile @@ -4,6 +4,15 @@ WORKDIR /app COPY . . RUN CGO_ENABLED=0 GOOS=$TARGETOS GOARCH=$TARGETARCH go build -ldflags="-s -w" -o /out/localaik ./cmd/localaik +FROM alpine:3@sha256:28bd5fe8b56d1bd048e5babf5b10710ebe0bae67db86916198a6eec434943f8b AS proxy +RUN apk add --no-cache ca-certificates poppler-utils tini +COPY --from=proxy-builder /out/localaik /usr/local/bin/localaik +ENV PORT=8090 +HEALTHCHECK --interval=5s --timeout=3s --start-period=5s \ + CMD wget -q -O - "http://127.0.0.1:${PORT:-8090}/health" >/dev/null 2>&1 || exit 1 +EXPOSE 8090 +ENTRYPOINT ["tini", "--", "localaik"] + # Upstream does not ship semver for the server image; pin by digest for reproducible multi-arch builds. # Logical tag at pin time: server (includes llama-server --mmproj for Gemma 3 vision). Bump digest to upgrade. FROM ghcr.io/ggml-org/llama.cpp@sha256:80910e898e5d9a6b46ca9d1b4674d3e15faf6d32b9692eb6011ccd34b2cb8a06 diff --git a/Makefile b/Makefile index 56d02dd..2d08e9b 100644 --- a/Makefile +++ b/Makefile @@ -10,7 +10,7 @@ GOFILES := $(shell find cmd internal integration -name '*.go' -type f | sort) export GOCACHE -.PHONY: help fmt fmt-check lint test-unit test-integration test build docker-build docker-up docker-down +.PHONY: help fmt fmt-check lint test-unit test-integration test build docker-build docker-build-proxy docker-up docker-down help: @printf '%s\n' \ @@ -22,6 +22,7 @@ help: 'make test Run lint, unit tests, and integration tests' \ 'make build Build the localaik binary' \ 'make docker-build Build the Docker image' \ + 'make docker-build-proxy Build the proxy-only image' \ 'make docker-up Start the Docker image on PORT' \ 'make docker-down Stop and remove the Docker container' @@ -48,6 +49,9 @@ build: docker-build: @docker build -t "$(IMAGE)" . +docker-build-proxy: + @docker build --target proxy -t "$(IMAGE)-proxy" . + docker-up: @if [[ "$(BUILD_IMAGE)" == "1" ]]; then $(MAKE) docker-build IMAGE="$(IMAGE)"; fi @docker rm -f "$(CONTAINER_NAME)" >/dev/null 2>&1 || true diff --git a/README.md b/README.md index 5db182c..cf55cc8 100644 --- a/README.md +++ b/README.md @@ -11,7 +11,7 @@ A local compatibility server for the Gemini, OpenAI, and Anthropic APIs. Run one ## Motivation -Testing code that calls Gemini, OpenAI, or Anthropic is painful: real API calls are slow, cost money, and need network access. localaik gives you a single Docker container that speaks all three protocols backed by a local model — no API key, no internet, deterministic enough for CI. +Testing code that calls Gemini, OpenAI, or Anthropic is painful: real API calls are slow, cost money, and need network access. localaik gives you a single Docker container that speaks all three protocols backed by a local model, or the `proxy` tag if you already run your own model server. The model-bundled tags need no API key and no internet, and are deterministic enough for CI. ## Architecture @@ -133,9 +133,11 @@ client := anthropic.NewClient( | --------------------- | ------------------ | ---------- | | `latest`, `gemma3-4b` | Gemma 3 4B Q4_K_M | ~3 GB | | `gemma3-12b` | Gemma 3 12B Q4_K_M | ~7 GB | +| `proxy` | none (you supply) | ~41 MB | -Version-pinned tags follow the pattern `v0.1.1-gemma3-4b`, `v0.1.1-gemma3-12b`. +Version-pinned tags follow the pattern `v0.1.1-gemma3-4b`, `v0.1.1-gemma3-12b`, +`v0.1.1-proxy`. The `proxy` tag is never published as `latest`. ## Tuning (v0.1.3 onwards) @@ -182,6 +184,43 @@ services: | `LK_MLOCK` | 0 (off) | Lock model in RAM (`1` to enable) | +## Bring your own model server (`:proxy`) + +If you already run llama.cpp, vLLM, or anything else that speaks the OpenAI +chat-completions API, the `proxy` tag gives you the translation layer alone. It +contains no model and no inference engine. + +```bash +docker run -d -p 127.0.0.1:8090:8090 \ + -e LK_UPSTREAM=http://llama.internal:8080/v1 \ + gokhalh/localaik:proxy +``` + +| Env var | Default | Description | +| --- | --- | --- | +| `LK_UPSTREAM` | `http://127.0.0.1:8080/v1` | Base URL of your model server | +| `LK_UPSTREAM_AUTH_HEADER` | unset | A full header line sent to your server, for example `Authorization: Bearer abc123` | +| `PORT` | `8090` | Port localaik listens on | + +`LK_UPSTREAM_AUTH_HEADER` is sent only to your upstream. Credentials that +clients send to localaik are still discarded and never forwarded. It is attached +only to requests whose host matches `LK_UPSTREAM`, and while it is set a +redirect from your upstream is returned to the caller rather than followed. + +`/health` returns 503 until your upstream answers, so existing healthchecks and +CI wait loops work unchanged. + +### Security + +`:proxy` has a different risk profile from the model-bundled tags. Those keep +llama.cpp bound to localhost inside the container, so the only thing reachable +is a disposable local model. `:proxy` forwards into infrastructure you care +about, and localaik does not authenticate its callers by design. + +**Anyone who can reach port 8090 can use your model server without +credentials.** Bind to localhost and do not publish the port on a shared +network. localaik is a testing tool, not a gateway. + ## Implemented routes @@ -357,12 +396,15 @@ docker build \ --build-arg MMPROJ_URL=... \ --build-arg MMPROJ_SHA256=... \ -t gokhalh/localaik:custom . + +# Proxy only, no model or inference engine +docker build --target proxy -t gokhalh/localaik:proxy . ``` ## Limitations - Intended for tests and development, not production -- Image size is dominated by model weights +- Image size is dominated by model weights (not applicable to the `proxy` tag, which ships none) - Cold starts can take tens of seconds while the model loads - PDF rendering adds latency per page diff --git a/cmd/localaik/main.go b/cmd/localaik/main.go index ccfe049..48fccbd 100644 --- a/cmd/localaik/main.go +++ b/cmd/localaik/main.go @@ -11,20 +11,28 @@ import ( "github.com/harshaneel/localaik/internal/server" ) -func main() { - defaultPort := os.Getenv("PORT") - if defaultPort == "" { - defaultPort = "8090" +func resolveFlagDefault(envName, fallback string) string { + if value := os.Getenv(envName); value != "" { + return value } + return fallback +} - port := flag.String("port", defaultPort, "port to listen on") - upstream := flag.String("upstream", "http://127.0.0.1:8080/v1", "upstream OpenAI-compatible base URL") +func main() { + port := flag.String("port", resolveFlagDefault("PORT", "8090"), "port to listen on") + upstream := flag.String("upstream", resolveFlagDefault("LK_UPSTREAM", "http://127.0.0.1:8080/v1"), "upstream OpenAI-compatible base URL") flag.Parse() + authHeader := os.Getenv("LK_UPSTREAM_AUTH_HEADER") + if authHeader != "" && !server.ValidUpstreamAuthHeader(authHeader) { + log.Printf("localaik: LK_UPSTREAM_AUTH_HEADER is set but is not a valid \"Name: value\" header line; no credential will be sent upstream") + } + handler, err := server.New(server.Config{ - UpstreamBaseURL: *upstream, - HTTPClient: &http.Client{}, - PDFRenderer: pdf.NewExecRenderer("pdftoppm"), + UpstreamBaseURL: *upstream, + UpstreamAuthHeader: authHeader, + HTTPClient: &http.Client{}, + PDFRenderer: pdf.NewExecRenderer("pdftoppm"), }) if err != nil { log.Fatalf("localaik: %v", err) diff --git a/cmd/localaik/main_test.go b/cmd/localaik/main_test.go new file mode 100644 index 0000000..6e37a1f --- /dev/null +++ b/cmd/localaik/main_test.go @@ -0,0 +1,40 @@ +package main + +import ( + "testing" + + "github.com/harshaneel/localaik/internal/server" +) + +func TestResolveFlagDefaultPrefersEnv(t *testing.T) { + t.Setenv("LK_TEST_VALUE", "from-env") + + if got := resolveFlagDefault("LK_TEST_VALUE", "fallback"); got != "from-env" { + t.Fatalf("resolveFlagDefault = %q, want from-env", got) + } +} + +func TestResolveFlagDefaultFallsBack(t *testing.T) { + t.Setenv("LK_TEST_VALUE", "") + + if got := resolveFlagDefault("LK_TEST_VALUE", "fallback"); got != "fallback" { + t.Fatalf("resolveFlagDefault = %q, want fallback", got) + } +} + +func TestResolveFlagDefaultUnsetFallsBack(t *testing.T) { + if got := resolveFlagDefault("LK_DEFINITELY_UNSET_VALUE", "fallback"); got != "fallback" { + t.Fatalf("resolveFlagDefault = %q, want fallback", got) + } +} + +// The startup warning must be driven by the same predicate the transport uses; +// server.ValidUpstreamAuthHeader owns the table of cases. +func TestStartupWarningUsesTheServerPredicate(t *testing.T) { + if server.ValidUpstreamAuthHeader("Authorization: Bearer token123") != true { + t.Fatal("a valid header line was rejected") + } + if server.ValidUpstreamAuthHeader("InvalidHeader NoColon") != false { + t.Fatal("a header line with no colon was accepted") + } +} diff --git a/go.mod b/go.mod index 6545104..f8009e4 100644 --- a/go.mod +++ b/go.mod @@ -5,6 +5,7 @@ go 1.25 require ( github.com/anthropics/anthropic-sdk-go v1.61.0 github.com/openai/openai-go/v3 v3.36.0 + golang.org/x/net v0.41.0 google.golang.org/genai v1.57.0 ) @@ -29,7 +30,6 @@ require ( go.opencensus.io v0.24.0 // indirect go.yaml.in/yaml/v4 v4.0.0-rc.2 // indirect golang.org/x/crypto v0.40.0 // indirect - golang.org/x/net v0.41.0 // indirect golang.org/x/sync v0.16.0 // indirect golang.org/x/sys v0.35.0 // indirect golang.org/x/text v0.27.0 // indirect diff --git a/integration/proxy_image_test.go b/integration/proxy_image_test.go new file mode 100644 index 0000000..d8e8e3a --- /dev/null +++ b/integration/proxy_image_test.go @@ -0,0 +1,219 @@ +//go:build docker_integration + +package integration + +import ( + "bytes" + "encoding/base64" + "encoding/json" + "fmt" + "image/png" + "io" + "net" + "net/http" + "net/http/httptest" + "os/exec" + "strings" + "sync" + "testing" + "time" +) + +// Exercises the built proxy image against a stub upstream, proving all three +// protocol surfaces round-trip without an inference stack in the container. +func TestProxyImageRoundTripsAllProtocols(t *testing.T) { + image := "localaik:proxy-integration" + + build := exec.Command("docker", "build", "--target", "proxy", "-t", image, "..") + if out, err := build.CombinedOutput(); err != nil { + t.Fatalf("docker build failed: %v\n%s", err, out) + } + + var authMu sync.Mutex + var seenAuth string + var lastChatBody []byte + stub := httptest.NewUnstartedServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + authMu.Lock() + seenAuth = r.Header.Get("X-Proxy-Token") + if strings.HasSuffix(r.URL.Path, "/chat/completions") { + lastChatBody, _ = io.ReadAll(r.Body) + } + authMu.Unlock() + w.Header().Set("Content-Type", "application/json") + if strings.HasSuffix(r.URL.Path, "/tokenize") { + _, _ = w.Write([]byte(`{"tokens":[1,2,3]}`)) + return + } + _, _ = w.Write([]byte(`{"choices":[{"index":0,"message":{"role":"assistant","content":"stubbed"},"finish_reason":"stop"}]}`)) + })) + + // NewUnstartedServer already bound a loopback-only listener; close it and + // swap in one bound to 0.0.0.0 so the container can reach it. + stub.Listener.Close() + listener, err := net.Listen("tcp", "0.0.0.0:0") + if err != nil { + t.Fatalf("listen: %v", err) + } + stub.Listener = listener + stub.Start() + defer stub.Close() + + stubPort := listener.Addr().(*net.TCPAddr).Port + upstream := fmt.Sprintf("http://host.docker.internal:%d/v1", stubPort) + + // Best-effort: clear any container left behind by a prior run that + // crashed before its own cleanup ran, so the name doesn't collide. + _ = exec.Command("docker", "rm", "-f", "proxy-integration").Run() + + run := exec.Command("docker", "run", "-d", "--name", "proxy-integration", + "--add-host", "host.docker.internal:host-gateway", + "-p", "18097:8090", + "-e", "LK_UPSTREAM="+upstream, + "-e", "LK_UPSTREAM_AUTH_HEADER=X-Proxy-Token: integration-secret", + image) + if out, err := run.CombinedOutput(); err != nil { + t.Fatalf("docker run failed: %v\n%s", err, out) + } + defer exec.Command("docker", "rm", "-f", "proxy-integration").Run() + + waitForHealth(t, "http://127.0.0.1:18097/health") + + cases := []struct { + name string + path string + body string + expectedKey string + }{ + {"openai", "/v1/chat/completions", `{"model":"m","messages":[{"role":"user","content":"hi"}]}`, "choices"}, + {"gemini", "/v1beta/models/m:generateContent", `{"contents":[{"role":"user","parts":[{"text":"hi"}]}]}`, "candidates"}, + {"anthropic", "/v1/messages", `{"max_tokens":16,"messages":[{"role":"user","content":"hi"}]}`, "content"}, + } + + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + resp, err := http.Post("http://127.0.0.1:18097"+tc.path, "application/json", strings.NewReader(tc.body)) + if err != nil { + t.Fatalf("post: %v", err) + } + defer resp.Body.Close() + + if resp.StatusCode != http.StatusOK { + t.Fatalf("status = %d, want 200", resp.StatusCode) + } + var decoded map[string]any + if err := json.NewDecoder(resp.Body).Decode(&decoded); err != nil { + t.Fatalf("decode: %v", err) + } + // The expected key catches a response shaped by the wrong handler. + if _, ok := decoded[tc.expectedKey]; !ok { + t.Fatalf("response missing %q key, got keys %v", tc.expectedKey, mapKeys(decoded)) + } + if tc.name == "anthropic" { + if role, _ := decoded["role"].(string); role != "assistant" { + t.Fatalf("role = %q, want assistant", role) + } + } + }) + } + + // Alpine's poppler-utils is a different build from the full image's Debian + // one, so the PDF-to-PNG path has to be proven inside this container. + t.Run("pdf_to_png", func(t *testing.T) { + payload, err := json.Marshal(map[string]any{ + "contents": []any{map[string]any{ + "role": "user", + "parts": []any{ + map[string]any{"text": "Read this document."}, + map[string]any{"inlineData": map[string]string{ + "mimeType": "application/pdf", + "data": base64.StdEncoding.EncodeToString(buildSimplePDF([]string{ + "NAME: ALICE", + "CITY: BOSTON", + })), + }}, + }, + }}, + }) + if err != nil { + t.Fatalf("marshal request: %v", err) + } + + resp, err := http.Post("http://127.0.0.1:18097/v1beta/models/m:generateContent", "application/json", bytes.NewReader(payload)) + if err != nil { + t.Fatalf("post: %v", err) + } + defer resp.Body.Close() + + if resp.StatusCode != http.StatusOK { + body, _ := io.ReadAll(resp.Body) + t.Fatalf("status = %d, want 200; body=%s", resp.StatusCode, body) + } + + authMu.Lock() + body := string(lastChatBody) + authMu.Unlock() + + if strings.Contains(body, "application/pdf") { + t.Fatal("upstream received the raw PDF instead of rendered pages") + } + + const prefix = "data:image/png;base64," + start := strings.Index(body, prefix) + if start == -1 { + t.Fatalf("upstream received no rendered PNG page; body=%s", truncateForLog(body)) + } + encoded := body[start+len(prefix):] + if end := strings.IndexByte(encoded, '"'); end != -1 { + encoded = encoded[:end] + } + page, err := base64.StdEncoding.DecodeString(encoded) + if err != nil { + t.Fatalf("rendered page was not valid base64: %v", err) + } + config, err := png.DecodeConfig(bytes.NewReader(page)) + if err != nil { + t.Fatalf("rendered page was not a valid PNG: %v", err) + } + if config.Width == 0 || config.Height == 0 { + t.Fatalf("rendered page is %dx%d", config.Width, config.Height) + } + }) + + authMu.Lock() + got := seenAuth + authMu.Unlock() + if got != "integration-secret" { + t.Fatalf("upstream saw X-Proxy-Token = %q, want integration-secret", got) + } +} + +func truncateForLog(body string) string { + const limit = 300 + if len(body) <= limit { + return body + } + return body[:limit] + "...(truncated)" +} + +func mapKeys(m map[string]any) []string { + keys := make([]string, 0, len(m)) + for k := range m { + keys = append(keys, k) + } + return keys +} + +func waitForHealth(t *testing.T, url string) { + t.Helper() + for i := 0; i < 60; i++ { + resp, err := http.Get(url) + if err == nil { + resp.Body.Close() + if resp.StatusCode == http.StatusOK { + return + } + } + time.Sleep(time.Second) + } + t.Fatalf("never became healthy: %s", url) +} diff --git a/internal/server/auth_integration_test.go b/internal/server/auth_integration_test.go new file mode 100644 index 0000000..4216fdb --- /dev/null +++ b/internal/server/auth_integration_test.go @@ -0,0 +1,183 @@ +package server + +import ( + "bytes" + "context" + "net/http" + "net/http/httptest" + "testing" + + "github.com/harshaneel/localaik/internal/pdf" + openaip "github.com/harshaneel/localaik/internal/protocol/openai" +) + +// Every upstream route must carry the proxy's credential and none of the +// caller's. +func TestUpstreamAuthHeaderReachesEveryUpstreamPath(t *testing.T) { + cases := []struct { + name string + method string + path string + body string + }{ + {"openai_chat", http.MethodPost, "/v1/chat/completions", `{"model":"m","messages":[]}`}, + {"openai_models", http.MethodGet, "/v1/models", ""}, + {"openai_completions", http.MethodPost, "/v1/completions", `{"prompt":"hello"}`}, + {"openai_model_get", http.MethodGet, "/v1/models/gpt-4", ""}, + {"gemini_generate", http.MethodPost, "/v1beta/models/m:generateContent", `{"contents":[{"role":"user","parts":[{"text":"hi"}]}]}`}, + {"gemini_count_tokens", http.MethodPost, "/v1beta/models/m:countTokens", `{"contents":[{"parts":[{"text":"hi"}]}]}`}, + {"gemini_models_list", http.MethodGet, "/v1beta/models", ""}, + {"gemini_model_get", http.MethodGet, "/v1beta/models/gemini-2.5-pro", ""}, + {"anthropic_messages", http.MethodPost, "/v1/messages", `{"max_tokens":8,"messages":[{"role":"user","content":"hi"}]}`}, + {"anthropic_count_tokens", http.MethodPost, "/v1/messages/count_tokens", `{"messages":[{"role":"user","content":"hi"}]}`}, + {"health", http.MethodGet, "/health", ""}, + } + + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + var called bool + var seenAuth, seenClientAuth, seenAPIKey, seenGoogKey string + + upstream := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + called = true + seenAuth = r.Header.Get("X-Proxy-Token") + seenClientAuth = r.Header.Get("Authorization") + seenAPIKey = r.Header.Get("X-Api-Key") + seenGoogKey = r.Header.Get("X-Goog-Api-Key") + + switch r.URL.Path { + case "/tokenize": + writeJSON(w, http.StatusOK, map[string]any{"tokens": []int{1, 2}}) + case "/v1/models": + writeJSON(w, http.StatusOK, openaip.ModelList{Object: "list", Data: []openaip.Model{{ID: "m"}}}) + case "/v1/models/gpt-4": + writeJSON(w, http.StatusOK, openaip.Model{ID: "gpt-4"}) + case "/health": + w.WriteHeader(http.StatusOK) + default: + writeJSON(w, http.StatusOK, openaip.ChatCompletionResponse{ + Choices: []openaip.Choice{{Message: openaip.Message{Content: "ok"}, FinishReason: "stop"}}, + }) + } + }) + + srv, err := New(Config{ + UpstreamBaseURL: "http://upstream.test/v1", + UpstreamAuthHeader: "X-Proxy-Token: upstream-secret", + HTTPClient: &http.Client{Transport: roundTripHandler{handler: upstream}}, + PDFRenderer: pdf.RendererFunc(func(context.Context, []byte) ([][]byte, error) { return nil, nil }), + }) + if err != nil { + t.Fatalf("New returned error: %v", err) + } + + var reader *bytes.Buffer + if tc.body != "" { + reader = bytes.NewBufferString(tc.body) + } else { + reader = bytes.NewBuffer(nil) + } + req := httptest.NewRequest(tc.method, tc.path, reader) + req.Header.Set("Authorization", "Bearer client-secret") + req.Header.Set("X-Api-Key", "client-anthropic-key") + req.Header.Set("X-Goog-Api-Key", "client-google-key") + + rec := httptest.NewRecorder() + srv.ServeHTTP(rec, req) + + if rec.Code != http.StatusOK { + t.Fatalf("status = %d, want 200; body=%s", rec.Code, rec.Body.String()) + } + if !called { + t.Fatal("upstream was never called, so the header check proves nothing") + } + if seenAuth != "upstream-secret" { + t.Fatalf("X-Proxy-Token = %q, want the proxy credential", seenAuth) + } + if seenClientAuth != "" || seenAPIKey != "" || seenGoogKey != "" { + t.Fatalf("client credentials leaked upstream: auth=%q apikey=%q googkey=%q", seenClientAuth, seenAPIKey, seenGoogKey) + } + }) + } +} + +func TestNoUpstreamAuthHeaderWhenUnset(t *testing.T) { + var seen http.Header + + upstream := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + seen = r.Header.Clone() + writeJSON(w, http.StatusOK, openaip.ChatCompletionResponse{ + Choices: []openaip.Choice{{Message: openaip.Message{Content: "ok"}, FinishReason: "stop"}}, + }) + }) + + srv := newTestServer(t, upstream) + + req := httptest.NewRequest(http.MethodPost, "/v1/messages", bytes.NewBufferString(`{"max_tokens":8,"messages":[{"role":"user","content":"hi"}]}`)) + rec := httptest.NewRecorder() + srv.ServeHTTP(rec, req) + + if rec.Code != http.StatusOK { + t.Fatalf("status = %d, want 200; body=%s", rec.Code, rec.Body.String()) + } + if got := seen.Get("Authorization"); got != "" { + t.Fatalf("Authorization = %q, want none when no credential is configured", got) + } +} + +// The Gemini streaming route builds its own request; confirm the credential is +// present there too. +func TestUpstreamAuthHeaderOnStreamingRoute(t *testing.T) { + var seen string + + upstream := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + seen = r.Header.Get("X-Proxy-Token") + w.Header().Set("Content-Type", "text/event-stream") + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte("data: {\"choices\":[{\"delta\":{\"content\":\"hi\"},\"finish_reason\":\"stop\"}]}\n\ndata: [DONE]\n\n")) + }) + + srv, err := New(Config{ + UpstreamBaseURL: "http://upstream.test/v1", + UpstreamAuthHeader: "X-Proxy-Token: upstream-secret", + HTTPClient: &http.Client{Transport: roundTripHandler{handler: upstream}}, + PDFRenderer: pdf.RendererFunc(func(context.Context, []byte) ([][]byte, error) { return nil, nil }), + }) + if err != nil { + t.Fatalf("New returned error: %v", err) + } + + body := `{"contents":[{"role":"user","parts":[{"text":"hi"}]}]}` + req := httptest.NewRequest(http.MethodPost, "/v1beta/models/m:streamGenerateContent", bytes.NewBufferString(body)) + rec := httptest.NewRecorder() + srv.ServeHTTP(rec, req) + + if rec.Code != http.StatusOK { + t.Fatalf("status = %d, want 200; body=%s", rec.Code, rec.Body.String()) + } + if seen != "upstream-secret" { + t.Fatalf("X-Proxy-Token = %q on the streaming route, want the proxy credential", seen) + } + if !bytes.Contains(rec.Body.Bytes(), []byte("data:")) { + t.Fatalf("response body missing data: frame; got %s", rec.Body.String()) + } +} + +func TestNewDoesNotMutateCallerClient(t *testing.T) { + sentinelTransport := &http.Transport{DisableCompression: true} + clientToPass := &http.Client{Transport: sentinelTransport} + + _, err := New(Config{ + UpstreamBaseURL: "http://upstream.test/v1", + UpstreamAuthHeader: "X-Proxy-Token: upstream-secret", + HTTPClient: clientToPass, + PDFRenderer: pdf.RendererFunc(func(context.Context, []byte) ([][]byte, error) { return nil, nil }), + }) + if err != nil { + t.Fatalf("New returned error: %v", err) + } + + if clientToPass.Transport != sentinelTransport { + t.Fatal("New mutated the caller's http.Client; it should have made a copy") + } +} diff --git a/internal/server/server.go b/internal/server/server.go index 3254b6b..47fd268 100644 --- a/internal/server/server.go +++ b/internal/server/server.go @@ -16,9 +16,10 @@ import ( ) type Config struct { - UpstreamBaseURL string - HTTPClient *http.Client - PDFRenderer pdf.Renderer + UpstreamBaseURL string + UpstreamAuthHeader string + HTTPClient *http.Client + PDFRenderer pdf.Renderer } type Server struct { @@ -48,6 +49,17 @@ func New(cfg Config) (*Server, error) { if client == nil { client = &http.Client{} } + if cfg.UpstreamAuthHeader != "" { + clone := *client + clone.Transport = newUpstreamAuthTransport(clone.Transport, cfg.UpstreamAuthHeader, parsed.Hostname()) + // Scoped here so the bundled-model images keep the stdlib default. The + // transport pins hostname only, so relaxing this reopens port and + // scheme leaks. + clone.CheckRedirect = func(*http.Request, []*http.Request) error { + return http.ErrUseLastResponse + } + client = &clone + } renderer := cfg.PDFRenderer if renderer == nil { diff --git a/internal/server/upstreamauth.go b/internal/server/upstreamauth.go new file mode 100644 index 0000000..a036db2 --- /dev/null +++ b/internal/server/upstreamauth.go @@ -0,0 +1,65 @@ +package server + +import ( + "net/http" + "strings" + + "golang.org/x/net/http/httpguts" +) + +type upstreamAuthTransport struct { + base http.RoundTripper + host string + name string + value string +} + +// ValidUpstreamAuthHeader reports whether header is a line the transport will +// actually send. Startup warnings must use this, not a second predicate. +func ValidUpstreamAuthHeader(header string) bool { + _, _, ok := parseUpstreamAuthHeader(header) + return ok +} + +// Only the first colon separates the pair, so values may contain further ones. +func parseUpstreamAuthHeader(header string) (string, string, bool) { + name, value, found := strings.Cut(header, ":") + name = strings.TrimSpace(name) + value = strings.TrimSpace(value) + if !found || name == "" || value == "" { + return "", "", false + } + // net/http rejects these at the wire on every request, with an error that + // never mentions the env var that caused it. + if !httpguts.ValidHeaderFieldName(name) || !httpguts.ValidHeaderFieldValue(value) { + return "", "", false + } + return name, value, true +} + +// newUpstreamAuthTransport returns base unchanged when header is not a usable +// "Name: value" line, so a misconfigured value cannot silently drop requests. +func newUpstreamAuthTransport(base http.RoundTripper, header, host string) http.RoundTripper { + if base == nil { + base = http.DefaultTransport + } + + name, value, ok := parseUpstreamAuthHeader(header) + if !ok { + return base + } + + return &upstreamAuthTransport{base: base, host: host, name: name, value: value} +} + +func (t *upstreamAuthTransport) RoundTrip(req *http.Request) (*http.Response, error) { + // A 3xx re-enters here with the target's URL, so setting the header for + // every host would hand the credential to whatever the redirect names. + if !strings.EqualFold(req.URL.Hostname(), t.host) { + return t.base.RoundTrip(req) + } + + clone := req.Clone(req.Context()) + clone.Header.Set(t.name, t.value) + return t.base.RoundTrip(clone) +} diff --git a/internal/server/upstreamauth_redirect_test.go b/internal/server/upstreamauth_redirect_test.go new file mode 100644 index 0000000..71d3639 --- /dev/null +++ b/internal/server/upstreamauth_redirect_test.go @@ -0,0 +1,139 @@ +package server + +import ( + "bytes" + "context" + "net/http" + "net/http/httptest" + "strings" + "sync" + "testing" + + "github.com/harshaneel/localaik/internal/pdf" +) + +// redirectingTransport answers upstream.test with a 302 to location and every +// other host with a marker body, recording the headers each host received. +type redirectingTransport struct { + location string + + mu sync.Mutex + seen map[string]http.Header +} + +func (r *redirectingTransport) RoundTrip(req *http.Request) (*http.Response, error) { + r.mu.Lock() + if r.seen == nil { + r.seen = make(map[string]http.Header) + } + r.seen[req.URL.Hostname()] = req.Header.Clone() + r.mu.Unlock() + + recorder := httptest.NewRecorder() + if req.URL.Hostname() == "upstream.test" { + recorder.Header().Set("Location", r.location) + recorder.WriteHeader(http.StatusFound) + return recorder.Result(), nil + } + + recorder.Header().Set("Content-Type", "application/json") + recorder.WriteHeader(http.StatusOK) + _, _ = recorder.WriteString(`{"leaked":true}`) + return recorder.Result(), nil +} + +func (r *redirectingTransport) headers(host string) (http.Header, bool) { + r.mu.Lock() + defer r.mu.Unlock() + header, ok := r.seen[host] + return header, ok +} + +// A 3xx re-enters the transport with the target's URL, so a transport that sets +// the credential unconditionally hands it to whatever host the redirect names. +func TestUpstreamAuthTransportWithholdsCredentialFromOtherHosts(t *testing.T) { + for _, name := range []string{"Authorization", "X-Proxy-Token"} { + t.Run(name, func(t *testing.T) { + base := &redirectingTransport{location: "http://redirect.test/v1/models"} + client := &http.Client{ + Transport: newUpstreamAuthTransport(base, name+": upstream-secret", "upstream.test"), + } + + resp, err := client.Get("http://upstream.test/v1/models") + if err != nil { + t.Fatalf("Get returned error: %v", err) + } + defer resp.Body.Close() + + configured, ok := base.headers("upstream.test") + if !ok { + t.Fatal("configured upstream was never called") + } + if got := configured.Get(name); got != "upstream-secret" { + t.Fatalf("configured upstream saw %s = %q, want the credential", name, got) + } + + target, ok := base.headers("redirect.test") + if !ok { + t.Fatal("redirect target was never reached, so this test proves nothing") + } + if got := target.Get(name); got != "" { + t.Fatalf("redirect target received the credential: %s = %q", name, got) + } + }) + } +} + +func TestCredentialedClientDoesNotFollowUpstreamRedirects(t *testing.T) { + base := &redirectingTransport{location: "http://redirect.test/v1/chat/completions"} + + srv, err := New(Config{ + UpstreamBaseURL: "http://upstream.test/v1", + UpstreamAuthHeader: "X-Proxy-Token: upstream-secret", + HTTPClient: &http.Client{Transport: base}, + PDFRenderer: pdf.RendererFunc(func(context.Context, []byte) ([][]byte, error) { return nil, nil }), + }) + if err != nil { + t.Fatalf("New returned error: %v", err) + } + + req := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", bytes.NewBufferString(`{"model":"m","messages":[]}`)) + rec := httptest.NewRecorder() + srv.ServeHTTP(rec, req) + + if rec.Code != http.StatusFound { + t.Fatalf("status = %d, want the 302 passed through to the caller", rec.Code) + } + if strings.Contains(rec.Body.String(), "leaked") { + t.Fatalf("caller received the redirect target's body: %s", rec.Body.String()) + } + if _, ok := base.headers("redirect.test"); ok { + t.Fatal("credentialed client followed the redirect") + } +} + +// The model-bundled images configure no credential and must keep the stdlib's +// redirect handling. +func TestClientWithoutCredentialStillFollowsRedirects(t *testing.T) { + base := &redirectingTransport{location: "http://redirect.test/v1/chat/completions"} + + srv, err := New(Config{ + UpstreamBaseURL: "http://upstream.test/v1", + HTTPClient: &http.Client{Transport: base}, + PDFRenderer: pdf.RendererFunc(func(context.Context, []byte) ([][]byte, error) { return nil, nil }), + }) + if err != nil { + t.Fatalf("New returned error: %v", err) + } + + req := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", bytes.NewBufferString(`{"model":"m","messages":[]}`)) + rec := httptest.NewRecorder() + srv.ServeHTTP(rec, req) + + if rec.Code != http.StatusOK { + t.Fatalf("status = %d, want 200 after following the redirect", rec.Code) + } + if !strings.Contains(rec.Body.String(), "leaked") { + t.Fatalf("redirect was not followed; body = %s", rec.Body.String()) + } +} diff --git a/internal/server/upstreamauth_test.go b/internal/server/upstreamauth_test.go new file mode 100644 index 0000000..b94a008 --- /dev/null +++ b/internal/server/upstreamauth_test.go @@ -0,0 +1,123 @@ +package server + +import ( + "net/http" + "net/http/httptest" + "testing" +) + +type capturingTransport struct { + seen http.Header +} + +func (c *capturingTransport) RoundTrip(req *http.Request) (*http.Response, error) { + c.seen = req.Header.Clone() + recorder := httptest.NewRecorder() + recorder.WriteHeader(http.StatusOK) + return recorder.Result(), nil +} + +// The single definition of a usable credential line. A capturing transport +// bypasses net/http's wire-level field checks, so these run at the predicate. +func TestValidUpstreamAuthHeader(t *testing.T) { + cases := []struct { + name string + header string + want bool + }{ + {"typical", "Authorization: Bearer token123", true}, + {"value keeps later colons", "Authorization: Bearer a:b", true}, + {"untrimmed", " X-Api-Key : abc123 ", true}, + {"empty", "", false}, + {"whitespace only", " : ", false}, + {"no colon", "InvalidHeader NoColon", false}, + {"no name", ": value", false}, + {"no value", "Name:", false}, + {"space in name", "Bad Name: secret", false}, + {"newline in value", "Authorization: Bearer a\nX-Evil: b", false}, + {"carriage return in value", "Authorization: Bearer a\rb", false}, + {"null in value", "Authorization: Bearer a\x00b", false}, + } + + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + if got := ValidUpstreamAuthHeader(tc.header); got != tc.want { + t.Fatalf("ValidUpstreamAuthHeader(%q) = %v, want %v", tc.header, got, tc.want) + } + }) + } +} + +// A header the predicate rejects must never reach the wire, where net/http +// would fail every request with an error that does not name the env var. +func TestRejectedHeaderIsNeverSentUpstream(t *testing.T) { + for _, header := range []string{"Bad Name: secret", "Authorization: Bearer a\nX-Evil: b"} { + capture := &capturingTransport{} + transport := newUpstreamAuthTransport(capture, header, "upstream.test") + + req := httptest.NewRequest(http.MethodGet, "http://upstream.test/v1/models", nil) + if _, err := transport.RoundTrip(req); err != nil { + t.Fatalf("RoundTrip returned error for %q: %v", header, err) + } + if len(capture.seen) != 0 { + t.Fatalf("header %q produced %v, want none", header, capture.seen) + } + } +} + +func TestUpstreamAuthTransportAddsHeader(t *testing.T) { + capture := &capturingTransport{} + transport := newUpstreamAuthTransport(capture, "Authorization: Bearer secret", "upstream.test") + + req := httptest.NewRequest(http.MethodGet, "http://upstream.test/v1/models", nil) + if _, err := transport.RoundTrip(req); err != nil { + t.Fatalf("RoundTrip returned error: %v", err) + } + + if got := capture.seen.Get("Authorization"); got != "Bearer secret" { + t.Fatalf("Authorization = %q, want %q", got, "Bearer secret") + } +} + +func TestUpstreamAuthTransportTrimsWhitespace(t *testing.T) { + capture := &capturingTransport{} + transport := newUpstreamAuthTransport(capture, " X-Api-Key : abc123 ", "upstream.test") + + req := httptest.NewRequest(http.MethodGet, "http://upstream.test/v1/models", nil) + if _, err := transport.RoundTrip(req); err != nil { + t.Fatalf("RoundTrip returned error: %v", err) + } + + if got := capture.seen.Get("X-Api-Key"); got != "abc123" { + t.Fatalf("X-Api-Key = %q, want %q", got, "abc123") + } +} + +func TestUpstreamAuthTransportIgnoresMalformedHeader(t *testing.T) { + for _, header := range []string{"", " ", "NoColonHere", ": novalue", "Name:"} { + capture := &capturingTransport{} + transport := newUpstreamAuthTransport(capture, header, "upstream.test") + + req := httptest.NewRequest(http.MethodGet, "http://upstream.test/v1/models", nil) + if _, err := transport.RoundTrip(req); err != nil { + t.Fatalf("RoundTrip returned error for %q: %v", header, err) + } + if len(capture.seen) != 0 { + t.Fatalf("header %q produced %v, want none", header, capture.seen) + } + } +} + +func TestUpstreamAuthTransportDoesNotMutateCallerRequest(t *testing.T) { + capture := &capturingTransport{} + transport := newUpstreamAuthTransport(capture, "Authorization: Bearer secret", "upstream.test") + + req := httptest.NewRequest(http.MethodGet, "http://upstream.test/v1/models", nil) + if _, err := transport.RoundTrip(req); err != nil { + t.Fatalf("RoundTrip returned error: %v", err) + } + + if req.Header.Get("Authorization") != "" { + t.Fatal("RoundTrip mutated the caller's request") + } +}