296 lines
10 KiB
Go
296 lines
10 KiB
Go
package openai
|
|
|
|
import (
|
|
"crypto/sha256"
|
|
"encoding/hex"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"strings"
|
|
"testing"
|
|
|
|
"iop/packages/go/config"
|
|
iop "iop/proto/gen/iop"
|
|
)
|
|
|
|
func sha256Hex(raw string) string {
|
|
sum := sha256.Sum256([]byte(raw))
|
|
return hex.EncodeToString(sum[:])
|
|
}
|
|
|
|
func TestRoutesResolvePrincipalTokenMapping(t *testing.T) {
|
|
const rawToken = "sk-alice-raw-token"
|
|
cfg := config.EdgeOpenAIConf{
|
|
PrincipalTokens: []config.OpenAIPrincipalTokenConf{
|
|
{TokenRef: "iop-tok-alice", TokenHashSHA256: sha256Hex(rawToken), PrincipalRef: "user:alice", PrincipalAlias: "alice"},
|
|
},
|
|
Models: []string{"model-a"},
|
|
}
|
|
srv := NewServer(cfg, &fakeRunService{}, nil)
|
|
|
|
for _, tc := range []struct {
|
|
name string
|
|
auth string
|
|
wantStatus int
|
|
}{
|
|
{name: "missing", wantStatus: http.StatusUnauthorized},
|
|
{name: "wrong token", auth: "Bearer wrong-token", wantStatus: http.StatusUnauthorized},
|
|
{name: "valid hash-mapped token", auth: "Bearer " + rawToken, wantStatus: http.StatusOK},
|
|
} {
|
|
t.Run(tc.name, func(t *testing.T) {
|
|
req := httptest.NewRequest(http.MethodGet, "/v1/models", nil)
|
|
if tc.auth != "" {
|
|
req.Header.Set("Authorization", tc.auth)
|
|
}
|
|
w := httptest.NewRecorder()
|
|
|
|
srv.routes().ServeHTTP(w, req)
|
|
|
|
if w.Code != tc.wantStatus {
|
|
t.Fatalf("status: got %d body=%s", w.Code, w.Body.String())
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestRoutesPrincipalTokenMappingFallsBackToLegacyBearer(t *testing.T) {
|
|
cfg := config.EdgeOpenAIConf{
|
|
PrincipalTokens: []config.OpenAIPrincipalTokenConf{
|
|
{TokenRef: "iop-tok-alice", TokenHashSHA256: sha256Hex("alice-token"), PrincipalRef: "user:alice"},
|
|
},
|
|
BearerToken: "legacy-secret",
|
|
Models: []string{"model-a"},
|
|
}
|
|
srv := NewServer(cfg, &fakeRunService{}, nil)
|
|
|
|
req := httptest.NewRequest(http.MethodGet, "/v1/models", nil)
|
|
req.Header.Set("Authorization", "Bearer legacy-secret")
|
|
w := httptest.NewRecorder()
|
|
|
|
srv.routes().ServeHTTP(w, req)
|
|
|
|
if w.Code != http.StatusOK {
|
|
t.Fatalf("status: got %d body=%s", w.Code, w.Body.String())
|
|
}
|
|
}
|
|
|
|
func TestResolvePrincipalAllowsMultipleTokensForSamePrincipal(t *testing.T) {
|
|
cfg := config.EdgeOpenAIConf{
|
|
PrincipalTokens: []config.OpenAIPrincipalTokenConf{
|
|
{TokenRef: "alice-app-a", TokenHashSHA256: sha256Hex("alice-app-a-token"), PrincipalRef: "user:alice", PrincipalAlias: "alice"},
|
|
{TokenRef: "alice-app-b", TokenHashSHA256: sha256Hex("alice-app-b-token"), PrincipalRef: "user:alice", PrincipalAlias: "alice"},
|
|
},
|
|
}
|
|
|
|
first, ok := resolvePrincipal(cfg, "Bearer alice-app-a-token")
|
|
if !ok {
|
|
t.Fatal("expected first app token to authenticate")
|
|
}
|
|
if first.PrincipalRef != "user:alice" || first.PrincipalAlias != "alice" || first.TokenRef != "alice-app-a" {
|
|
t.Fatalf("unexpected first principal: %+v", first)
|
|
}
|
|
|
|
second, ok := resolvePrincipal(cfg, "Bearer alice-app-b-token")
|
|
if !ok {
|
|
t.Fatal("expected second app token to authenticate")
|
|
}
|
|
if second.PrincipalRef != "user:alice" || second.PrincipalAlias != "alice" || second.TokenRef != "alice-app-b" {
|
|
t.Fatalf("unexpected second principal: %+v", second)
|
|
}
|
|
}
|
|
|
|
func TestHealthzDoesNotRequireBearerTokenWithPrincipalTokensConfigured(t *testing.T) {
|
|
cfg := config.EdgeOpenAIConf{
|
|
PrincipalTokens: []config.OpenAIPrincipalTokenConf{
|
|
{TokenRef: "iop-tok-alice", TokenHashSHA256: sha256Hex("alice-token"), PrincipalRef: "user:alice"},
|
|
},
|
|
}
|
|
srv := NewServer(cfg, &fakeRunService{}, nil)
|
|
|
|
req := httptest.NewRequest(http.MethodGet, "/healthz", nil)
|
|
w := httptest.NewRecorder()
|
|
|
|
srv.routes().ServeHTTP(w, req)
|
|
|
|
if w.Code != http.StatusOK {
|
|
t.Fatalf("status: got %d body=%s", w.Code, w.Body.String())
|
|
}
|
|
}
|
|
|
|
func TestChatCompletionsAddsPrincipalMetadataFromBearerToken(t *testing.T) {
|
|
fake := &fakeRunService{events: make(chan *iop.RunEvent, 2)}
|
|
fake.events <- &iop.RunEvent{Type: "delta", Delta: "ok"}
|
|
fake.events <- &iop.RunEvent{Type: "complete"}
|
|
|
|
const rawToken = "sk-alice-raw-token"
|
|
cfg := config.EdgeOpenAIConf{
|
|
Adapter: "ollama",
|
|
PrincipalTokens: []config.OpenAIPrincipalTokenConf{
|
|
{TokenRef: "iop-tok-alice", TokenHashSHA256: sha256Hex(rawToken), PrincipalRef: "user:alice", PrincipalAlias: "alice"},
|
|
},
|
|
}
|
|
srv := NewServer(cfg, fake, nil)
|
|
req := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader(`{
|
|
"model":"from-request",
|
|
"messages":[{"role":"user","content":"hi"}],
|
|
"metadata":{"iop_principal_ref":"spoofed","user":"someone-else"}
|
|
}`))
|
|
req.Header.Set("Authorization", "Bearer "+rawToken)
|
|
w := httptest.NewRecorder()
|
|
|
|
srv.routes().ServeHTTP(w, req)
|
|
|
|
if w.Code != http.StatusOK {
|
|
t.Fatalf("status: got %d body=%s", w.Code, w.Body.String())
|
|
}
|
|
if fake.req.Metadata["iop_principal_ref"] != "user:alice" {
|
|
t.Fatalf("iop_principal_ref must come from the authenticated context, not caller metadata: got %q", fake.req.Metadata["iop_principal_ref"])
|
|
}
|
|
if fake.req.Metadata["iop_principal_alias"] != "alice" {
|
|
t.Fatalf("iop_principal_alias: got %q", fake.req.Metadata["iop_principal_alias"])
|
|
}
|
|
if fake.req.Metadata["iop_token_ref"] != "iop-tok-alice" {
|
|
t.Fatalf("iop_token_ref: got %q", fake.req.Metadata["iop_token_ref"])
|
|
}
|
|
if fake.req.Metadata["iop_principal_source"] != principalSourceToken {
|
|
t.Fatalf("iop_principal_source: got %q, want %q", fake.req.Metadata["iop_principal_source"], principalSourceToken)
|
|
}
|
|
if fake.req.Metadata["user"] != "someone-else" {
|
|
t.Fatalf("caller metadata.user should still pass through as opaque metadata (not used as identity): %+v", fake.req.Metadata)
|
|
}
|
|
}
|
|
|
|
func TestRoutesDoNotUseMetadataUserForPrincipal(t *testing.T) {
|
|
fake := &fakeRunService{events: make(chan *iop.RunEvent, 2)}
|
|
fake.events <- &iop.RunEvent{Type: "delta", Delta: "ok"}
|
|
fake.events <- &iop.RunEvent{Type: "complete"}
|
|
|
|
// No principal_tokens and no bearer_token configured: the endpoint is
|
|
// open (unauthenticated), so no caller identity should be resolved even
|
|
// though the caller supplies metadata.user.
|
|
srv := NewServer(config.EdgeOpenAIConf{Adapter: "ollama"}, fake, nil)
|
|
req := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader(`{
|
|
"model":"from-request",
|
|
"messages":[{"role":"user","content":"hi"}],
|
|
"metadata":{"user":"claimed-identity"}
|
|
}`))
|
|
w := httptest.NewRecorder()
|
|
|
|
srv.routes().ServeHTTP(w, req)
|
|
|
|
if w.Code != http.StatusOK {
|
|
t.Fatalf("status: got %d body=%s", w.Code, w.Body.String())
|
|
}
|
|
if _, ok := fake.req.Metadata["iop_principal_ref"]; ok {
|
|
t.Fatalf("iop_principal_ref must not be derived from metadata.user: %+v", fake.req.Metadata)
|
|
}
|
|
if fake.req.Metadata["iop_principal_source"] != principalSourceUnknown {
|
|
t.Fatalf("iop_principal_source: got %q, want %q", fake.req.Metadata["iop_principal_source"], principalSourceUnknown)
|
|
}
|
|
if fake.req.Metadata["user"] != "claimed-identity" {
|
|
t.Fatalf("caller metadata.user should still pass through as opaque metadata: %+v", fake.req.Metadata)
|
|
}
|
|
}
|
|
|
|
func TestResponsesAddsPrincipalMetadataFromBearerToken(t *testing.T) {
|
|
fake := &fakeRunService{events: make(chan *iop.RunEvent, 2)}
|
|
fake.events <- &iop.RunEvent{Type: "delta", Delta: "ok"}
|
|
fake.events <- &iop.RunEvent{Type: "complete"}
|
|
|
|
const rawToken = "sk-bob-raw-token"
|
|
cfg := config.EdgeOpenAIConf{
|
|
Adapter: "ollama",
|
|
PrincipalTokens: []config.OpenAIPrincipalTokenConf{
|
|
{TokenRef: "iop-tok-bob", TokenHashSHA256: sha256Hex(rawToken), PrincipalRef: "user:bob"},
|
|
},
|
|
}
|
|
srv := NewServer(cfg, fake, nil)
|
|
req := httptest.NewRequest(http.MethodPost, "/v1/responses", strings.NewReader(`{
|
|
"model":"from-request",
|
|
"input":"hi"
|
|
}`))
|
|
req.Header.Set("Authorization", "Bearer "+rawToken)
|
|
w := httptest.NewRecorder()
|
|
|
|
srv.routes().ServeHTTP(w, req)
|
|
|
|
if w.Code != http.StatusOK {
|
|
t.Fatalf("status: got %d body=%s", w.Code, w.Body.String())
|
|
}
|
|
if fake.req.Metadata["iop_principal_ref"] != "user:bob" {
|
|
t.Fatalf("iop_principal_ref: got %q, want user:bob", fake.req.Metadata["iop_principal_ref"])
|
|
}
|
|
if fake.req.Metadata["iop_token_ref"] != "iop-tok-bob" {
|
|
t.Fatalf("iop_token_ref: got %q", fake.req.Metadata["iop_token_ref"])
|
|
}
|
|
if fake.req.Metadata["iop_principal_alias"] != "" {
|
|
t.Fatalf("iop_principal_alias should be empty when unset: got %q", fake.req.Metadata["iop_principal_alias"])
|
|
}
|
|
}
|
|
|
|
func TestProviderTunnelAddsPrincipalMetadataFromBearerToken(t *testing.T) {
|
|
fake := &fakeRunService{tunnelFrames: staticProviderTunnelFrames(`{"ok":true}`)}
|
|
catalog := []config.ModelCatalogEntry{{
|
|
ID: "ornith:35b",
|
|
Providers: map[string]string{"prov-vllm": "Ornith-1.0-35B"},
|
|
}}
|
|
|
|
const rawToken = "sk-carol-raw-token"
|
|
cfg := config.EdgeOpenAIConf{
|
|
PrincipalTokens: []config.OpenAIPrincipalTokenConf{
|
|
{TokenRef: "iop-tok-carol", TokenHashSHA256: sha256Hex(rawToken), PrincipalRef: "user:carol", PrincipalAlias: "carol"},
|
|
},
|
|
}
|
|
srv := NewServer(cfg, fake, nil)
|
|
srv.SetModelCatalog(catalog)
|
|
|
|
req := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader(`{
|
|
"model":"ornith:35b",
|
|
"messages":[{"role":"user","content":"hello"}]
|
|
}`))
|
|
req.Header.Set("Authorization", "Bearer "+rawToken)
|
|
w := httptest.NewRecorder()
|
|
|
|
srv.routes().ServeHTTP(w, req)
|
|
|
|
if w.Code != http.StatusOK {
|
|
t.Fatalf("status: got %d body=%s", w.Code, w.Body.String())
|
|
}
|
|
tunnelReqs := fake.tunnelReqsSnapshot()
|
|
if len(tunnelReqs) != 1 {
|
|
t.Fatalf("expected 1 tunnel dispatch, got %d", len(tunnelReqs))
|
|
}
|
|
meta := tunnelReqs[0].Metadata
|
|
if meta["iop_principal_ref"] != "user:carol" {
|
|
t.Fatalf("iop_principal_ref: got %q, want user:carol", meta["iop_principal_ref"])
|
|
}
|
|
if meta["iop_principal_alias"] != "carol" {
|
|
t.Fatalf("iop_principal_alias: got %q", meta["iop_principal_alias"])
|
|
}
|
|
if meta["iop_token_ref"] != "iop-tok-carol" {
|
|
t.Fatalf("iop_token_ref: got %q", meta["iop_token_ref"])
|
|
}
|
|
}
|
|
|
|
func TestChatCompletionsPrincipalMetadataEmptyWhenUnauthenticatedDirectCall(t *testing.T) {
|
|
// Handlers invoked directly (bypassing withAuth, as most existing tests
|
|
// do) must not fail when no principal was ever resolved for the context.
|
|
fake := &fakeRunService{events: make(chan *iop.RunEvent, 2)}
|
|
fake.events <- &iop.RunEvent{Type: "delta", Delta: "ok"}
|
|
fake.events <- &iop.RunEvent{Type: "complete"}
|
|
|
|
srv := NewServer(config.EdgeOpenAIConf{Adapter: "ollama"}, fake, nil)
|
|
req := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader(`{
|
|
"model":"from-request",
|
|
"messages":[{"role":"user","content":"hi"}]
|
|
}`))
|
|
w := httptest.NewRecorder()
|
|
|
|
srv.handleChatCompletions(w, req)
|
|
|
|
if w.Code != http.StatusOK {
|
|
t.Fatalf("status: got %d body=%s", w.Code, w.Body.String())
|
|
}
|
|
if _, ok := fake.req.Metadata["iop_principal_source"]; ok {
|
|
t.Fatalf("no principal metadata should be added without an authenticated context: %+v", fake.req.Metadata)
|
|
}
|
|
}
|