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)