From 0a1d837f40eae5b54dcb8efe641782519a7a42d3 Mon Sep 17 00:00:00 2001 From: Josh Siviour Date: Thu, 6 Aug 2026 16:17:10 +1000 Subject: [PATCH] feat(bedrock): gate Ochre models behind per-account access grants (#741) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * feat(bedrock): gate Ochre models behind per-account access grants Self-host models were invokable by any authenticated account: tieredCatalog returned them unconditionally and no runtime router checked the caller. GetFoundationModel also described models the caller could not use. Add a bedrock-model-access JetStream KV bucket holding presence-only grants, resolved through a new AccessResolver. Access is deny-by-default: a gateway with no store configured falls back to DenyAllAccessResolver, so an unconfigured deployment serves nothing rather than everything. The system account bypasses grants, matching how quotas exempt it. All four runtime routers share one grantedCatalogEntry gate so they cannot drift on who may invoke what. List and describe agree: an ungranted model is absent from ListFoundationModels and reported as ResourceNotFoundException by GetFoundationModel, so the error cannot confirm a model exists. Administration lands on two surfaces: Grant/Revoke/ListModelAccess actions on the spinifex gateway service for the UI, and `spx admin ochre access` for operators and provisioning, which opens the bucket over NATS directly rather than requiring a SigV4 client the CLI does not have. * feat(bedrock): seed the admin account's model grants at gateway start Deny-by-default left a fresh install with a catalog nobody could see, and the first cut filled that gap from ansible — which is a dev tool, so any install that never ran it came up unusable. Seed from the product instead: on start awsgw grants the platform admin account the full catalog, once, behind a marker key in the same bucket. The marker makes it a bootstrap step rather than a policy: tenant accounts still begin with no access, and an operator who revokes a grant does not have it restored by the next restart. Grants are written before the marker so a crash midway repeats the idempotent grants rather than leaving a half-seeded account, and the marker is a conditional create so only the first node to reach it need win. Seeding is best-effort: a KV failure logs and leaves the gateway serving, and the next start retries because the marker is written only on success. * test(bedrock): cover the model-access CLI and gateway admin actions Route the access CLI runners through the same loadConfigAndConnectFn / ochreExit seam the weights runners use, so grant, revoke and list can be driven against an embedded JetStream server, and add gateway tests for the grant admin actions including the non-admin denial. --- cmd/spinifex/cmd/admin_ochre.go | 178 +++++++++++++ cmd/spinifex/cmd/admin_ochre_test.go | 158 ++++++++++++ spinifex/gateway/bedrock.go | 33 ++- spinifex/gateway/bedrock/access.go | 238 +++++++++++++++++ spinifex/gateway/bedrock/access_test.go | 243 ++++++++++++++++++ spinifex/gateway/bedrock/catalog.go | 93 +++++-- spinifex/gateway/bedrock/catalog_test.go | 99 +++++-- spinifex/gateway/bedrock/contract_test.go | 4 +- spinifex/gateway/bedrock/converse_stream.go | 8 +- .../gateway/bedrock/converse_stream_test.go | 10 +- spinifex/gateway/bedrock/invoke.go | 58 +++-- spinifex/gateway/bedrock/invoke_stream.go | 17 +- .../gateway/bedrock/invoke_stream_test.go | 6 +- spinifex/gateway/bedrock/invoke_test.go | 20 +- spinifex/gateway/bedrock/provider.go | 45 ++-- spinifex/gateway/bedrock/router_test.go | 20 +- spinifex/gateway/bedrock_request_test.go | 60 ++++- spinifex/gateway/bedrockruntime.go | 17 +- spinifex/gateway/gateway.go | 9 + spinifex/gateway/spinifex.go | 42 ++- spinifex/gateway/spinifex_modelaccess_test.go | 128 +++++++++ spinifex/services/awsgw/awsgw.go | 21 ++ tests/e2e/bedrock/main_test.go | 7 + tests/integration/harness.go | 10 + 24 files changed, 1379 insertions(+), 145 deletions(-) create mode 100644 spinifex/gateway/bedrock/access.go create mode 100644 spinifex/gateway/bedrock/access_test.go create mode 100644 spinifex/gateway/spinifex_modelaccess_test.go diff --git a/cmd/spinifex/cmd/admin_ochre.go b/cmd/spinifex/cmd/admin_ochre.go index a8f8fb6de..c3e7e509a 100644 --- a/cmd/spinifex/cmd/admin_ochre.go +++ b/cmd/spinifex/cmd/admin_ochre.go @@ -9,6 +9,7 @@ import ( "os/exec" "path" "path/filepath" + "sort" "strings" "time" @@ -706,3 +707,180 @@ func runOchreWeightsRemove(cmd *cobra.Command, _ []string) { } fmt.Printf("✅ Removed staged-weights entry for %s (snapshot %s and source objects untouched).\n", modelID, entry.SnapshotID) } + +var adminOchreAccessCmd = &cobra.Command{ + Use: "access", + Short: "Manage per-account model access grants", + Long: `Manage which accounts may see and invoke which Ochre models. + +Access is deny-by-default: an account with no grants sees an empty model +catalog and every invocation is refused. Grants are cluster state held in the +bedrock-model-access KV bucket, so these commands need a running cluster but +may be run from any node.`, +} + +var adminOchreAccessGrantCmd = &cobra.Command{ + Use: "grant", + Short: "Grant an account access to a model", + Run: runOchreAccessGrant, +} + +var adminOchreAccessRevokeCmd = &cobra.Command{ + Use: "revoke", + Short: "Revoke an account's access to a model", + Run: runOchreAccessRevoke, +} + +var adminOchreAccessListCmd = &cobra.Command{ + Use: "list", + Short: "List the models an account has been granted", + Run: runOchreAccessList, +} + +var adminOchreModelsCmd = &cobra.Command{ + Use: "models", + Short: "List every model in the platform catalog", + Long: `List the full platform catalog, ignoring grants. This is what an operator can +grant from; it is not what any account can see.`, + Run: runOchreModels, +} + +func init() { + ochreCmd.AddCommand(adminOchreAccessCmd) + ochreCmd.AddCommand(adminOchreModelsCmd) + adminOchreAccessCmd.AddCommand(adminOchreAccessGrantCmd) + adminOchreAccessCmd.AddCommand(adminOchreAccessRevokeCmd) + adminOchreAccessCmd.AddCommand(adminOchreAccessListCmd) + + for _, c := range []*cobra.Command{adminOchreAccessGrantCmd, adminOchreAccessRevokeCmd} { + c.Flags().String("account-id", "", "12-digit account ID to change (required)") + c.Flags().String("model-id", "", "Model ID to change (e.g. meta.llama3-2-1b-instruct-v1:0)") + c.Flags().Bool("all-models", false, "Apply to every model in the platform catalog") + if err := c.MarkFlagRequired("account-id"); err != nil { + panic(err) + } + } + + adminOchreAccessListCmd.Flags().String("account-id", "", "12-digit account ID to inspect (required)") + if err := adminOchreAccessListCmd.MarkFlagRequired("account-id"); err != nil { + panic(err) + } +} + +// ochreAccessStore connects to the cluster and returns the grant store along +// with a cleanup that closes the connection. +func ochreAccessStore() (*gateway_bedrock.ModelAccessStore, func(), error) { + cfg, nc, err := loadConfigAndConnectFn() + if err != nil { + return nil, nil, fmt.Errorf("connect to cluster: %w", err) + } + js, err := jetstream.New(nc) + if err != nil { + nc.Close() + return nil, nil, fmt.Errorf("jetstream context: %w", err) + } + return gateway_bedrock.NewModelAccessStore(js, len(cfg.Nodes)), func() { nc.Close() }, nil +} + +// ochreTargetModels resolves the --model-id / --all-models pair into the model +// set to act on. Exactly one of the two must be given: defaulting either way +// would make a mistyped flag silently change more or less than intended. +func ochreTargetModels(cmd *cobra.Command) ([]string, error) { + modelID, _ := cmd.Flags().GetString("model-id") + allModels, _ := cmd.Flags().GetBool("all-models") + + switch { + case allModels && modelID != "": + return nil, fmt.Errorf("--model-id and --all-models are mutually exclusive") + case allModels: + return gateway_bedrock.CatalogModelIDs(), nil + case modelID != "": + return []string{modelID}, nil + default: + return nil, fmt.Errorf("one of --model-id or --all-models is required") + } +} + +func runOchreAccessGrant(cmd *cobra.Command, _ []string) { + runOchreAccessChange(cmd, true) +} + +func runOchreAccessRevoke(cmd *cobra.Command, _ []string) { + runOchreAccessChange(cmd, false) +} + +// runOchreAccessChange applies a grant or revoke across the resolved model set. +// Both operations are idempotent, so re-running after a partial failure is safe +// and this is callable from provisioning that runs on every deploy. +func runOchreAccessChange(cmd *cobra.Command, grant bool) { + accountID, _ := cmd.Flags().GetString("account-id") + + models, err := ochreTargetModels(cmd) + if err != nil { + fmt.Fprintf(os.Stderr, "Error: %v\n", err) + ochreExit(1) + return + } + + store, cleanup, err := ochreAccessStore() + if err != nil { + fmt.Fprintf(os.Stderr, "Error: %v\n", err) + ochreExit(1) + return + } + defer cleanup() + + ctx := context.Background() + verb := "Granted" + for _, modelID := range models { + if grant { + err = store.Grant(ctx, accountID, modelID) + } else { + verb = "Revoked" + err = store.Revoke(ctx, accountID, modelID) + } + if err != nil { + fmt.Fprintf(os.Stderr, "Error: %v\n", err) + ochreExit(1) + return + } + fmt.Printf("✅ %s %s → %s\n", verb, accountID, modelID) + } +} + +func runOchreAccessList(cmd *cobra.Command, _ []string) { + accountID, _ := cmd.Flags().GetString("account-id") + + store, cleanup, err := ochreAccessStore() + if err != nil { + fmt.Fprintf(os.Stderr, "Error: %v\n", err) + ochreExit(1) + return + } + defer cleanup() + + models, err := store.List(context.Background(), accountID) + if err != nil { + fmt.Fprintf(os.Stderr, "Error: %v\n", err) + ochreExit(1) + return + } + if len(models) == 0 { + fmt.Printf("Account %s has no model grants (it can see and invoke nothing).\n", accountID) + return + } + + // KV key order is not meaningful, so sort for a stable, diffable listing. + sort.Strings(models) + for _, modelID := range models { + fmt.Println(modelID) + } +} + +func runOchreModels(_ *cobra.Command, _ []string) { + models := gateway_bedrock.CatalogModelIDs() + sort.Strings(models) + for _, modelID := range models { + fmt.Println(modelID) + } +} diff --git a/cmd/spinifex/cmd/admin_ochre_test.go b/cmd/spinifex/cmd/admin_ochre_test.go index b6b9ce592..325446ba1 100644 --- a/cmd/spinifex/cmd/admin_ochre_test.go +++ b/cmd/spinifex/cmd/admin_ochre_test.go @@ -1000,3 +1000,161 @@ func TestRegisterWeightsSnapshot_WriteFailureSurfaces(t *testing.T) { require.Error(t, err) assert.Contains(t, err.Error(), "register snapshot metadata") } + +// ochreFlagCmd builds a bare command carrying the model-selection flags, so +// the resolver can be exercised without the cluster connection the real +// subcommands make. +func ochreFlagCmd(t *testing.T, modelID string, allModels bool) *cobra.Command { + t.Helper() + cmd := &cobra.Command{} + cmd.Flags().String("model-id", modelID, "") + cmd.Flags().Bool("all-models", allModels, "") + return cmd +} + +func TestOchreTargetModels_SingleModel(t *testing.T) { + models, err := ochreTargetModels(ochreFlagCmd(t, selfHostModelID, false)) + require.NoError(t, err) + assert.Equal(t, []string{selfHostModelID}, models) +} + +func TestOchreTargetModels_AllModels(t *testing.T) { + models, err := ochreTargetModels(ochreFlagCmd(t, "", true)) + require.NoError(t, err) + assert.Equal(t, gateway_bedrock.CatalogModelIDs(), models) + assert.NotEmpty(t, models) +} + +// TestOchreTargetModels_NeitherFlag pins the refusal to guess: silently +// defaulting to the whole catalog would over-grant on a typo. +func TestOchreTargetModels_NeitherFlag(t *testing.T) { + _, err := ochreTargetModels(ochreFlagCmd(t, "", false)) + require.Error(t, err) + assert.Contains(t, err.Error(), "--model-id or --all-models") +} + +func TestOchreTargetModels_BothFlags(t *testing.T) { + _, err := ochreTargetModels(ochreFlagCmd(t, selfHostModelID, true)) + require.Error(t, err) + assert.Contains(t, err.Error(), "mutually exclusive") +} + +// ochreAccessCmd builds a command carrying the access subcommands' full flag +// set, so the grant/revoke/list runners can be driven without the cobra +// command tree. +func ochreAccessCmd(t *testing.T, accountID, modelID string, allModels bool) *cobra.Command { + t.Helper() + cmd := ochreFlagCmd(t, modelID, allModels) + cmd.Flags().String("account-id", accountID, "") + return cmd +} + +const accessTestAccount = "000000000001" + +// stubAccessConnect dials a fresh connection per call, because each access +// runner closes the connection it was handed; a single shared one would leave +// the next runner in a sequence talking to a closed conn. +func stubAccessConnect(t *testing.T, clientURL string) { + t.Helper() + orig := loadConfigAndConnectFn + t.Cleanup(func() { loadConfigAndConnectFn = orig }) + loadConfigAndConnectFn = func() (*config.ClusterConfig, *nats.Conn, error) { + nc, err := nats.Connect(clientURL) + if err != nil { + return nil, nil, err + } + t.Cleanup(nc.Close) + return fakeClusterConfig(), nc, nil + } +} + +// TestRunOchreAccessChange_ModelFlagFailureExits1 pins that the flag refusal +// happens before any cluster connection: an operator who mistypes a flag gets +// told so without a NATS dial. +func TestRunOchreAccessChange_ModelFlagFailureExits1(t *testing.T) { + stubConnect(t, nil, nil, errors.New("connect must not be reached")) + + cmd := ochreAccessCmd(t, accessTestAccount, "", false) + code := withOchreExitCapture(t, func() { runOchreAccessGrant(cmd, nil) }) + assert.Equal(t, 1, code) +} + +func TestRunOchreAccessGrant_ConnectFailureExits1(t *testing.T) { + stubConnect(t, nil, nil, errors.New("dial failed")) + + cmd := ochreAccessCmd(t, accessTestAccount, selfHostModelID, false) + code := withOchreExitCapture(t, func() { runOchreAccessGrant(cmd, nil) }) + assert.Equal(t, 1, code) +} + +func TestRunOchreAccessList_ConnectFailureExits1(t *testing.T) { + stubConnect(t, nil, nil, errors.New("dial failed")) + + cmd := ochreAccessCmd(t, accessTestAccount, "", false) + code := withOchreExitCapture(t, func() { runOchreAccessList(cmd, nil) }) + assert.Equal(t, 1, code) +} + +// TestRunOchreAccessGrantListRevoke drives the three runners in sequence +// against an embedded JetStream server, so the CLI is checked against the same +// grant store the gateway resolves against rather than a fake. +func TestRunOchreAccessGrantListRevoke(t *testing.T) { + srv, _, _ := testutil.StartTestJetStream(t) + stubAccessConnect(t, srv.ClientURL()) + + grantCmd := ochreAccessCmd(t, accessTestAccount, selfHostModelID, false) + out := captureStdout(t, func() { runOchreAccessGrant(grantCmd, nil) }) + assert.Contains(t, out, "Granted") + assert.Contains(t, out, selfHostModelID) + + listCmd := ochreAccessCmd(t, accessTestAccount, "", false) + out = captureStdout(t, func() { runOchreAccessList(listCmd, nil) }) + assert.Contains(t, out, selfHostModelID) + assert.NotContains(t, out, providerModelID, "list must show only the granted model") + + revokeCmd := ochreAccessCmd(t, accessTestAccount, selfHostModelID, false) + out = captureStdout(t, func() { runOchreAccessRevoke(revokeCmd, nil) }) + assert.Contains(t, out, "Revoked") + + out = captureStdout(t, func() { runOchreAccessList(listCmd, nil) }) + assert.Contains(t, out, "no model grants") +} + +// TestRunOchreAccessGrant_AllModelsGrantsWholeCatalog covers the --all-models +// fan-out, which is what provisioning uses. +func TestRunOchreAccessGrant_AllModelsGrantsWholeCatalog(t *testing.T) { + srv, _, _ := testutil.StartTestJetStream(t) + stubAccessConnect(t, srv.ClientURL()) + + grantCmd := ochreAccessCmd(t, accessTestAccount, "", true) + captureStdout(t, func() { runOchreAccessGrant(grantCmd, nil) }) + + listCmd := ochreAccessCmd(t, accessTestAccount, "", false) + out := captureStdout(t, func() { runOchreAccessList(listCmd, nil) }) + for _, modelID := range gateway_bedrock.CatalogModelIDs() { + assert.Contains(t, out, modelID) + } +} + +// TestRunOchreAccessList_OtherAccountSeesNothing is the deny-by-default rule +// as an operator observes it: an account nobody granted has an empty listing. +func TestRunOchreAccessList_OtherAccountSeesNothing(t *testing.T) { + srv, _, _ := testutil.StartTestJetStream(t) + stubAccessConnect(t, srv.ClientURL()) + + grantCmd := ochreAccessCmd(t, accessTestAccount, selfHostModelID, false) + captureStdout(t, func() { runOchreAccessGrant(grantCmd, nil) }) + + listCmd := ochreAccessCmd(t, "000000000002", "", false) + out := captureStdout(t, func() { runOchreAccessList(listCmd, nil) }) + assert.Contains(t, out, "no model grants") +} + +// TestRunOchreModels_PrintsSortedCatalog checks the listing operators consult +// before granting: every catalog ID, one per line. +func TestRunOchreModels_PrintsSortedCatalog(t *testing.T) { + out := captureStdout(t, func() { runOchreModels(nil, nil) }) + for _, modelID := range gateway_bedrock.CatalogModelIDs() { + assert.Contains(t, out, modelID) + } +} diff --git a/spinifex/gateway/bedrock.go b/spinifex/gateway/bedrock.go index 1c29d105c..14c26f890 100644 --- a/spinifex/gateway/bedrock.go +++ b/spinifex/gateway/bedrock.go @@ -26,23 +26,25 @@ type bedrockRoute struct { // bedrockRouteHandler invokes a per-action bedrock (control-plane) gateway // function. params holds the regex capture groups, PathUnescape'd. resolver // is gw.bedrockResolver(): the configured credential store, or a no-op -// fallback. loggingStore is gw.bedrockLoggingConfigStore(). -type bedrockRouteHandler func(ctx context.Context, accountID string, params []string, body []byte, resolver gateway_bedrock.CredentialResolver, loggingStore *gateway_bedrock.LoggingConfigStore) (any, error) +// fallback. loggingStore is gw.bedrockLoggingConfigStore(). access is +// gw.bedrockAccessResolver(): the configured grant store, or a deny-all +// fallback. +type bedrockRouteHandler func(ctx context.Context, accountID string, params []string, body []byte, resolver gateway_bedrock.CredentialResolver, loggingStore *gateway_bedrock.LoggingConfigStore, access gateway_bedrock.AccessResolver) (any, error) // bedrockRoutes is the dispatch table. More-specific paths must precede // less-specific ones with the same prefix so the regex matcher picks the // deeper route first. var bedrockRoutes = []bedrockRoute{ {"GET", regexp.MustCompile(`^/foundation-models$`), "ListFoundationModels", - func(ctx context.Context, acct string, p []string, b []byte, resolver gateway_bedrock.CredentialResolver, _ *gateway_bedrock.LoggingConfigStore) (any, error) { - return gateway_bedrock.ListFoundationModels(ctx, acct, resolver, new(bedrock.ListFoundationModelsInput)) + func(ctx context.Context, acct string, p []string, b []byte, resolver gateway_bedrock.CredentialResolver, _ *gateway_bedrock.LoggingConfigStore, access gateway_bedrock.AccessResolver) (any, error) { + return gateway_bedrock.ListFoundationModels(ctx, acct, resolver, access, new(bedrock.ListFoundationModelsInput)) }}, {"GET", regexp.MustCompile(`^/foundation-models/([^/]+)$`), "GetFoundationModel", - func(ctx context.Context, acct string, p []string, b []byte, _ gateway_bedrock.CredentialResolver, _ *gateway_bedrock.LoggingConfigStore) (any, error) { - return gateway_bedrock.GetFoundationModel(ctx, acct, p[0]) + func(ctx context.Context, acct string, p []string, b []byte, _ gateway_bedrock.CredentialResolver, _ *gateway_bedrock.LoggingConfigStore, access gateway_bedrock.AccessResolver) (any, error) { + return gateway_bedrock.GetFoundationModel(ctx, acct, p[0], access) }}, {"PUT", regexp.MustCompile(`^/logging/modelinvocations$`), "PutModelInvocationLoggingConfiguration", - func(ctx context.Context, acct string, p []string, b []byte, _ gateway_bedrock.CredentialResolver, store *gateway_bedrock.LoggingConfigStore) (any, error) { + func(ctx context.Context, acct string, p []string, b []byte, _ gateway_bedrock.CredentialResolver, store *gateway_bedrock.LoggingConfigStore, _ gateway_bedrock.AccessResolver) (any, error) { input := new(bedrock.PutModelInvocationLoggingConfigurationInput) if len(b) > 0 { if err := json.Unmarshal(b, input); err != nil { @@ -52,11 +54,11 @@ var bedrockRoutes = []bedrockRoute{ return gateway_bedrock.PutModelInvocationLoggingConfiguration(ctx, acct, store, input) }}, {"GET", regexp.MustCompile(`^/logging/modelinvocations$`), "GetModelInvocationLoggingConfiguration", - func(ctx context.Context, acct string, p []string, b []byte, _ gateway_bedrock.CredentialResolver, store *gateway_bedrock.LoggingConfigStore) (any, error) { + func(ctx context.Context, acct string, p []string, b []byte, _ gateway_bedrock.CredentialResolver, store *gateway_bedrock.LoggingConfigStore, _ gateway_bedrock.AccessResolver) (any, error) { return gateway_bedrock.GetModelInvocationLoggingConfiguration(ctx, acct, store, new(bedrock.GetModelInvocationLoggingConfigurationInput)) }}, {"DELETE", regexp.MustCompile(`^/logging/modelinvocations$`), "DeleteModelInvocationLoggingConfiguration", - func(ctx context.Context, acct string, p []string, b []byte, _ gateway_bedrock.CredentialResolver, store *gateway_bedrock.LoggingConfigStore) (any, error) { + func(ctx context.Context, acct string, p []string, b []byte, _ gateway_bedrock.CredentialResolver, store *gateway_bedrock.LoggingConfigStore, _ gateway_bedrock.AccessResolver) (any, error) { return gateway_bedrock.DeleteModelInvocationLoggingConfiguration(ctx, acct, store, new(bedrock.DeleteModelInvocationLoggingConfigurationInput)) }}, } @@ -121,7 +123,7 @@ func (gw *GatewayConfig) Bedrock_Request(w http.ResponseWriter, r *http.Request) return errors.New(awserrors.ErrorInvalidParameterValue) } - output, err := handler(r.Context(), accountID, params, body, gw.bedrockResolver(), gw.bedrockLoggingConfigStore()) + output, err := handler(r.Context(), accountID, params, body, gw.bedrockResolver(), gw.bedrockLoggingConfigStore(), gw.bedrockAccessResolver()) if err != nil { return err } @@ -160,6 +162,17 @@ func (gw *GatewayConfig) bedrockRecorder() gateway_bedrock.Recorder { return gateway_bedrock.NoopRecorder } +// bedrockAccessResolver returns gw.BedrockAccess as an AccessResolver, or the +// deny-all fallback when no grant store is configured. Model access is +// deny-by-default, so an unconfigured gateway advertises and serves no models +// rather than all of them. +func (gw *GatewayConfig) bedrockAccessResolver() gateway_bedrock.AccessResolver { + if gw.BedrockAccess != nil { + return gw.BedrockAccess + } + return gateway_bedrock.DenyAllAccessResolver +} + // bedrockEndpointResolver returns an EndpointResolver over the configured // pinned self-host endpoints (gw.BedrockEndpoints). A nil/empty map resolves // nothing, so self-host models return ModelNotReady until configured. diff --git a/spinifex/gateway/bedrock/access.go b/spinifex/gateway/bedrock/access.go new file mode 100644 index 000000000..7b1c269ff --- /dev/null +++ b/spinifex/gateway/bedrock/access.go @@ -0,0 +1,238 @@ +package gateway_bedrock + +import ( + "context" + "encoding/base64" + "errors" + "fmt" + "strings" + "sync" + + "github.com/mulgadc/spinifex/spinifex/kvutil" + "github.com/mulgadc/spinifex/spinifex/utils" + "github.com/nats-io/nats.go/jetstream" +) + +// modelAccessBucket is the cluster-replicated KV bucket holding per-account +// model-access grants. A grant is presence-only: the key exists iff the +// account may use the model, so the value carries no meaning. +const modelAccessBucket = "bedrock-model-access" + +// modelAccessHistory keeps one revision; a grant is a set membership, not a +// series, so history buys nothing. +const modelAccessHistory = 1 + +// modelAccessPrefix namespaces grant keys within the bucket, leaving room for +// other per-account access state alongside them. +const modelAccessPrefix = "grants" + +// accessKey returns the KV key for accountID's grant on modelID. Model IDs +// contain ':' (e.g. "anthropic.claude-3-5-sonnet-20240620-v1:0"), which NATS +// rejects in a KV key, so the model segment is base64url-encoded: unambiguous +// and reversible, which List relies on. +func accessKey(accountID, modelID string) string { + return fmt.Sprintf("%s/%s/%s", modelAccessPrefix, accountID, base64.RawURLEncoding.EncodeToString([]byte(modelID))) +} + +// modelAccessSeedPrefix namespaces the one-shot seeding markers. It is kept +// out of modelAccessPrefix so a marker can never be read back as a grant. +const modelAccessSeedPrefix = "seeded" + +// accessSeedKey returns the key marking accountID's initial grants as seeded. +func accessSeedKey(accountID string) string { + return fmt.Sprintf("%s/%s", modelAccessSeedPrefix, accountID) +} + +// accountGrantPrefix returns the key prefix covering every grant held by +// accountID. +func accountGrantPrefix(accountID string) string { + return fmt.Sprintf("%s/%s/", modelAccessPrefix, accountID) +} + +// AccessResolver reports whether an account may see and invoke a model. +// Access is deny-by-default: a model is usable only where a grant exists, so +// an unconfigured deployment serves nothing rather than everything. +type AccessResolver interface { + Granted(ctx context.Context, accountID, modelID string) (bool, error) +} + +// ModelAccessStore resolves per-account model grants from the +// bedrock-model-access JetStream KV bucket. +type ModelAccessStore struct { + js jetstream.JetStream + replicas int + + mu sync.Mutex + kv jetstream.KeyValue +} + +var _ AccessResolver = (*ModelAccessStore)(nil) + +// NewModelAccessStore constructs a ModelAccessStore over the cluster's +// JetStream client, replicated across replicas nodes. +func NewModelAccessStore(js jetstream.JetStream, replicas int) *ModelAccessStore { + return &ModelAccessStore{js: js, replicas: replicas} +} + +// bucket lazily opens (or creates) the bedrock-model-access KV bucket, +// caching the handle, mirroring WeightsStore.bucket. +func (s *ModelAccessStore) bucket(ctx context.Context) (jetstream.KeyValue, error) { + s.mu.Lock() + defer s.mu.Unlock() + if s.kv != nil { + return s.kv, nil + } + kv, err := kvutil.GetOrCreateBucketWithReplicas(ctx, s.js, modelAccessBucket, modelAccessHistory, s.replicas) + if err != nil { + return nil, err + } + s.kv = kv + return kv, nil +} + +// Granted reports whether accountID holds a grant on modelID. The system +// account bypasses grants entirely, matching how handlers_quota exempts it +// from every quota dimension. +func (s *ModelAccessStore) Granted(ctx context.Context, accountID, modelID string) (bool, error) { + if accountID == utils.GlobalAccountID { + return true, nil + } + kv, err := s.bucket(ctx) + if err != nil { + return false, err + } + switch _, err := kv.Get(ctx, accessKey(accountID, modelID)); { + case err == nil: + return true, nil + case errors.Is(err, jetstream.ErrKeyNotFound): + return false, nil + default: + return false, fmt.Errorf("kv get grant for %s/%s: %w", accountID, modelID, err) + } +} + +// Grant gives accountID access to modelID. It is idempotent: re-granting an +// existing grant is a no-op rather than an error. +func (s *ModelAccessStore) Grant(ctx context.Context, accountID, modelID string) error { + kv, err := s.bucket(ctx) + if err != nil { + return err + } + if _, err := kv.Put(ctx, accessKey(accountID, modelID), nil); err != nil { + return fmt.Errorf("kv put grant for %s/%s: %w", accountID, modelID, err) + } + return nil +} + +// Revoke removes accountID's access to modelID. Revoking a grant that does +// not exist succeeds, so callers need not check first. +func (s *ModelAccessStore) Revoke(ctx context.Context, accountID, modelID string) error { + kv, err := s.bucket(ctx) + if err != nil { + return err + } + if err := kv.Delete(ctx, accessKey(accountID, modelID)); err != nil && !errors.Is(err, jetstream.ErrKeyNotFound) { + return fmt.Errorf("kv delete grant for %s/%s: %w", accountID, modelID, err) + } + return nil +} + +// SeedAccountGrants grants accountID every model in modelIDs, once per +// deployment, and reports whether it did. It exists so a fresh install has a +// working operator account without the deny-by-default rule being softened: +// only the account named here is seeded, and only until an operator changes +// it, since the marker means a later revoke is never undone by a restart. +// +// Grants are written before the marker, so a crash midway leaves the marker +// absent and the next start repeats the (idempotent) grants rather than +// leaving a half-seeded account. The marker itself is a conditional create: +// every node runs this at startup and only the first need win. +func (s *ModelAccessStore) SeedAccountGrants(ctx context.Context, accountID string, modelIDs []string) (bool, error) { + kv, err := s.bucket(ctx) + if err != nil { + return false, err + } + + switch _, err := kv.Get(ctx, accessSeedKey(accountID)); { + case err == nil: + return false, nil + case errors.Is(err, jetstream.ErrKeyNotFound): + // Not seeded yet — fall through and seed. + default: + return false, fmt.Errorf("kv get seed marker for %s: %w", accountID, err) + } + + for _, modelID := range modelIDs { + if err := s.Grant(ctx, accountID, modelID); err != nil { + return false, err + } + } + + if _, err := kv.Create(ctx, accessSeedKey(accountID), nil); err != nil && !errors.Is(err, jetstream.ErrKeyExists) { + return false, fmt.Errorf("kv create seed marker for %s: %w", accountID, err) + } + return true, nil +} + +// List returns the model IDs accountID holds grants on, in no particular +// order. Keys that do not decode are skipped rather than failing the call, so +// one malformed key cannot hide an account's whole grant set. +func (s *ModelAccessStore) List(ctx context.Context, accountID string) ([]string, error) { + kv, err := s.bucket(ctx) + if err != nil { + return nil, err + } + keys, err := kv.Keys(ctx) + if err != nil { + if errors.Is(err, jetstream.ErrNoKeysFound) { + return nil, nil + } + return nil, fmt.Errorf("kv keys for %s: %w", accountID, err) + } + + prefix := accountGrantPrefix(accountID) + var models []string + for _, key := range keys { + encoded, ok := strings.CutPrefix(key, prefix) + if !ok { + continue + } + modelID, err := base64.RawURLEncoding.DecodeString(encoded) + if err != nil { + continue + } + models = append(models, string(modelID)) + } + return models, nil +} + +// ModelAccessChange is the response to a grant or revoke: it echoes what was +// changed so a caller driving the API from a script can log the effect without +// a follow-up list. +type ModelAccessChange struct { + AccountID string `json:"AccountId"` + ModelID string `json:"ModelId"` +} + +// ModelAccessList is the response to a grant listing. +type ModelAccessList struct { + AccountID string `json:"AccountId"` + ModelIDs []string `json:"ModelIds"` +} + +// denyAllAccessResolver is the zero-value fallback used wherever no +// ModelAccessStore is configured. It grants nothing, so a gateway built +// without an access store serves no models rather than all of them. +type denyAllAccessResolver struct{} + +var _ AccessResolver = (*denyAllAccessResolver)(nil) + +func (denyAllAccessResolver) Granted(_ context.Context, _, _ string) (bool, error) { + return false, nil +} + +// DenyAllAccessResolver grants no model to any account. It is the fallback for +// a nil GatewayConfig.BedrockAccess: the insecure direction of this default is +// "nothing works", which surfaces immediately, rather than "everything works", +// which surfaces as a breach. +var DenyAllAccessResolver AccessResolver = denyAllAccessResolver{} diff --git a/spinifex/gateway/bedrock/access_test.go b/spinifex/gateway/bedrock/access_test.go new file mode 100644 index 000000000..041a35f4b --- /dev/null +++ b/spinifex/gateway/bedrock/access_test.go @@ -0,0 +1,243 @@ +package gateway_bedrock + +import ( + "context" + "errors" + "testing" + + "github.com/mulgadc/spinifex/spinifex/awserrors" + "github.com/mulgadc/spinifex/spinifex/testutil" + "github.com/mulgadc/spinifex/spinifex/utils" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// Shipped catalog entries, one per tier, referenced by the access tests. +// Naming them here keeps a catalog change to one edit per tier rather than one +// per assertion. +const ( + selfHostTestModel = "meta.llama3-2-1b-instruct-v1:0" + anthropicTestModel = "anthropic.claude-3-5-sonnet-20240620-v1:0" +) + +// grantSet is an AccessResolver granting exactly the model IDs it contains. +// The shipped routers are exercised through it so a test can express "this +// account may use that model" without standing up JetStream. +type grantSet map[string]bool + +var _ AccessResolver = grantSet(nil) + +func (g grantSet) Granted(_ context.Context, _, modelID string) (bool, error) { + return g[modelID], nil +} + +// grantAll is an AccessResolver granting every model, for tests whose subject +// is something other than access control. +type grantAll struct{} + +var _ AccessResolver = grantAll{} + +func (grantAll) Granted(_ context.Context, _, _ string) (bool, error) { return true, nil } + +// failingAccess reports an error from every check, covering the path where the +// grant store itself is unhealthy. +type failingAccess struct{} + +var _ AccessResolver = failingAccess{} + +func (failingAccess) Granted(_ context.Context, _, _ string) (bool, error) { + return false, errors.New("kv unavailable") +} + +func TestDenyAllAccessResolver_GrantsNothing(t *testing.T) { + for _, modelID := range []string{selfHostTestModel, anthropicTestModel, "nonexistent.model-v1:0"} { + granted, err := DenyAllAccessResolver.Granted(context.Background(), "000000000001", modelID) + require.NoError(t, err) + assert.False(t, granted, "model %q must not be granted by the deny-all fallback", modelID) + } +} + +func TestAccessKey_EncodesModelID(t *testing.T) { + // Model IDs contain ':', which NATS rejects in a KV key, so the model + // segment must be encoded rather than interpolated raw. + key := accessKey("000000000001", "anthropic.claude-3-5-sonnet-20240620-v1:0") + assert.NotContains(t, key, ":") + assert.Greater(t, len(key), len(accountGrantPrefix("000000000001"))) +} + +// TestModelAccessStore_GrantResolveRevoke exercises the real JetStream KV +// path: lazy bucket create, Grant, the hit and miss branches of Granted, and +// Revoke. +func TestModelAccessStore_GrantResolveRevoke(t *testing.T) { + _, _, js := testutil.StartTestJetStream(t) + store := NewModelAccessStore(js, 1) + ctx := context.Background() + + granted, err := store.Granted(ctx, "000000000001", selfHostTestModel) + require.NoError(t, err) + assert.False(t, granted, "a model must be denied before it is granted") + + require.NoError(t, store.Grant(ctx, "000000000001", selfHostTestModel)) + granted, err = store.Granted(ctx, "000000000001", selfHostTestModel) + require.NoError(t, err) + assert.True(t, granted) + + // A grant is scoped to one account and one model. + granted, err = store.Granted(ctx, "000000000002", selfHostTestModel) + require.NoError(t, err) + assert.False(t, granted, "a grant must not leak to another account") + granted, err = store.Granted(ctx, "000000000001", anthropicTestModel) + require.NoError(t, err) + assert.False(t, granted, "a grant must not leak to another model") + + require.NoError(t, store.Revoke(ctx, "000000000001", selfHostTestModel)) + granted, err = store.Granted(ctx, "000000000001", selfHostTestModel) + require.NoError(t, err) + assert.False(t, granted) +} + +func TestModelAccessStore_GrantIsIdempotent(t *testing.T) { + _, _, js := testutil.StartTestJetStream(t) + store := NewModelAccessStore(js, 1) + ctx := context.Background() + + require.NoError(t, store.Grant(ctx, "000000000001", selfHostTestModel)) + require.NoError(t, store.Grant(ctx, "000000000001", selfHostTestModel)) + + models, err := store.List(ctx, "000000000001") + require.NoError(t, err) + assert.Equal(t, []string{selfHostTestModel}, models, "a repeated grant must not duplicate the entry") +} + +// TestModelAccessStore_RevokeMissingGrantSucceeds covers the ErrKeyNotFound +// branch: callers revoke without checking first, so a missing grant is not an +// error. +func TestModelAccessStore_RevokeMissingGrantSucceeds(t *testing.T) { + _, _, js := testutil.StartTestJetStream(t) + store := NewModelAccessStore(js, 1) + + assert.NoError(t, store.Revoke(context.Background(), "000000000001", selfHostTestModel)) +} + +// TestModelAccessStore_List_RoundTripsModelIDs proves the key encoding is +// reversible, including the ':' that forced it. +func TestModelAccessStore_List_RoundTripsModelIDs(t *testing.T) { + _, _, js := testutil.StartTestJetStream(t) + store := NewModelAccessStore(js, 1) + ctx := context.Background() + + require.NoError(t, store.Grant(ctx, "000000000001", selfHostTestModel)) + require.NoError(t, store.Grant(ctx, "000000000001", anthropicTestModel)) + require.NoError(t, store.Grant(ctx, "000000000002", selfHostTestModel)) + + models, err := store.List(ctx, "000000000001") + require.NoError(t, err) + assert.ElementsMatch(t, []string{selfHostTestModel, anthropicTestModel}, models) + + models, err = store.List(ctx, "000000000002") + require.NoError(t, err) + assert.Equal(t, []string{selfHostTestModel}, models, "List must not return another account's grants") +} + +func TestModelAccessStore_ListEmptyAccount(t *testing.T) { + _, _, js := testutil.StartTestJetStream(t) + store := NewModelAccessStore(js, 1) + + models, err := store.List(context.Background(), "000000000009") + require.NoError(t, err) + assert.Empty(t, models) +} + +// TestModelAccessStore_SystemAccountBypassesGrants mirrors how +// handlers_quota exempts the system account from every quota dimension. +func TestModelAccessStore_SystemAccountBypassesGrants(t *testing.T) { + _, _, js := testutil.StartTestJetStream(t) + store := NewModelAccessStore(js, 1) + + granted, err := store.Granted(context.Background(), utils.GlobalAccountID, selfHostTestModel) + require.NoError(t, err) + assert.True(t, granted, "the system account must bypass grants") +} + +// TestGrantedCatalogEntry_ErrorClasses pins the distinction the runtime paths +// depend on: an unknown model is not found, a known-but-ungranted model is +// access denied, and a broken grant store propagates rather than failing open. +func TestGrantedCatalogEntry_ErrorClasses(t *testing.T) { + ctx := context.Background() + + _, err := grantedCatalogEntry(ctx, "000000000001", "nonexistent.model-v1:0", grantAll{}) + require.Error(t, err) + assert.Equal(t, awserrors.ErrorResourceNotFoundException, err.Error()) + + _, err = grantedCatalogEntry(ctx, "000000000001", selfHostTestModel, grantSet{}) + require.Error(t, err) + assert.Equal(t, awserrors.ErrorAccessDeniedException, err.Error()) + + _, err = grantedCatalogEntry(ctx, "000000000001", selfHostTestModel, failingAccess{}) + require.Error(t, err) + assert.Equal(t, "kv unavailable", err.Error()) + + entry, err := grantedCatalogEntry(ctx, "000000000001", selfHostTestModel, grantSet{selfHostTestModel: true}) + require.NoError(t, err) + assert.Equal(t, selfHostTestModel, entry.ModelID) +} + +// TestSeedAccountGrants_SeedsOnceThenRespectsRevoke is the whole point of the +// seed marker: a fresh install gets a usable operator account, but an operator +// who then revokes a grant does not have it restored by the next restart. +func TestSeedAccountGrants_SeedsOnceThenRespectsRevoke(t *testing.T) { + _, _, js := testutil.StartTestJetStream(t) + store := NewModelAccessStore(js, 1) + ctx := context.Background() + const account = "000000000001" + + seeded, err := store.SeedAccountGrants(ctx, account, CatalogModelIDs()) + require.NoError(t, err) + assert.True(t, seeded, "first call must seed") + + models, err := store.List(ctx, account) + require.NoError(t, err) + assert.ElementsMatch(t, CatalogModelIDs(), models) + + require.NoError(t, store.Revoke(ctx, account, selfHostTestModel)) + + // A restart re-runs the seed; the marker must make it a no-op. + seeded, err = store.SeedAccountGrants(ctx, account, CatalogModelIDs()) + require.NoError(t, err) + assert.False(t, seeded, "second call must not re-seed") + + granted, err := store.Granted(ctx, account, selfHostTestModel) + require.NoError(t, err) + assert.False(t, granted, "a revoked grant must stay revoked across restarts") +} + +// TestSeedAccountGrants_LeavesOtherAccountsDenied keeps the seed narrow: it is +// a bootstrap convenience for the operator's account, not a way for every +// account to end up with the catalog. +func TestSeedAccountGrants_LeavesOtherAccountsDenied(t *testing.T) { + _, _, js := testutil.StartTestJetStream(t) + store := NewModelAccessStore(js, 1) + ctx := context.Background() + + _, err := store.SeedAccountGrants(ctx, "000000000001", CatalogModelIDs()) + require.NoError(t, err) + + granted, err := store.Granted(ctx, "000000000002", selfHostTestModel) + require.NoError(t, err) + assert.False(t, granted, "seeding the admin account must not grant a tenant account") +} + +// TestSeedAccountGrants_MarkerIsNotAGrant guards the key namespaces: a marker +// read back as a grant would show up as a bogus model ID in List. +func TestSeedAccountGrants_MarkerIsNotAGrant(t *testing.T) { + _, _, js := testutil.StartTestJetStream(t) + store := NewModelAccessStore(js, 1) + ctx := context.Background() + + _, err := store.SeedAccountGrants(ctx, "000000000001", nil) + require.NoError(t, err) + + models, err := store.List(ctx, "000000000001") + require.NoError(t, err) + assert.Empty(t, models, "the seed marker must not be listed as a grant") +} diff --git a/spinifex/gateway/bedrock/catalog.go b/spinifex/gateway/bedrock/catalog.go index af51ee4a4..5e0b46b2c 100644 --- a/spinifex/gateway/bedrock/catalog.go +++ b/spinifex/gateway/bedrock/catalog.go @@ -90,6 +90,18 @@ var catalog = []catalogEntry{ }, } +// CatalogModelIDs returns every model ID the platform knows about, ignoring +// grants and credential tiering. It exists for administration — granting an +// account the whole catalog — and must not be used to answer an account's own +// ListFoundationModels, which is grant-filtered. +func CatalogModelIDs() []string { + ids := make([]string, 0, len(catalog)) + for _, entry := range catalog { + ids = append(ids, entry.ModelID) + } + return ids +} + // CredentialResolver resolves accountID's usable provider credential for // vendor: a per-account key, else an optional platform default. key is only // meaningful when ok is true. @@ -97,14 +109,20 @@ type CredentialResolver interface { Resolve(ctx context.Context, accountID, vendor string) (key string, ok bool, err error) } -// tieredCatalog returns the catalog entries advertised to accountID: -// self-host entries only where a weights snapshot resolves, provider entries -// only when resolver finds a usable credential. A resolve error (as opposed -// to a clean not-found) is an internal fault, not a servability verdict, and -// aborts the whole list rather than silently thinning it out. -func tieredCatalog(ctx context.Context, accountID string, resolver CredentialResolver) ([]catalogEntry, error) { +// tieredCatalog returns the catalog entries advertised to accountID: those the +// account holds a grant on, and among those, self-host entries only where a +// weights snapshot resolves and provider entries only where resolver finds a +// usable credential. Access is checked first, so a model the account was never +// granted stays hidden even when a platform-default credential would otherwise +// serve it. A weights resolve error is an internal fault, not a servability +// verdict, and aborts the whole list rather than silently thinning it out. +func tieredCatalog(ctx context.Context, accountID string, resolver CredentialResolver, access AccessResolver) ([]catalogEntry, error) { var out []catalogEntry for _, entry := range catalog { + granted, err := access.Granted(ctx, accountID, entry.ModelID) + if err != nil || !granted { + continue + } if entry.Provider == tierSelfHost { _, resolvable, err := currentWeightsResolver().Resolve(ctx, entry.ModelID) if err != nil { @@ -161,8 +179,9 @@ func modelARN(modelID string) string { return fmt.Sprintf(modelARNFormat, modelID) } -// lookupCatalogEntry finds a catalog entry by exact modelId, ignoring tier -// gating (used by the runtime router, which returns its own error class). +// lookupCatalogEntry finds a catalog entry by exact modelId, ignoring both +// tier gating and access grants. Callers on a request path must use +// grantedCatalogEntry instead; this is the raw table lookup underneath it. func lookupCatalogEntry(modelID string) (catalogEntry, bool) { for _, entry := range catalog { if entry.ModelID == modelID { @@ -203,11 +222,36 @@ func LookupServingSpec(modelID string) (spec ServingSpec, found, selfHost bool) }, true, true } -// ListFoundationModels returns the deployment-tiered catalog: self-host -// entries where a weights snapshot resolves, provider entries where a -// credential resolves. -func ListFoundationModels(ctx context.Context, accountID string, resolver CredentialResolver, _ *bedrock.ListFoundationModelsInput) (*bedrock.ListFoundationModelsOutput, error) { - entries, err := tieredCatalog(ctx, accountID, resolver) +// grantedCatalogEntry resolves modelID to its catalog entry and checks +// accountID's grant on it. Every runtime path shares this one gate so the four +// routers cannot drift apart on who may invoke what. +// +// An unknown model and an ungranted one are deliberately distinguishable here +// (ResourceNotFoundException vs AccessDeniedException) because the caller has +// already been told the model exists by its own catalog listing; the describe +// path collapses both to ResourceNotFoundException instead, where that is not +// true. +func grantedCatalogEntry(ctx context.Context, accountID, modelID string, access AccessResolver) (catalogEntry, error) { + entry, ok := lookupCatalogEntry(modelID) + if !ok { + return catalogEntry{}, errors.New(awserrors.ErrorResourceNotFoundException) + } + granted, err := access.Granted(ctx, accountID, modelID) + if err != nil { + return catalogEntry{}, err + } + if !granted { + return catalogEntry{}, errors.New(awserrors.ErrorAccessDeniedException) + } + return entry, nil +} + +// ListFoundationModels returns the catalog visible to accountID: models it +// holds a grant on, with self-host entries further filtered to those whose +// weights snapshot resolves and provider entries to those whose credential +// resolves. +func ListFoundationModels(ctx context.Context, accountID string, resolver CredentialResolver, access AccessResolver, _ *bedrock.ListFoundationModelsInput) (*bedrock.ListFoundationModelsOutput, error) { + entries, err := tieredCatalog(ctx, accountID, resolver, access) if err != nil { return nil, err } @@ -218,15 +262,28 @@ func ListFoundationModels(ctx context.Context, accountID string, resolver Creden return &bedrock.ListFoundationModelsOutput{ModelSummaries: summaries}, nil } -// GetFoundationModel looks up a single model by exact modelId. Unknown -// models, and self-host models with no resolvable weights snapshot, return -// ResourceNotFoundException. A weights resolve error is an internal fault, -// not a not-found verdict, and is surfaced (and logged) as such. -func GetFoundationModel(ctx context.Context, _ string, modelID string) (*bedrock.GetFoundationModelOutput, error) { +// GetFoundationModel looks up a single model by exact modelId, gated by the +// caller's grant. An ungranted model, and a self-host model with no resolvable +// weights snapshot, are both reported as ResourceNotFoundException rather than +// AccessDeniedException so describe agrees with list: a model the account +// cannot see does not exist as far as this API is concerned, and the error does +// not confirm the model's existence to an account probing for it. A weights +// resolve error is an internal fault, not a not-found verdict, and is surfaced +// (and logged) as such. +func GetFoundationModel(ctx context.Context, accountID string, modelID string, access AccessResolver) (*bedrock.GetFoundationModelOutput, error) { entry, ok := lookupCatalogEntry(modelID) if !ok { return nil, errors.New(awserrors.ErrorResourceNotFoundException) } + // Grant first: an ungranted account must not learn whether the model is + // servable, only that it is not there. + granted, err := access.Granted(ctx, accountID, modelID) + if err != nil { + return nil, err + } + if !granted { + return nil, errors.New(awserrors.ErrorResourceNotFoundException) + } if entry.Provider == tierSelfHost { _, resolvable, err := currentWeightsResolver().Resolve(ctx, entry.ModelID) if err != nil { diff --git a/spinifex/gateway/bedrock/catalog_test.go b/spinifex/gateway/bedrock/catalog_test.go index e39c21676..aa47ba272 100644 --- a/spinifex/gateway/bedrock/catalog_test.go +++ b/spinifex/gateway/bedrock/catalog_test.go @@ -63,64 +63,127 @@ func modelIDs(out *bedrock.ListFoundationModelsOutput) []string { } func TestListFoundationModels_SelfHostIncludedWhenWeightsResolve(t *testing.T) { - withWeightsResolver(t, stubWeightsResolver{ok: map[string]bool{"meta.llama3-2-1b-instruct-v1:0": true}}) - out, err := ListFoundationModels(context.Background(), "000000000001", stubResolver{ok: map[string]bool{}}, &bedrock.ListFoundationModelsInput{}) + withWeightsResolver(t, stubWeightsResolver{ok: map[string]bool{selfHostTestModel: true}}) + out, err := ListFoundationModels(context.Background(), "000000000001", stubResolver{ok: map[string]bool{}}, + grantAll{}, &bedrock.ListFoundationModelsInput{}) require.NoError(t, err) - assert.Contains(t, modelIDs(out), "meta.llama3-2-1b-instruct-v1:0") + assert.Contains(t, modelIDs(out), selfHostTestModel) } func TestListFoundationModels_SelfHostExcludedWhenWeightsUnresolvable(t *testing.T) { withWeightsResolver(t, stubWeightsResolver{ok: map[string]bool{}}) - out, err := ListFoundationModels(context.Background(), "000000000001", stubResolver{ok: map[string]bool{}}, &bedrock.ListFoundationModelsInput{}) + out, err := ListFoundationModels(context.Background(), "000000000001", stubResolver{ok: map[string]bool{}}, + grantAll{}, &bedrock.ListFoundationModelsInput{}) require.NoError(t, err) - assert.NotContains(t, modelIDs(out), "meta.llama3-2-1b-instruct-v1:0") + assert.NotContains(t, modelIDs(out), selfHostTestModel) } -func TestListFoundationModels_ProviderIncludedWhenResolvable(t *testing.T) { - withWeightsResolver(t, stubWeightsResolver{ok: map[string]bool{}}) - out, err := ListFoundationModels(context.Background(), "000000000001", stubResolver{ok: map[string]bool{"anthropic": true}}, &bedrock.ListFoundationModelsInput{}) +func TestListFoundationModels_SelfHostIncludedWhenGranted(t *testing.T) { + withWeightsResolver(t, stubWeightsResolver{ok: map[string]bool{selfHostTestModel: true}}) + out, err := ListFoundationModels(context.Background(), "000000000001", stubResolver{ok: map[string]bool{}}, + grantSet{selfHostTestModel: true}, &bedrock.ListFoundationModelsInput{}) require.NoError(t, err) - assert.Contains(t, modelIDs(out), "anthropic.claude-3-5-sonnet-20240620-v1:0") + assert.Contains(t, modelIDs(out), selfHostTestModel) +} + +// TestListFoundationModels_SelfHostExcludedWhenUngranted is the behaviour +// change: self-host entries used to be advertised to every account +// unconditionally, so this is the regression guard for that. The weights +// resolve here, so the exclusion can only be the missing grant. +func TestListFoundationModels_SelfHostExcludedWhenUngranted(t *testing.T) { + withWeightsResolver(t, stubWeightsResolver{ok: map[string]bool{selfHostTestModel: true}}) + out, err := ListFoundationModels(context.Background(), "000000000001", stubResolver{ok: map[string]bool{}}, + grantSet{}, &bedrock.ListFoundationModelsInput{}) + require.NoError(t, err) + assert.NotContains(t, modelIDs(out), selfHostTestModel) + assert.Empty(t, modelIDs(out), "an account with no grants must see an empty catalog") +} + +func TestListFoundationModels_ProviderIncludedWhenGrantedAndResolvable(t *testing.T) { + out, err := ListFoundationModels(context.Background(), "000000000001", stubResolver{ok: map[string]bool{"anthropic": true}}, + grantSet{anthropicTestModel: true}, &bedrock.ListFoundationModelsInput{}) + require.NoError(t, err) + assert.Contains(t, modelIDs(out), anthropicTestModel) } func TestListFoundationModels_ProviderExcludedWhenUnresolvable(t *testing.T) { withWeightsResolver(t, stubWeightsResolver{ok: map[string]bool{}}) - out, err := ListFoundationModels(context.Background(), "000000000001", stubResolver{ok: map[string]bool{}}, &bedrock.ListFoundationModelsInput{}) + out, err := ListFoundationModels(context.Background(), "000000000001", stubResolver{ok: map[string]bool{}}, + grantSet{anthropicTestModel: true}, &bedrock.ListFoundationModelsInput{}) require.NoError(t, err) - assert.NotContains(t, modelIDs(out), "anthropic.claude-3-5-sonnet-20240620-v1:0") + assert.NotContains(t, modelIDs(out), anthropicTestModel) } func TestGetFoundationModel_KnownModelWithResolvableWeights(t *testing.T) { - withWeightsResolver(t, stubWeightsResolver{ok: map[string]bool{"meta.llama3-2-1b-instruct-v1:0": true}}) - out, err := GetFoundationModel(context.Background(), "000000000001", "meta.llama3-2-1b-instruct-v1:0") + withWeightsResolver(t, stubWeightsResolver{ok: map[string]bool{selfHostTestModel: true}}) + out, err := GetFoundationModel(context.Background(), "000000000001", selfHostTestModel, grantAll{}) require.NoError(t, err) require.NotNil(t, out.ModelDetails) - assert.Equal(t, "meta.llama3-2-1b-instruct-v1:0", *out.ModelDetails.ModelId) + assert.Equal(t, selfHostTestModel, *out.ModelDetails.ModelId) } func TestGetFoundationModel_SelfHostWithUnresolvableWeightsReturnsNotFound(t *testing.T) { withWeightsResolver(t, stubWeightsResolver{ok: map[string]bool{}}) - _, err := GetFoundationModel(context.Background(), "000000000001", "meta.llama3-2-1b-instruct-v1:0") + _, err := GetFoundationModel(context.Background(), "000000000001", selfHostTestModel, grantAll{}) require.Error(t, err) assert.Equal(t, "ResourceNotFoundException", err.Error()) } +// TestListFoundationModels_GrantDoesNotOverrideCredentialTier keeps the two +// filters independent: a grant says the account may use the model, not that +// the platform can reach it. +func TestListFoundationModels_ProviderExcludedWhenUngrantedButResolvable(t *testing.T) { + out, err := ListFoundationModels(context.Background(), "000000000001", stubResolver{ok: map[string]bool{"anthropic": true}}, + grantSet{}, &bedrock.ListFoundationModelsInput{}) + require.NoError(t, err) + assert.NotContains(t, modelIDs(out), anthropicTestModel, + "a platform-default credential must not advertise a model the account was never granted") +} + +// TestCatalogModelIDs_CoversWholeCatalog keeps the admin-facing listing in step +// with the catalog: a model missing here would be ungrantable via --all-models +// and so invisible to every account. +func TestCatalogModelIDs_CoversWholeCatalog(t *testing.T) { + ids := CatalogModelIDs() + require.Len(t, ids, len(catalog)) + assert.Contains(t, ids, selfHostTestModel) + assert.Contains(t, ids, anthropicTestModel) +} + +func TestGetFoundationModel_KnownGrantedModel(t *testing.T) { + withWeightsResolver(t, stubWeightsResolver{ok: map[string]bool{selfHostTestModel: true}}) + out, err := GetFoundationModel(context.Background(), "000000000001", selfHostTestModel, grantSet{selfHostTestModel: true}) + require.NoError(t, err) + require.NotNil(t, out.ModelDetails) + assert.Equal(t, selfHostTestModel, *out.ModelDetails.ModelId) +} + func TestGetFoundationModel_UnknownModelReturnsNotFound(t *testing.T) { - _, err := GetFoundationModel(context.Background(), "000000000001", "does-not-exist") + _, err := GetFoundationModel(context.Background(), "000000000001", "does-not-exist", grantAll{}) + require.Error(t, err) + assert.Equal(t, "ResourceNotFoundException", err.Error()) +} + +// TestGetFoundationModel_UngrantedModelReturnsNotFound pins describe to the +// same answer as list: an ungranted model is reported as absent rather than +// forbidden, so the error cannot be used to confirm the model exists. +func TestGetFoundationModel_UngrantedModelReturnsNotFound(t *testing.T) { + _, err := GetFoundationModel(context.Background(), "000000000001", selfHostTestModel, grantSet{}) require.Error(t, err) assert.Equal(t, "ResourceNotFoundException", err.Error()) } func TestGetFoundationModel_WeightsResolveErrorIsNotResourceNotFound(t *testing.T) { withWeightsResolver(t, errWeightsResolver{}) - _, err := GetFoundationModel(context.Background(), "000000000001", "meta.llama3-2-1b-instruct-v1:0") + _, err := GetFoundationModel(context.Background(), "000000000001", selfHostTestModel, grantAll{}) require.Error(t, err) assert.NotEqual(t, "ResourceNotFoundException", err.Error()) } func TestListFoundationModels_WeightsResolveErrorPropagates(t *testing.T) { withWeightsResolver(t, errWeightsResolver{}) - _, err := ListFoundationModels(context.Background(), "000000000001", stubResolver{ok: map[string]bool{}}, &bedrock.ListFoundationModelsInput{}) + _, err := ListFoundationModels(context.Background(), "000000000001", stubResolver{ok: map[string]bool{}}, + grantAll{}, &bedrock.ListFoundationModelsInput{}) require.Error(t, err) assert.NotEqual(t, "ResourceNotFoundException", err.Error()) } diff --git a/spinifex/gateway/bedrock/contract_test.go b/spinifex/gateway/bedrock/contract_test.go index ed582ff7f..c4a0479cf 100644 --- a/spinifex/gateway/bedrock/contract_test.go +++ b/spinifex/gateway/bedrock/contract_test.go @@ -56,7 +56,7 @@ func TestContract_ConverseStream_SelfHost_RealSDKConsumerDecodesFrames(t *testin if !assert.NoError(t, err) { return } - err = ConverseStream(r.Context(), w, "000000000001", m[1], body, nil, endpoints, nil) + err = ConverseStream(r.Context(), w, "000000000001", m[1], body, nil, endpoints, nil, grantAll{}) assert.NoError(t, err) }) @@ -200,7 +200,7 @@ func TestContract_InvokeModelWithResponseStream_SelfHost_RealSDKConsumerDecodesF if !assert.NoError(t, err) { return } - err = InvokeModelWithResponseStream(r.Context(), w, "000000000001", m[1], body, nil, endpoints, r.Header.Get("Content-Type"), nil) + err = InvokeModelWithResponseStream(r.Context(), w, "000000000001", m[1], body, nil, endpoints, r.Header.Get("Content-Type"), nil, grantAll{}) assert.NoError(t, err) }) diff --git a/spinifex/gateway/bedrock/converse_stream.go b/spinifex/gateway/bedrock/converse_stream.go index 81224761d..9bdf12d9b 100644 --- a/spinifex/gateway/bedrock/converse_stream.go +++ b/spinifex/gateway/bedrock/converse_stream.go @@ -95,12 +95,12 @@ type converseStreamSource interface { // ConverseStream is the bedrock-runtime ConverseStream entry point used by // the gateway route table. Unlike the JSON-dispatch handlers it owns w -// directly: a pre-stream failure (unknown model, unresolved credential, -// upstream connect error) returns an awserrors code for the normal +// directly: a pre-stream failure (unknown model, ungranted model, unresolved +// credential, upstream connect error) returns an awserrors code for the normal // ErrorHandler envelope. Once the first frame is written it always returns // nil — any later failure is an in-band exception event, since the HTTP // status can no longer change. -func ConverseStream(ctx context.Context, w http.ResponseWriter, accountID, modelID string, body []byte, resolver CredentialResolver, endpointResolver EndpointResolver, recorder Recorder) error { +func ConverseStream(ctx context.Context, w http.ResponseWriter, accountID, modelID string, body []byte, resolver CredentialResolver, endpointResolver EndpointResolver, recorder Recorder, access AccessResolver) error { if recorder == nil { recorder = NoopRecorder } @@ -116,7 +116,7 @@ func ConverseStream(ctx context.Context, w http.ResponseWriter, accountID, model entry, _ := lookupCatalogEntry(modelID) // Router.ConverseStream below re-validates; only its Provider tag is needed here. - src, err := NewRouter(resolver, endpointResolver, recorder).ConverseStream(ctx, accountID, modelID, input) + src, err := NewRouter(resolver, endpointResolver, recorder, access).ConverseStream(ctx, accountID, modelID, input) if err != nil { return err } diff --git a/spinifex/gateway/bedrock/converse_stream_test.go b/spinifex/gateway/bedrock/converse_stream_test.go index 2a4a3a6ec..6616ed8f2 100644 --- a/spinifex/gateway/bedrock/converse_stream_test.go +++ b/spinifex/gateway/bedrock/converse_stream_test.go @@ -154,7 +154,7 @@ func TestConverseStream_UnknownModelReturnsResourceNotFoundPreHeader(t *testing. rec := httptest.NewRecorder() body := []byte(`{"messages":[{"role":"user","content":[{"text":"hi"}]}]}`) - err := ConverseStream(context.Background(), rec, "000000000001", "does.not-exist-v1:0", body, nil, nil, nil) + err := ConverseStream(context.Background(), rec, "000000000001", "does.not-exist-v1:0", body, nil, nil, nil, grantAll{}) require.Error(t, err) assert.Equal(t, awserrors.ErrorResourceNotFoundException, err.Error()) // A pre-stream failure must not have written anything: the gateway's @@ -164,7 +164,7 @@ func TestConverseStream_UnknownModelReturnsResourceNotFoundPreHeader(t *testing. func TestConverseStream_MalformedBodyReturnsValidationException(t *testing.T) { rec := httptest.NewRecorder() - err := ConverseStream(context.Background(), rec, "000000000001", "meta.llama3-2-1b-instruct-v1:0", []byte("{not-json"), nil, nil, nil) + err := ConverseStream(context.Background(), rec, "000000000001", "meta.llama3-2-1b-instruct-v1:0", []byte("{not-json"), nil, nil, nil, grantAll{}) require.Error(t, err) assert.Equal(t, awserrors.ErrorValidationException, err.Error()) } @@ -186,7 +186,7 @@ func TestConverseStream_SelfHostHappyPath_WritesFramedTaxonomy(t *testing.T) { }) require.NoError(t, err) - err = ConverseStream(context.Background(), rec, "000000000001", modelID, body, nil, NewStaticEndpointResolver(map[string]string{modelID: ts.URL}), nil) + err = ConverseStream(context.Background(), rec, "000000000001", modelID, body, nil, NewStaticEndpointResolver(map[string]string{modelID: ts.URL}), nil, grantAll{}) require.NoError(t, err) assert.Equal(t, http.StatusOK, rec.Code) @@ -301,7 +301,7 @@ func TestConverseStream_ClientDisconnectAbortsUpstreamRequest(t *testing.T) { done := make(chan struct{}) go func() { defer close(done) - _ = ConverseStream(ctx, rec, "000000000001", modelID, body, nil, NewStaticEndpointResolver(map[string]string{modelID: ts.URL}), nil) + _ = ConverseStream(ctx, rec, "000000000001", modelID, body, nil, NewStaticEndpointResolver(map[string]string{modelID: ts.URL}), nil, grantAll{}) }() select { @@ -337,7 +337,7 @@ func TestConverseStream_NonFlusherWriter_ReturnsPreHeaderError(t *testing.T) { }) require.NoError(t, err) - err = ConverseStream(context.Background(), w, "000000000001", modelID, body, nil, NewStaticEndpointResolver(map[string]string{modelID: ts.URL}), nil) + err = ConverseStream(context.Background(), w, "000000000001", modelID, body, nil, NewStaticEndpointResolver(map[string]string{modelID: ts.URL}), nil, grantAll{}) require.Error(t, err) assert.Equal(t, awserrors.ErrorInternalError, err.Error()) assert.False(t, w.wroteHeader) diff --git a/spinifex/gateway/bedrock/invoke.go b/spinifex/gateway/bedrock/invoke.go index 6ba425930..efd246747 100644 --- a/spinifex/gateway/bedrock/invoke.go +++ b/spinifex/gateway/bedrock/invoke.go @@ -25,13 +25,14 @@ type InvokeRouter struct { resolver CredentialResolver endpointResolver EndpointResolver recorder Recorder + access AccessResolver } -// NewInvokeRouter constructs an InvokeRouter. A nil resolver, -// endpointResolver, or recorder falls back to a no-op implementation, so an -// InvokeRouter is always safe to use even before the real stores are wired -// in. -func NewInvokeRouter(resolver CredentialResolver, endpointResolver EndpointResolver, recorder Recorder) *InvokeRouter { +// NewInvokeRouter constructs an InvokeRouter. A nil resolver, endpointResolver, +// or recorder falls back to a no-op implementation, and a nil access falls back +// to denying every model, so an InvokeRouter is always safe to use even before +// the real stores are wired in. +func NewInvokeRouter(resolver CredentialResolver, endpointResolver EndpointResolver, recorder Recorder, access AccessResolver) *InvokeRouter { if resolver == nil { resolver = NoopCredentialResolver } @@ -41,13 +42,17 @@ func NewInvokeRouter(resolver CredentialResolver, endpointResolver EndpointResol if recorder == nil { recorder = NoopRecorder } - return &InvokeRouter{resolver: resolver, endpointResolver: endpointResolver, recorder: recorder} + if access == nil { + access = DenyAllAccessResolver + } + return &InvokeRouter{resolver: resolver, endpointResolver: endpointResolver, recorder: recorder, access: access} } // InvokeModel routes modelID to its family adapter via the catalog. Unknown -// modelIds and unresolvable vendors return ResourceNotFoundException; a -// vendor with no resolvable credential returns AccessDeniedException. Every -// exit records an InvocationRecord via the deferred closure. +// modelIds and unresolvable vendors return ResourceNotFoundException; an +// ungranted model, or a vendor with no resolvable credential, returns +// AccessDeniedException. Every exit records an InvocationRecord via the +// deferred closure. func (rt *InvokeRouter) InvokeModel(ctx context.Context, accountID, modelID string, body []byte) (respBody []byte, contentType string, err error) { requestID := uuid.NewString() start := time.Now() @@ -71,9 +76,10 @@ func (rt *InvokeRouter) InvokeModel(ctx context.Context, accountID, modelID stri }) }() - entry, ok := lookupCatalogEntry(modelID) - if !ok { - err = errors.New(awserrors.ErrorResourceNotFoundException) + // Assigns the named err, so the deferred closure records a denied + // invocation the same as any other failed one. + entry, err := grantedCatalogEntry(ctx, accountID, modelID, rt.access) + if err != nil { return nil, "", err } backend = entry.Provider @@ -110,10 +116,10 @@ func (rt *InvokeRouter) InvokeModel(ctx context.Context, accountID, modelID stri } // InvokeModel is the bedrock-runtime InvokeModel entry point used by the -// gateway route table. resolver, endpointResolver, and recorder may be nil; -// NewInvokeRouter supplies no-op fallbacks. -func InvokeModel(ctx context.Context, accountID, modelID string, body []byte, resolver CredentialResolver, endpointResolver EndpointResolver, recorder Recorder) ([]byte, string, error) { - return NewInvokeRouter(resolver, endpointResolver, recorder).InvokeModel(ctx, accountID, modelID, body) +// gateway route table. resolver, endpointResolver, recorder and access may be +// nil; NewInvokeRouter supplies no-op (and, for access, deny-all) fallbacks. +func InvokeModel(ctx context.Context, accountID, modelID string, body []byte, resolver CredentialResolver, endpointResolver EndpointResolver, recorder Recorder, access AccessResolver) ([]byte, string, error) { + return NewInvokeRouter(resolver, endpointResolver, recorder, access).InvokeModel(ctx, accountID, modelID, body) } // InvokeStreamAdapter is the optional streaming capability an InvokeAdapter @@ -130,29 +136,33 @@ type InvokeStreamAdapter interface { type InvokeStreamRouter struct { resolver CredentialResolver endpointResolver EndpointResolver + access AccessResolver } // NewInvokeStreamRouter constructs an InvokeStreamRouter. A nil resolver or -// endpointResolver falls back to a resolver/resolver that finds nothing, so -// an InvokeStreamRouter is always safe to use even before the real stores -// are wired in. -func NewInvokeStreamRouter(resolver CredentialResolver, endpointResolver EndpointResolver) *InvokeStreamRouter { +// endpointResolver falls back to a resolver/resolver that finds nothing, and a +// nil access falls back to denying every model, so an InvokeStreamRouter is +// always safe to use even before the real stores are wired in. +func NewInvokeStreamRouter(resolver CredentialResolver, endpointResolver EndpointResolver, access AccessResolver) *InvokeStreamRouter { if resolver == nil { resolver = NoopCredentialResolver } if endpointResolver == nil { endpointResolver = NewStaticEndpointResolver(nil) } - return &InvokeStreamRouter{resolver: resolver, endpointResolver: endpointResolver} + if access == nil { + access = DenyAllAccessResolver + } + return &InvokeStreamRouter{resolver: resolver, endpointResolver: endpointResolver, access: access} } // InvokeModelWithResponseStream routes modelID to its family adapter via the // catalog, exactly like InvokeRouter.InvokeModel, then requires the resolved // adapter to also implement InvokeStreamAdapter. func (rt *InvokeStreamRouter) InvokeModelWithResponseStream(ctx context.Context, accountID, modelID string, body []byte) (invokeStreamSource, error) { - entry, ok := lookupCatalogEntry(modelID) - if !ok { - return nil, errors.New(awserrors.ErrorResourceNotFoundException) + entry, err := grantedCatalogEntry(ctx, accountID, modelID, rt.access) + if err != nil { + return nil, err } var a InvokeAdapter diff --git a/spinifex/gateway/bedrock/invoke_stream.go b/spinifex/gateway/bedrock/invoke_stream.go index 84c44dc38..9553e5bc8 100644 --- a/spinifex/gateway/bedrock/invoke_stream.go +++ b/spinifex/gateway/bedrock/invoke_stream.go @@ -23,14 +23,15 @@ type invokeStreamSource interface { // InvokeModelWithResponseStream is the bedrock-runtime // InvokeModelWithResponseStream entry point used by the gateway route table. // Unlike the JSON-dispatch handlers it owns w directly: a pre-stream failure -// (unknown model, unresolved credential, upstream connect error) returns an -// awserrors code for the normal ErrorHandler envelope. Once the first frame -// is written it always returns nil — any later failure is an in-band -// exception event, since the HTTP status can no longer change. +// (unknown model, ungranted model, unresolved credential, upstream connect +// error) returns an awserrors code for the normal ErrorHandler envelope. Once +// the first frame is written it always returns nil — any later failure is an +// in-band exception event, since the HTTP status can no longer change. // requestContentType is the client's declared Content-Type, logged only. -// resolver, endpointResolver, and recorder may be nil; NewInvokeStreamRouter -// and the internal NoopRecorder fallback keep this call safe either way. -func InvokeModelWithResponseStream(ctx context.Context, w http.ResponseWriter, accountID, modelID string, body []byte, resolver CredentialResolver, endpointResolver EndpointResolver, requestContentType string, recorder Recorder) error { +// resolver, endpointResolver, recorder and access may be nil; +// NewInvokeStreamRouter and the internal NoopRecorder fallback keep this call +// safe either way. +func InvokeModelWithResponseStream(ctx context.Context, w http.ResponseWriter, accountID, modelID string, body []byte, resolver CredentialResolver, endpointResolver EndpointResolver, requestContentType string, recorder Recorder, access AccessResolver) error { if recorder == nil { recorder = NoopRecorder } @@ -39,7 +40,7 @@ func InvokeModelWithResponseStream(ctx context.Context, w http.ResponseWriter, a entry, _ := lookupCatalogEntry(modelID) // InvokeStreamRouter below re-validates; only its Provider tag is needed here. - src, err := NewInvokeStreamRouter(resolver, endpointResolver).InvokeModelWithResponseStream(ctx, accountID, modelID, body) + src, err := NewInvokeStreamRouter(resolver, endpointResolver, access).InvokeModelWithResponseStream(ctx, accountID, modelID, body) if err != nil { return err } diff --git a/spinifex/gateway/bedrock/invoke_stream_test.go b/spinifex/gateway/bedrock/invoke_stream_test.go index c1052157e..6897de2de 100644 --- a/spinifex/gateway/bedrock/invoke_stream_test.go +++ b/spinifex/gateway/bedrock/invoke_stream_test.go @@ -119,7 +119,7 @@ func TestPumpInvokeStream_RecordsInvocationOnUpstreamFault(t *testing.T) { func TestInvokeModelWithResponseStream_UnknownModelReturnsResourceNotFoundPreHeader(t *testing.T) { rec := httptest.NewRecorder() - err := InvokeModelWithResponseStream(context.Background(), rec, "000000000001", "does.not-exist-v1:0", []byte(`{}`), nil, nil, "application/json", nil) + err := InvokeModelWithResponseStream(context.Background(), rec, "000000000001", "does.not-exist-v1:0", []byte(`{}`), nil, nil, "application/json", nil, grantAll{}) require.Error(t, err) assert.Equal(t, awserrors.ErrorResourceNotFoundException, err.Error()) assert.Equal(t, 0, rec.Body.Len()) @@ -137,7 +137,7 @@ func TestInvokeModelWithResponseStream_SelfHostHappyPath_WritesChunkFrames(t *te rec := httptest.NewRecorder() body := []byte(`{"prompt":"hello","max_gen_len":128}`) - err := InvokeModelWithResponseStream(context.Background(), rec, "000000000001", modelID, body, nil, NewStaticEndpointResolver(map[string]string{modelID: ts.URL}), "application/json", nil) + err := InvokeModelWithResponseStream(context.Background(), rec, "000000000001", modelID, body, nil, NewStaticEndpointResolver(map[string]string{modelID: ts.URL}), "application/json", nil, grantAll{}) require.NoError(t, err) assert.Equal(t, http.StatusOK, rec.Code) @@ -162,7 +162,7 @@ func TestInvokeModelWithResponseStream_NonFlusherWriter_ReturnsPreHeaderError(t w := &nonFlusherWriter{} body := []byte(`{"prompt":"hello"}`) - err := InvokeModelWithResponseStream(context.Background(), w, "000000000001", modelID, body, nil, NewStaticEndpointResolver(map[string]string{modelID: ts.URL}), "application/json", nil) + err := InvokeModelWithResponseStream(context.Background(), w, "000000000001", modelID, body, nil, NewStaticEndpointResolver(map[string]string{modelID: ts.URL}), "application/json", nil, grantAll{}) require.Error(t, err) assert.Equal(t, awserrors.ErrorInternalError, err.Error()) assert.False(t, w.wroteHeader) diff --git a/spinifex/gateway/bedrock/invoke_test.go b/spinifex/gateway/bedrock/invoke_test.go index ae19a50a6..083409269 100644 --- a/spinifex/gateway/bedrock/invoke_test.go +++ b/spinifex/gateway/bedrock/invoke_test.go @@ -24,7 +24,7 @@ func TestInvokeRouter_SelfHostSuccess(t *testing.T) { defer ts.Close() modelID := "meta.llama3-2-1b-instruct-v1:0" - rt := NewInvokeRouter(nil, NewStaticEndpointResolver(map[string]string{modelID: ts.URL}), nil) + rt := NewInvokeRouter(nil, NewStaticEndpointResolver(map[string]string{modelID: ts.URL}), nil, grantAll{}) respBody, contentType, err := rt.InvokeModel(context.Background(), "000000000001", modelID, []byte(`{"prompt":"hello"}`)) require.NoError(t, err) @@ -36,21 +36,21 @@ func TestInvokeRouter_SelfHostSuccess(t *testing.T) { } func TestInvokeRouter_UnknownModelReturnsResourceNotFound(t *testing.T) { - rt := NewInvokeRouter(nil, nil, nil) + rt := NewInvokeRouter(nil, nil, nil, grantAll{}) _, _, err := rt.InvokeModel(context.Background(), "000000000001", "does.not-exist-v1:0", []byte(`{}`)) require.Error(t, err) assert.Equal(t, awserrors.ErrorResourceNotFoundException, err.Error()) } func TestInvokeRouter_AnthropicNoCredentialReturnsAccessDenied(t *testing.T) { - rt := NewInvokeRouter(stubCredentialResolver{ok: false}, nil, nil) + rt := NewInvokeRouter(stubCredentialResolver{ok: false}, nil, nil, grantAll{}) _, _, err := rt.InvokeModel(context.Background(), "000000000001", "anthropic.claude-3-5-sonnet-20240620-v1:0", []byte(`{}`)) require.Error(t, err) assert.Equal(t, awserrors.ErrorAccessDeniedException, err.Error()) } func TestInvokeRouter_SelfHostNoEndpointReturnsModelNotReady(t *testing.T) { - rt := NewInvokeRouter(nil, nil, nil) + rt := NewInvokeRouter(nil, nil, nil, grantAll{}) _, _, err := rt.InvokeModel(context.Background(), "000000000001", "meta.llama3-2-1b-instruct-v1:0", []byte(`{"prompt":"hello"}`)) require.Error(t, err) assert.Equal(t, awserrors.ErrorModelNotReadyException, err.Error()) @@ -68,7 +68,7 @@ func TestInvokeModel_PackageEntryPoint_SelfHostSuccess(t *testing.T) { defer ts.Close() modelID := "meta.llama3-2-1b-instruct-v1:0" - respBody, contentType, err := InvokeModel(context.Background(), "000000000001", modelID, []byte(`{"prompt":"hello"}`), nil, NewStaticEndpointResolver(map[string]string{modelID: ts.URL}), nil) + respBody, contentType, err := InvokeModel(context.Background(), "000000000001", modelID, []byte(`{"prompt":"hello"}`), nil, NewStaticEndpointResolver(map[string]string{modelID: ts.URL}), nil, grantAll{}) require.NoError(t, err) assert.Equal(t, "application/json", contentType) @@ -78,7 +78,7 @@ func TestInvokeModel_PackageEntryPoint_SelfHostSuccess(t *testing.T) { } func TestInvokeModel_PackageEntryPoint_UnknownModel(t *testing.T) { - _, _, err := InvokeModel(context.Background(), "000000000001", "does.not-exist-v1:0", []byte(`{}`), nil, nil, nil) + _, _, err := InvokeModel(context.Background(), "000000000001", "does.not-exist-v1:0", []byte(`{}`), nil, nil, nil, grantAll{}) require.Error(t, err) assert.Equal(t, awserrors.ErrorResourceNotFoundException, err.Error()) } @@ -92,7 +92,7 @@ func TestInvokeStreamRouter_SelfHostSuccess(t *testing.T) { defer ts.Close() modelID := "meta.llama3-2-1b-instruct-v1:0" - rt := NewInvokeStreamRouter(nil, NewStaticEndpointResolver(map[string]string{modelID: ts.URL})) + rt := NewInvokeStreamRouter(nil, NewStaticEndpointResolver(map[string]string{modelID: ts.URL}), grantAll{}) src, err := rt.InvokeModelWithResponseStream(context.Background(), "000000000001", modelID, []byte(`{"prompt":"hello"}`)) require.NoError(t, err) @@ -103,21 +103,21 @@ func TestInvokeStreamRouter_SelfHostSuccess(t *testing.T) { } func TestInvokeStreamRouter_UnknownModelReturnsResourceNotFound(t *testing.T) { - rt := NewInvokeStreamRouter(nil, nil) + rt := NewInvokeStreamRouter(nil, nil, grantAll{}) _, err := rt.InvokeModelWithResponseStream(context.Background(), "000000000001", "does.not-exist-v1:0", []byte(`{}`)) require.Error(t, err) assert.Equal(t, awserrors.ErrorResourceNotFoundException, err.Error()) } func TestInvokeStreamRouter_AnthropicNoCredentialReturnsAccessDenied(t *testing.T) { - rt := NewInvokeStreamRouter(stubCredentialResolver{ok: false}, nil) + rt := NewInvokeStreamRouter(stubCredentialResolver{ok: false}, nil, grantAll{}) _, err := rt.InvokeModelWithResponseStream(context.Background(), "000000000001", "anthropic.claude-3-5-sonnet-20240620-v1:0", []byte(`{}`)) require.Error(t, err) assert.Equal(t, awserrors.ErrorAccessDeniedException, err.Error()) } func TestInvokeStreamRouter_SelfHostNoEndpointReturnsModelNotReady(t *testing.T) { - rt := NewInvokeStreamRouter(nil, nil) + rt := NewInvokeStreamRouter(nil, nil, grantAll{}) _, err := rt.InvokeModelWithResponseStream(context.Background(), "000000000001", "meta.llama3-2-1b-instruct-v1:0", []byte(`{"prompt":"hello"}`)) require.Error(t, err) assert.Equal(t, awserrors.ErrorModelNotReadyException, err.Error()) diff --git a/spinifex/gateway/bedrock/provider.go b/spinifex/gateway/bedrock/provider.go index 07575959f..68dc8f2d5 100644 --- a/spinifex/gateway/bedrock/provider.go +++ b/spinifex/gateway/bedrock/provider.go @@ -34,12 +34,14 @@ type Router struct { resolver CredentialResolver endpointResolver EndpointResolver recorder Recorder + access AccessResolver } -// NewRouter constructs a Router. A nil resolver, endpointResolver, or -// recorder falls back to a no-op implementation, so a Router is always safe -// to use even before the real stores are wired in. -func NewRouter(resolver CredentialResolver, endpointResolver EndpointResolver, recorder Recorder) *Router { +// NewRouter constructs a Router. A nil resolver, endpointResolver, or recorder +// falls back to a no-op implementation, and a nil access falls back to denying +// every model, so a Router is always safe to use even before the real stores +// are wired in. +func NewRouter(resolver CredentialResolver, endpointResolver EndpointResolver, recorder Recorder, access AccessResolver) *Router { if resolver == nil { resolver = NoopCredentialResolver } @@ -49,14 +51,18 @@ func NewRouter(resolver CredentialResolver, endpointResolver EndpointResolver, r if recorder == nil { recorder = NoopRecorder } - return &Router{resolver: resolver, endpointResolver: endpointResolver, recorder: recorder} + if access == nil { + access = DenyAllAccessResolver + } + return &Router{resolver: resolver, endpointResolver: endpointResolver, recorder: recorder, access: access} } // Converse routes modelID to its provider via the catalog. Unknown modelIds -// and unresolvable vendors return ResourceNotFoundException; a vendor with no -// resolvable credential returns AccessDeniedException. Every exit records an -// InvocationRecord via the deferred closure, matching pumpConverseStream's -// treatment of the streaming path. +// and unresolvable vendors return ResourceNotFoundException; an ungranted +// model, or a vendor with no resolvable credential, returns +// AccessDeniedException. Every exit records an InvocationRecord via the +// deferred closure, matching pumpConverseStream's treatment of the streaming +// path. func (rt *Router) Converse(ctx context.Context, accountID, modelID string, input *bedrockruntime.ConverseInput) (out *bedrockruntime.ConverseOutput, err error) { requestID := uuid.NewString() start := time.Now() @@ -96,9 +102,10 @@ func (rt *Router) Converse(ctx context.Context, accountID, modelID string, input }) }() - entry, ok := lookupCatalogEntry(modelID) - if !ok { - err = errors.New(awserrors.ErrorResourceNotFoundException) + // Assigns the named err, so the deferred closure records a denied + // invocation the same as any other failed one. + entry, err := grantedCatalogEntry(ctx, accountID, modelID, rt.access) + if err != nil { return nil, err } backend = entry.Provider @@ -135,10 +142,10 @@ func (rt *Router) Converse(ctx context.Context, accountID, modelID string, input } // Converse is the bedrock-runtime Converse entry point used by the gateway -// route table. resolver, endpointResolver, and recorder may be nil; NewRouter -// supplies no-op fallbacks. -func Converse(ctx context.Context, accountID, modelID string, input *bedrockruntime.ConverseInput, resolver CredentialResolver, endpointResolver EndpointResolver, recorder Recorder) (*bedrockruntime.ConverseOutput, error) { - return NewRouter(resolver, endpointResolver, recorder).Converse(ctx, accountID, modelID, input) +// route table. resolver, endpointResolver, recorder and access may be nil; +// NewRouter supplies no-op (and, for access, deny-all) fallbacks. +func Converse(ctx context.Context, accountID, modelID string, input *bedrockruntime.ConverseInput, resolver CredentialResolver, endpointResolver EndpointResolver, recorder Recorder, access AccessResolver) (*bedrockruntime.ConverseOutput, error) { + return NewRouter(resolver, endpointResolver, recorder, access).Converse(ctx, accountID, modelID, input) } // ConverseStreamProvider is the optional streaming capability a Provider may @@ -166,9 +173,9 @@ func converseStreamToConverseInput(input *bedrockruntime.ConverseStreamInput) *b // Converse, then requires the resolved provider to also implement // ConverseStreamProvider. func (rt *Router) ConverseStream(ctx context.Context, accountID, modelID string, input *bedrockruntime.ConverseStreamInput) (converseStreamSource, error) { - entry, ok := lookupCatalogEntry(modelID) - if !ok { - return nil, errors.New(awserrors.ErrorResourceNotFoundException) + entry, err := grantedCatalogEntry(ctx, accountID, modelID, rt.access) + if err != nil { + return nil, err } var p Provider diff --git a/spinifex/gateway/bedrock/router_test.go b/spinifex/gateway/bedrock/router_test.go index f1009ab4b..18856dcff 100644 --- a/spinifex/gateway/bedrock/router_test.go +++ b/spinifex/gateway/bedrock/router_test.go @@ -46,7 +46,7 @@ func TestRouter_Converse_SelfHostSuccess(t *testing.T) { defer ts.Close() modelID := "meta.llama3-2-1b-instruct-v1:0" - rt := NewRouter(nil, NewStaticEndpointResolver(map[string]string{modelID: ts.URL}), nil) + rt := NewRouter(nil, NewStaticEndpointResolver(map[string]string{modelID: ts.URL}), nil, grantAll{}) out, err := rt.Converse(context.Background(), "000000000001", modelID, converseInput()) require.NoError(t, err) @@ -55,14 +55,14 @@ func TestRouter_Converse_SelfHostSuccess(t *testing.T) { } func TestRouter_Converse_UnknownModelReturnsResourceNotFound(t *testing.T) { - rt := NewRouter(nil, nil, nil) + rt := NewRouter(nil, nil, nil, grantAll{}) _, err := rt.Converse(context.Background(), "000000000001", "does.not-exist-v1:0", converseInput()) require.Error(t, err) assert.Equal(t, awserrors.ErrorResourceNotFoundException, err.Error()) } func TestRouter_Converse_AnthropicNoCredentialReturnsAccessDenied(t *testing.T) { - rt := NewRouter(stubCredentialResolver{ok: false}, nil, nil) + rt := NewRouter(stubCredentialResolver{ok: false}, nil, nil, grantAll{}) _, err := rt.Converse(context.Background(), "000000000001", "anthropic.claude-3-5-sonnet-20240620-v1:0", converseInput()) require.Error(t, err) assert.Equal(t, awserrors.ErrorAccessDeniedException, err.Error()) @@ -80,14 +80,14 @@ func TestConverse_PackageEntryPoint_SelfHostSuccess(t *testing.T) { defer ts.Close() modelID := "meta.llama3-2-1b-instruct-v1:0" - out, err := Converse(context.Background(), "000000000001", modelID, converseInput(), nil, NewStaticEndpointResolver(map[string]string{modelID: ts.URL}), nil) + out, err := Converse(context.Background(), "000000000001", modelID, converseInput(), nil, NewStaticEndpointResolver(map[string]string{modelID: ts.URL}), nil, grantAll{}) require.NoError(t, err) require.NotNil(t, out.Output.Message) assert.Equal(t, "via package Converse", *out.Output.Message.Content[0].Text) } func TestConverse_PackageEntryPoint_UnknownModel(t *testing.T) { - _, err := Converse(context.Background(), "000000000001", "does.not-exist-v1:0", converseInput(), nil, nil, nil) + _, err := Converse(context.Background(), "000000000001", "does.not-exist-v1:0", converseInput(), nil, nil, nil, grantAll{}) require.Error(t, err) assert.Equal(t, awserrors.ErrorResourceNotFoundException, err.Error()) } @@ -109,7 +109,7 @@ func TestRouter_ConverseStream_SelfHostSuccess(t *testing.T) { defer ts.Close() modelID := "meta.llama3-2-1b-instruct-v1:0" - rt := NewRouter(nil, NewStaticEndpointResolver(map[string]string{modelID: ts.URL}), nil) + rt := NewRouter(nil, NewStaticEndpointResolver(map[string]string{modelID: ts.URL}), nil, grantAll{}) src, err := rt.ConverseStream(context.Background(), "000000000001", modelID, converseStreamInput()) require.NoError(t, err) @@ -120,28 +120,28 @@ func TestRouter_ConverseStream_SelfHostSuccess(t *testing.T) { } func TestRouter_ConverseStream_UnknownModelReturnsResourceNotFound(t *testing.T) { - rt := NewRouter(nil, nil, nil) + rt := NewRouter(nil, nil, nil, grantAll{}) _, err := rt.ConverseStream(context.Background(), "000000000001", "does.not-exist-v1:0", converseStreamInput()) require.Error(t, err) assert.Equal(t, awserrors.ErrorResourceNotFoundException, err.Error()) } func TestRouter_ConverseStream_AnthropicNoCredentialReturnsAccessDenied(t *testing.T) { - rt := NewRouter(stubCredentialResolver{ok: false}, nil, nil) + rt := NewRouter(stubCredentialResolver{ok: false}, nil, nil, grantAll{}) _, err := rt.ConverseStream(context.Background(), "000000000001", "anthropic.claude-3-5-sonnet-20240620-v1:0", converseStreamInput()) require.Error(t, err) assert.Equal(t, awserrors.ErrorAccessDeniedException, err.Error()) } func TestRouter_ConverseStream_SelfHostNoEndpointReturnsModelNotReady(t *testing.T) { - rt := NewRouter(nil, nil, nil) + rt := NewRouter(nil, nil, nil, grantAll{}) _, err := rt.ConverseStream(context.Background(), "000000000001", "meta.llama3-2-1b-instruct-v1:0", converseStreamInput()) require.Error(t, err) assert.Equal(t, awserrors.ErrorModelNotReadyException, err.Error()) } func TestNewRouter_NilArgumentsFallBackToNoops(t *testing.T) { - rt := NewRouter(nil, nil, nil) + rt := NewRouter(nil, nil, nil, grantAll{}) require.NotNil(t, rt) // Self-host model with no endpoint resolver configured resolves nothing, diff --git a/spinifex/gateway/bedrock_request_test.go b/spinifex/gateway/bedrock_request_test.go index 7663b11e4..6352f6272 100644 --- a/spinifex/gateway/bedrock_request_test.go +++ b/spinifex/gateway/bedrock_request_test.go @@ -65,10 +65,19 @@ func (bedrockRequestWeightsStub) Resolve(_ context.Context, modelID string) (str return "snap-stub", true, nil } +// grantAllModels is an AccessResolver that grants every model, standing in for +// a provisioned grant store in tests whose subject is routing rather than +// access control. +type grantAllModels struct{} + +func (grantAllModels) Granted(_ context.Context, _, _ string) (bool, error) { return true, nil } + // newBedrockRequestGateway builds a GatewayConfig with a real NATS connection // (satisfying Bedrock_Request/BedrockRuntime_Request's nil-check only — no // NATS subject handling is required for these routes), no IAMService (so -// checkPolicy is a no-op), and BedrockEndpoints pinned at the given vLLM stub. +// checkPolicy is a no-op), BedrockEndpoints pinned at the given vLLM stub, and +// every model granted — access is deny-by-default, so without a resolver these +// routes would all return before reaching the code under test. // ListFoundationModels/GetFoundationModel read the weights resolver through a // package-level setter rather than a GatewayConfig field, so it is installed // here and reset on cleanup. @@ -83,6 +92,7 @@ func newBedrockRequestGateway(t *testing.T, vllmURL string) *GatewayConfig { BedrockEndpoints: map[string]string{ bedrockTestLlamaModelID: vllmURL, }, + BedrockAccess: grantAllModels{}, } } @@ -92,6 +102,54 @@ func bedrockRequestWithAccount(method, path, body string) *http.Request { return req.WithContext(ctx) } +// TestBedrockRequest_DenyByDefault covers the whole gateway path with no +// access resolver configured, which is what a deployment that has granted +// nothing looks like: the catalog is empty and every runtime route refuses. +// This is the end-to-end guard on deny-by-default, one layer above the +// per-function access tests in gateway/bedrock. +func TestBedrockRequest_DenyByDefault(t *testing.T) { + ts := newVLLMStub(t) + gw := newBedrockRequestGateway(t, ts.URL) + gw.BedrockAccess = nil // no grant store configured + + t.Run("list is empty", func(t *testing.T) { + req := bedrockRequestWithAccount(http.MethodGet, "/foundation-models", "") + w := httptest.NewRecorder() + require.NoError(t, gw.Bedrock_Request(w, req)) + require.Equal(t, http.StatusOK, w.Code) + + var out bedrock.ListFoundationModelsOutput + require.NoError(t, json.Unmarshal(w.Body.Bytes(), &out)) + assert.Empty(t, out.ModelSummaries) + }) + + t.Run("describe reports not found", func(t *testing.T) { + req := bedrockRequestWithAccount(http.MethodGet, "/foundation-models/"+bedrockTestLlamaModelID, "") + w := httptest.NewRecorder() + err := gw.Bedrock_Request(w, req) + require.Error(t, err) + assert.Equal(t, awserrors.ErrorResourceNotFoundException, err.Error()) + }) + + // Every runtime verb must refuse, not just the one that happens to be + // checked first — guessing a modelId must gain nothing on any path. + runtimeCases := []struct{ name, path, body string }{ + {"converse", "/model/" + bedrockTestLlamaModelID + "/converse", `{"messages":[]}`}, + {"invoke", "/model/" + bedrockTestLlamaModelID + "/invoke", `{"prompt":"hi"}`}, + {"converse-stream", "/model/" + bedrockTestLlamaModelID + "/converse-stream", `{"messages":[]}`}, + {"invoke-with-response-stream", "/model/" + bedrockTestLlamaModelID + "/invoke-with-response-stream", `{"prompt":"hi"}`}, + } + for _, tc := range runtimeCases { + t.Run(tc.name, func(t *testing.T) { + req := bedrockRequestWithAccount(http.MethodPost, tc.path, tc.body) + w := httptest.NewRecorder() + err := gw.BedrockRuntime_Request(w, req) + require.Error(t, err) + assert.Equal(t, awserrors.ErrorAccessDeniedException, err.Error()) + }) + } +} + func TestBedrockRequest_ListFoundationModels(t *testing.T) { ts := newVLLMStub(t) gw := newBedrockRequestGateway(t, ts.URL) diff --git a/spinifex/gateway/bedrockruntime.go b/spinifex/gateway/bedrockruntime.go index 18af32314..4cf084617 100644 --- a/spinifex/gateway/bedrockruntime.go +++ b/spinifex/gateway/bedrockruntime.go @@ -27,22 +27,23 @@ type bedrockRuntimeRoute struct { // gateway function. params holds the regex capture groups, PathUnescape'd. // resolver is gw.bedrockResolver() (credential store or no-op); endpoints is // gw.bedrockEndpointResolver() over the configured pinned self-host -// endpoints; recorder is gw.bedrockRecorder() (invocation recorder or no-op). -type bedrockRuntimeRouteHandler func(ctx context.Context, accountID string, params []string, body []byte, resolver gateway_bedrock.CredentialResolver, endpoints gateway_bedrock.EndpointResolver, recorder gateway_bedrock.Recorder) (any, error) +// endpoints; recorder is gw.bedrockRecorder() (invocation recorder or no-op); +// access is gw.bedrockAccessResolver() (grant store or deny-all). +type bedrockRuntimeRouteHandler func(ctx context.Context, accountID string, params []string, body []byte, resolver gateway_bedrock.CredentialResolver, endpoints gateway_bedrock.EndpointResolver, recorder gateway_bedrock.Recorder, access gateway_bedrock.AccessResolver) (any, error) // bedrockRuntimeRoutes is the dispatch table. InvokeModel has no handler // function here: BedrockRuntime_Request special-cases its action to bypass // the JSON-marshaling dispatch below, since its response is raw bytes. var bedrockRuntimeRoutes = []bedrockRuntimeRoute{ {"POST", regexp.MustCompile(`^/model/([^/]+)/converse$`), "Converse", - func(ctx context.Context, acct string, p []string, b []byte, resolver gateway_bedrock.CredentialResolver, endpoints gateway_bedrock.EndpointResolver, recorder gateway_bedrock.Recorder) (any, error) { + func(ctx context.Context, acct string, p []string, b []byte, resolver gateway_bedrock.CredentialResolver, endpoints gateway_bedrock.EndpointResolver, recorder gateway_bedrock.Recorder, access gateway_bedrock.AccessResolver) (any, error) { input := new(bedrockruntime.ConverseInput) if len(b) > 0 { if err := json.Unmarshal(b, input); err != nil { return nil, errors.New(awserrors.ErrorValidationException) } } - return gateway_bedrock.Converse(ctx, acct, p[0], input, resolver, endpoints, recorder) + return gateway_bedrock.Converse(ctx, acct, p[0], input, resolver, endpoints, recorder, access) }}, {"POST", regexp.MustCompile(`^/model/([^/]+)/invoke$`), "InvokeModel", nil}, {"POST", regexp.MustCompile(`^/model/([^/]+)/converse-stream$`), "ConverseStream", nil}, @@ -123,7 +124,7 @@ func (gw *GatewayConfig) BedrockRuntime_Request(w http.ResponseWriter, r *http.R // InvokeModel returns provider-native bytes, not a struct WriteJSONResponse // could marshal, so it writes its own response body directly. if action == "InvokeModel" { - respBody, contentType, err := gateway_bedrock.InvokeModel(r.Context(), accountID, params[0], body, gw.bedrockResolver(), gw.bedrockEndpointResolver(), gw.bedrockRecorder()) + respBody, contentType, err := gateway_bedrock.InvokeModel(r.Context(), accountID, params[0], body, gw.bedrockResolver(), gw.bedrockEndpointResolver(), gw.bedrockRecorder(), gw.bedrockAccessResolver()) if err != nil { return err } @@ -138,13 +139,13 @@ func (gw *GatewayConfig) BedrockRuntime_Request(w http.ResponseWriter, r *http.R // (-> ErrorHandler); once streaming starts they always return nil, // surfacing any further failure as an in-band exception event. if action == "ConverseStream" { - return gateway_bedrock.ConverseStream(r.Context(), w, accountID, params[0], body, gw.bedrockResolver(), gw.bedrockEndpointResolver(), gw.bedrockRecorder()) + return gateway_bedrock.ConverseStream(r.Context(), w, accountID, params[0], body, gw.bedrockResolver(), gw.bedrockEndpointResolver(), gw.bedrockRecorder(), gw.bedrockAccessResolver()) } if action == "InvokeModelWithResponseStream" { - return gateway_bedrock.InvokeModelWithResponseStream(r.Context(), w, accountID, params[0], body, gw.bedrockResolver(), gw.bedrockEndpointResolver(), r.Header.Get("Content-Type"), gw.bedrockRecorder()) + return gateway_bedrock.InvokeModelWithResponseStream(r.Context(), w, accountID, params[0], body, gw.bedrockResolver(), gw.bedrockEndpointResolver(), r.Header.Get("Content-Type"), gw.bedrockRecorder(), gw.bedrockAccessResolver()) } - output, err := handler(r.Context(), accountID, params, body, gw.bedrockResolver(), gw.bedrockEndpointResolver(), gw.bedrockRecorder()) + output, err := handler(r.Context(), accountID, params, body, gw.bedrockResolver(), gw.bedrockEndpointResolver(), gw.bedrockRecorder(), gw.bedrockAccessResolver()) if err != nil { return err } diff --git a/spinifex/gateway/gateway.go b/spinifex/gateway/gateway.go index e2bb6d6c5..a06182691 100644 --- a/spinifex/gateway/gateway.go +++ b/spinifex/gateway/gateway.go @@ -132,6 +132,15 @@ type GatewayConfig struct { // to a no-op recorder, so routes stay safe to exercise before the // invocation stream is wired in (e.g. unit tests of unrelated routes). BedrockRecorder gateway_bedrock.Recorder + // BedrockAccess resolves per-account model-access grants. Access is + // deny-by-default, so nil means no account may list or invoke any model. + // It is the AccessResolver interface rather than the concrete store so a + // test can inject a fixed grant set without standing up JetStream. + BedrockAccess gateway_bedrock.AccessResolver + // BedrockAccessAdmin is the writable side of the same grant store, used by + // the spinifex admin actions. Nil disables grant administration, which is + // how a gateway with an injected read-only resolver behaves. + BedrockAccessAdmin *gateway_bedrock.ModelAccessStore } var supportedServices = map[string]bool{ diff --git a/spinifex/gateway/spinifex.go b/spinifex/gateway/spinifex.go index f386c551c..d184d6863 100644 --- a/spinifex/gateway/spinifex.go +++ b/spinifex/gateway/spinifex.go @@ -8,16 +8,20 @@ import ( "github.com/mulgadc/spinifex/spinifex/admin" "github.com/mulgadc/spinifex/spinifex/awserrors" + gateway_bedrock "github.com/mulgadc/spinifex/spinifex/gateway/bedrock" gateway_spx "github.com/mulgadc/spinifex/spinifex/gateway/spx" ) // spinifexAdminActions lists actions that require admin account access. var spinifexAdminActions = map[string]bool{ - "GetVersion": true, - "GetNodes": true, - "GetVMs": true, - "GetStorageStatus": true, - "PromoteImage": true, + "GetVersion": true, + "GetNodes": true, + "GetVMs": true, + "GetStorageStatus": true, + "PromoteImage": true, + "GrantModelAccess": true, + "RevokeModelAccess": true, + "ListModelAccess": true, } func (gw *GatewayConfig) Spinifex_Request(w http.ResponseWriter, r *http.Request) error { @@ -79,6 +83,34 @@ func (gw *GatewayConfig) Spinifex_Request(w http.ResponseWriter, r *http.Request return errors.New(awserrors.ErrorMissingParameter) } output, err = gateway_spx.PromoteImage(ctx, gw.NATSConn, imageID, accountID) + case "GrantModelAccess", "RevokeModelAccess": + // Grants are gateway-owned KV state, like the bedrock credential store, + // so these need no NATS hop: the gateway writes the bucket it reads. + if gw.BedrockAccessAdmin == nil { + return errors.New(awserrors.ErrorServerInternal) + } + targetAccount, modelID := queryArgs["AccountId"], queryArgs["ModelId"] + if targetAccount == "" || modelID == "" { + return errors.New(awserrors.ErrorMissingParameter) + } + if action == "GrantModelAccess" { + err = gw.BedrockAccessAdmin.Grant(ctx, targetAccount, modelID) + } else { + err = gw.BedrockAccessAdmin.Revoke(ctx, targetAccount, modelID) + } + output = gateway_bedrock.ModelAccessChange{AccountID: targetAccount, ModelID: modelID} + case "ListModelAccess": + if gw.BedrockAccessAdmin == nil { + return errors.New(awserrors.ErrorServerInternal) + } + targetAccount := queryArgs["AccountId"] + if targetAccount == "" { + return errors.New(awserrors.ErrorMissingParameter) + } + var models []string + if models, err = gw.BedrockAccessAdmin.List(ctx, targetAccount); err == nil { + output = gateway_bedrock.ModelAccessList{AccountID: targetAccount, ModelIDs: models} + } default: return errors.New(awserrors.ErrorInvalidAction) } diff --git a/spinifex/gateway/spinifex_modelaccess_test.go b/spinifex/gateway/spinifex_modelaccess_test.go new file mode 100644 index 000000000..053a22219 --- /dev/null +++ b/spinifex/gateway/spinifex_modelaccess_test.go @@ -0,0 +1,128 @@ +package gateway + +import ( + "context" + "net/http" + "net/http/httptest" + "net/url" + "strings" + "testing" + + "github.com/mulgadc/spinifex/spinifex/admin" + "github.com/mulgadc/spinifex/spinifex/awserrors" + gateway_bedrock "github.com/mulgadc/spinifex/spinifex/gateway/bedrock" + "github.com/mulgadc/spinifex/spinifex/testutil" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +const accessTestModelID = "meta.llama3-2-1b-instruct-v1:0" + +// spinifexAccessRequest drives Spinifex_Request with the query args the model +// access actions read, which the Action-only helper cannot carry. +func spinifexAccessRequest(t *testing.T, gw *GatewayConfig, action, accountID string, args url.Values) *httptest.ResponseRecorder { + t.Helper() + args.Set("Action", action) + req := httptest.NewRequest(http.MethodPost, "/", strings.NewReader(args.Encode())) + req.Header.Set("Content-Type", "application/x-www-form-urlencoded") + ctx := req.Context() + ctx = context.WithValue(ctx, ctxAccountID, accountID) + ctx = context.WithValue(ctx, ctxIdentity, "admin") + ctx = context.WithValue(ctx, ctxService, "spinifex") + req = req.WithContext(ctx) + + w := httptest.NewRecorder() + if err := gw.Spinifex_Request(w, req); err != nil { + w.WriteHeader(http.StatusForbidden) + w.WriteString(err.Error()) + } + return w +} + +// accessGateway builds a gateway whose grant store is backed by an embedded +// JetStream server, so the admin actions are exercised against the same store +// the runtime resolves grants from. +func accessGateway(t *testing.T) *GatewayConfig { + t.Helper() + _, _, js := testutil.StartTestJetStream(t) + return &GatewayConfig{ + DisableLogging: true, + BedrockAccessAdmin: gateway_bedrock.NewModelAccessStore(js, 1), + } +} + +// TestSpinifex_ModelAccess_GrantListRevoke walks the operator's whole loop: +// an account starts with nothing, a grant shows up in the listing, and a +// revoke takes it away again. +func TestSpinifex_ModelAccess_GrantListRevoke(t *testing.T) { + gw := accessGateway(t) + adminAccount := admin.DefaultAccountID() + + w := spinifexAccessRequest(t, gw, "ListModelAccess", adminAccount, url.Values{"AccountId": {"000000000002"}}) + require.Equal(t, http.StatusOK, w.Code) + assert.NotContains(t, w.Body.String(), accessTestModelID) + + w = spinifexAccessRequest(t, gw, "GrantModelAccess", adminAccount, + url.Values{"AccountId": {"000000000002"}, "ModelId": {accessTestModelID}}) + require.Equal(t, http.StatusOK, w.Code) + assert.Contains(t, w.Body.String(), accessTestModelID) + + w = spinifexAccessRequest(t, gw, "ListModelAccess", adminAccount, url.Values{"AccountId": {"000000000002"}}) + require.Equal(t, http.StatusOK, w.Code) + assert.Contains(t, w.Body.String(), accessTestModelID) + + w = spinifexAccessRequest(t, gw, "RevokeModelAccess", adminAccount, + url.Values{"AccountId": {"000000000002"}, "ModelId": {accessTestModelID}}) + require.Equal(t, http.StatusOK, w.Code) + + w = spinifexAccessRequest(t, gw, "ListModelAccess", adminAccount, url.Values{"AccountId": {"000000000002"}}) + require.Equal(t, http.StatusOK, w.Code) + assert.NotContains(t, w.Body.String(), accessTestModelID) +} + +// TestSpinifex_ModelAccess_NonAdminDenied is the important one: an account +// that could grant itself models would make the whole access model decorative. +func TestSpinifex_ModelAccess_NonAdminDenied(t *testing.T) { + gw := accessGateway(t) + + for _, action := range []string{"GrantModelAccess", "RevokeModelAccess", "ListModelAccess"} { + w := spinifexAccessRequest(t, gw, action, "000000000002", + url.Values{"AccountId": {"000000000002"}, "ModelId": {accessTestModelID}}) + require.Equal(t, http.StatusForbidden, w.Code, "action %s", action) + assert.Contains(t, w.Body.String(), awserrors.ErrorAccessDenied, "action %s", action) + } +} + +func TestSpinifex_ModelAccess_MissingParameters(t *testing.T) { + gw := accessGateway(t) + adminAccount := admin.DefaultAccountID() + + cases := []struct { + action string + args url.Values + }{ + {"GrantModelAccess", url.Values{"ModelId": {accessTestModelID}}}, + {"GrantModelAccess", url.Values{"AccountId": {"000000000002"}}}, + {"RevokeModelAccess", url.Values{"ModelId": {accessTestModelID}}}, + {"ListModelAccess", url.Values{}}, + } + for _, tc := range cases { + w := spinifexAccessRequest(t, gw, tc.action, adminAccount, tc.args) + require.Equal(t, http.StatusForbidden, w.Code, "action %s", tc.action) + assert.Contains(t, w.Body.String(), awserrors.ErrorMissingParameter, "action %s", tc.action) + } +} + +// TestSpinifex_ModelAccess_NoStoreIsServerError covers a gateway built without +// a grant store: the admin actions refuse rather than panicking on nil. +func TestSpinifex_ModelAccess_NoStoreIsServerError(t *testing.T) { + gw := &GatewayConfig{DisableLogging: true} + adminAccount := admin.DefaultAccountID() + + for _, action := range []string{"GrantModelAccess", "RevokeModelAccess", "ListModelAccess"} { + w := spinifexAccessRequest(t, gw, action, adminAccount, + url.Values{"AccountId": {"000000000002"}, "ModelId": {accessTestModelID}}) + require.Equal(t, http.StatusForbidden, w.Code, "action %s", action) + assert.Contains(t, w.Body.String(), awserrors.ErrorServerInternal, "action %s", action) + } +} diff --git a/spinifex/services/awsgw/awsgw.go b/spinifex/services/awsgw/awsgw.go index 9b60210db..85f278951 100644 --- a/spinifex/services/awsgw/awsgw.go +++ b/spinifex/services/awsgw/awsgw.go @@ -308,6 +308,25 @@ func launchService(config *config.ClusterConfig) error { // are called from gateway/bedrock.go's fixed-arity route table. gateway_bedrock.SetWeightsResolver(gateway_bedrock.NewWeightsStore(js, len(config.Nodes))) + // Bedrock model access: grants live in the bedrock-model-access KV bucket + // and are deny-by-default, so a fresh deployment serves no models until an + // operator grants them (spx admin ochre access grant). + bedrockAccess := gateway_bedrock.NewModelAccessStore(js, len(config.Nodes)) + + // Deny-by-default would otherwise leave a fresh install with a catalog + // nobody can see, so seed the platform admin account — the operator's own + // account, created by spx admin init — with the full catalog on first + // start. Tenant accounts are unaffected and still begin with no access. + // Best-effort: the gateway must serve even if this fails, and because the + // marker is written only on success, the next start retries. + if seeded, err := bedrockAccess.SeedAccountGrants(context.Background(), admin.DefaultAccountID(), gateway_bedrock.CatalogModelIDs()); err != nil { + slog.Warn("Bedrock model access: seeding admin grants failed, will retry on next start", + "accountID", admin.DefaultAccountID(), "err", err) + } else if seeded { + slog.Info("Bedrock model access: seeded admin account with the model catalog", + "accountID", admin.DefaultAccountID(), "models", len(gateway_bedrock.CatalogModelIDs())) + } + // Bedrock self-host endpoints are pinned, so their // OpenAI-compatible base URLs come from static config. OCHRE_VLLM_ENDPOINTS // is a comma-separated list of modelId=baseURL pairs. @@ -364,6 +383,8 @@ func launchService(config *config.ClusterConfig) error { BedrockEndpoints: bedrockEndpoints, BedrockLoggingConfig: bedrockLoggingConfig, BedrockRecorder: bedrockRecorder, + BedrockAccess: bedrockAccess, + BedrockAccessAdmin: bedrockAccess, } // Rotate the ECR signing key on a 30-day cadence, retaining the previous keys diff --git a/tests/e2e/bedrock/main_test.go b/tests/e2e/bedrock/main_test.go index d92bc2a8d..a91c57d9b 100644 --- a/tests/e2e/bedrock/main_test.go +++ b/tests/e2e/bedrock/main_test.go @@ -6,6 +6,13 @@ // ListFoundationModels/GetFoundationModel against any running cluster and // gates real Converse calls on backend availability (self-host GPU or a live // Anthropic key), so CI parity holds without requiring either to be present. +// +// Model access is deny-by-default, so the cluster's test account must hold +// grants or every assertion here reduces to "empty catalog, access denied". +// awsgw seeds the platform admin account with the full catalog on first +// start, so a suite running as that account needs no setup; running as a +// tenant account needs +// `spx admin ochre access grant --account-id --all-models` first. package bedrock import ( diff --git a/tests/integration/harness.go b/tests/integration/harness.go index 8919b53b6..bd7f4f8e5 100644 --- a/tests/integration/harness.go +++ b/tests/integration/harness.go @@ -33,6 +33,7 @@ import ( "github.com/aws/aws-sdk-go/service/iam" "github.com/aws/aws-sdk-go/service/sts" "github.com/mulgadc/spinifex/spinifex/gateway" + gateway_bedrock "github.com/mulgadc/spinifex/spinifex/gateway/bedrock" gateway_ecr "github.com/mulgadc/spinifex/spinifex/gateway/ecr" gateway_ecrauth "github.com/mulgadc/spinifex/spinifex/gateway/ecrauth" handlers_ecr "github.com/mulgadc/spinifex/spinifex/handlers/ecr" @@ -157,6 +158,8 @@ func startGateway(t *testing.T, collector *conformanceCollector, opts ...Option) signingKey, verifyKeys, err := gateway_ecrauth.LoadOrCreateSigningKey(t.Context(), js, masterKey, 1) require.NoError(t, err) + bedrockAccess := gateway_bedrock.NewModelAccessStore(js, 1) + cfg := &gateway.GatewayConfig{ DisableLogging: true, NATSConn: nc, @@ -173,6 +176,13 @@ func startGateway(t *testing.T, collector *conformanceCollector, opts ...Option) ECRRegistry: gateway_ecr.NewRegistry(objectstore.NewMemoryObjectStore(), handlers_ecr.NewNATSMetaStore(nc), utils.GlobalAccountID), ECRTokenIssuer: gateway_ecrauth.NewIssuer(signingKey, testECRAudience), ECRTokenVerifier: gateway_ecrauth.NewVerifier(verifyKeys, testECRAudience), + // Ochre model access is deny-by-default, so without a grant store every + // bedrock route refuses. Tests sign as the system account, which the + // store exempts from grants exactly as it does in production, so the + // catalog is visible here without provisioning grants per test; a test + // signing as a tenant account would have to grant first. + BedrockAccess: bedrockAccess, + BedrockAccessAdmin: bedrockAccess, } for _, opt := range opts { opt(cfg)