diff --git a/internal/auth/store_factory.go b/internal/auth/store_factory.go index e9535709..addd21d4 100644 --- a/internal/auth/store_factory.go +++ b/internal/auth/store_factory.go @@ -2,6 +2,8 @@ package auth import ( "context" + "fmt" + "strings" "github.com/jackc/pgx/v5/pgxpool" "chat-agent-system/internal/config" @@ -13,7 +15,17 @@ type Stores struct { Sessions SessionStore } +func allowInMemoryAuthFallback(cfg *config.Config) bool { + if cfg == nil || !cfg.AllowInMemory { + return false + } + env := strings.ToLower(strings.TrimSpace(cfg.Env)) + return env == "dev" || env == "test" +} + // NewStores 根据数据库与 IAM 迁移状态选择 PG 或内存实现。 +// 仅当环境明确为 dev/test 且 AllowInMemory 时允许 MemoryStore; +// 环境空值、未知值、main/prod、pool 缺失或 IAM 表未就绪时 fail closed。 func NewStores(ctx context.Context, pool *pgxpool.Pool, cfg *config.Config) (*Stores, error) { if pool != nil && IAMTablesReady(ctx, pool) { pg, err := NewPgStore(pool, cfg) @@ -25,6 +37,19 @@ func NewStores(ctx context.Context, pool *pgxpool.Pool, cfg *config.Config) (*St } return &Stores{Auth: pg, Sessions: pg}, nil } - mem := NewMemoryStore(cfg) - return &Stores{Auth: mem, Sessions: NewMemorySessionStore()}, nil + if allowInMemoryAuthFallback(cfg) { + mem := NewMemoryStore(cfg) + return &Stores{Auth: mem, Sessions: NewMemorySessionStore()}, nil + } + env := "" + if cfg != nil { + env = strings.TrimSpace(cfg.Env) + } + reason := "IAM tables not ready or database pool unavailable" + if pool == nil { + reason = "database pool is nil" + } else if !IAMTablesReady(ctx, pool) { + reason = "IAM tables not ready" + } + return nil, fmt.Errorf("auth store fail-closed: %s (env=%q); memory IAM fallback requires env=dev|test and ALLOW_IN_MEMORY=true", reason, env) } diff --git a/internal/auth/store_factory_test.go b/internal/auth/store_factory_test.go new file mode 100644 index 00000000..59236546 --- /dev/null +++ b/internal/auth/store_factory_test.go @@ -0,0 +1,62 @@ +package auth + +import ( + "context" + "strings" + "testing" + + "chat-agent-system/internal/config" +) + +func TestNewStores_FailClosed_WhenIAMUnavailable(t *testing.T) { + ctx := context.Background() + cases := []struct { + name string + cfg *config.Config + }{ + {name: "main_nil_pool", cfg: &config.Config{Env: "main", AllowInMemory: false}}, + {name: "prod_nil_pool", cfg: &config.Config{Env: "prod", AllowInMemory: false}}, + {name: "empty_env_nil_pool", cfg: &config.Config{Env: "", AllowInMemory: true}}, + {name: "unknown_env_nil_pool", cfg: &config.Config{Env: "staging", AllowInMemory: true}}, + {name: "dev_without_opt_in", cfg: &config.Config{Env: "dev", AllowInMemory: false}}, + {name: "test_without_opt_in", cfg: &config.Config{Env: "test", AllowInMemory: false}}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + stores, err := NewStores(ctx, nil, tc.cfg) + if err == nil { + t.Fatalf("expected fail-closed error, got stores=%v", stores) + } + if stores != nil { + t.Fatalf("expected nil stores on failure, got %#v", stores) + } + if !strings.Contains(err.Error(), "IAM") && !strings.Contains(err.Error(), "memory") { + t.Fatalf("error should mention IAM/memory policy: %v", err) + } + }) + } +} + +func TestNewStores_AllowsMemoryOnlyWithDevTestOptIn(t *testing.T) { + ctx := context.Background() + for _, env := range []string{"dev", "test"} { + t.Run(env, func(t *testing.T) { + cfg := &config.Config{ + Env: env, + AllowInMemory: true, + AdminUsername: "admin", + AdminPassword: "admin-test-pass", + } + stores, err := NewStores(ctx, nil, cfg) + if err != nil { + t.Fatalf("expected memory fixture: %v", err) + } + if stores == nil || stores.Auth == nil || stores.Sessions == nil { + t.Fatal("expected auth+session stores") + } + if _, ok := stores.Auth.(*MemoryStore); !ok { + t.Fatalf("Auth should be MemoryStore, got %T", stores.Auth) + } + }) + } +}