diff --git a/cmd/hold/main.go b/cmd/hold/main.go index f9c6420..ae20c80 100644 --- a/cmd/hold/main.go +++ b/cmd/hold/main.go @@ -13,6 +13,7 @@ import ( "atcr.io/pkg/hold" "atcr.io/pkg/hold/oci" "atcr.io/pkg/hold/pds" + "atcr.io/pkg/hold/quota" "atcr.io/pkg/logging" "atcr.io/pkg/s3" @@ -98,6 +99,18 @@ func main() { os.Exit(1) } + // Initialize quota manager from quotas.yaml + quotaMgr, err := quota.NewManager("./quotas.yaml") + if err != nil { + slog.Error("Failed to load quota config", "error", err) + os.Exit(1) + } + if quotaMgr.IsEnabled() { + slog.Info("Quota enforcement enabled", "berths", quotaMgr.BerthCount(), "defaultBerth", quotaMgr.GetDefaultBerth()) + } else { + slog.Info("Quota enforcement disabled (no quotas.yaml found)") + } + // Create blob store adapter and XRPC handlers var ociHandler *oci.XRPCHandler if holdPDS != nil { @@ -116,10 +129,10 @@ func main() { } // Create PDS XRPC handler (ATProto endpoints) - xrpcHandler = pds.NewXRPCHandler(holdPDS, *s3Service, driver, broadcaster, nil) + xrpcHandler = pds.NewXRPCHandler(holdPDS, *s3Service, driver, broadcaster, nil, quotaMgr) // Create OCI XRPC handler (multipart upload endpoints) - ociHandler = oci.NewXRPCHandler(holdPDS, *s3Service, driver, cfg.Server.DisablePresignedURLs, cfg.Registration.EnableBlueskyPosts, nil) + ociHandler = oci.NewXRPCHandler(holdPDS, *s3Service, driver, cfg.Server.DisablePresignedURLs, cfg.Registration.EnableBlueskyPosts, nil, quotaMgr) } // Setup HTTP routes with chi router diff --git a/deploy/docker-compose.prod.yml b/deploy/docker-compose.prod.yml index 21d6e5a..f1d6032 100644 --- a/deploy/docker-compose.prod.yml +++ b/deploy/docker-compose.prod.yml @@ -123,6 +123,7 @@ services: volumes: # PDS data (carstore SQLite + signing keys) - atcr-hold-data:/var/lib/atcr-hold + - ./quotas.yaml:/quotas.yaml:ro networks: - atcr-network healthcheck: diff --git a/deploy/quotas.yaml b/deploy/quotas.yaml new file mode 100644 index 0000000..d92bef4 --- /dev/null +++ b/deploy/quotas.yaml @@ -0,0 +1,35 @@ +# ATCR Hold Service Quota Configuration +# Copy this file to quotas.yaml to enable quota enforcement. +# If quotas.yaml doesn't exist, quotas are disabled (unlimited for all users). + +# Berths define quota tiers using nautical crew ranks. +# Each berth has a quota limit specified in human-readable format. +# Supported units: B, KB, MB, GB, TB, PB (case-insensitive) +berths: + # Entry-level crew - suitable for new or casual users + deckhand: + quota: 5GB + + # Mid-level crew - for regular contributors + bosun: + quota: 50GB + + # Senior crew - for power users or trusted contributors + quartermaster: + quota: 100GB + + # You can add custom berths with any name: + # unlimited_crew: + # quota: 1TB + +defaults: + # Default berth assigned to new crew members who don't have an explicit berth. + # This berth must exist in the berths section above. + new_crew_berth: deckhand + +# Notes: +# - The hold captain (owner) always has unlimited quota regardless of berths. +# - Crew members can be assigned a specific berth in their crew record. +# - If a crew member's berth doesn't exist in config, they fall back to the default. +# - Quota is calculated per-user by summing unique blob sizes (deduplicated). +# - Quota is checked when pushing manifests (after blobs are already uploaded). diff --git a/docs/QUOTAS.md b/docs/QUOTAS.md index 0e1bdc1..232e71d 100644 --- a/docs/QUOTAS.md +++ b/docs/QUOTAS.md @@ -507,13 +507,48 @@ GET /xrpc/io.atcr.hold.getQuotaBreakdown - Storage by repository - Email/webhook notifications - Grace period before hard enforcement -### 3. Tiered Quotas +### 3. Berth-Based Quotas (Implemented) -| Tier | Limit | -|------|-------| -| Free | 10 GB | -| Pro | 100 GB | -| Enterprise | Unlimited | +ATCR uses nautical-themed "berths" for quota tiers, configured via `quotas.yaml`: + +```yaml +# quotas.yaml +berths: + deckhand: # Entry-level crew + quota: 5GB + bosun: # Mid-level crew + quota: 50GB + quartermaster: # High-level crew + quota: 100GB + +defaults: + new_crew_berth: deckhand # Default berth for new crew members +``` + +| Berth | Limit | Description | +|-------|-------|-------------| +| deckhand | 5 GB | Entry-level crew member | +| bosun | 50 GB | Mid-level crew member | +| quartermaster | 100 GB | Senior crew member | +| owner (captain) | Unlimited | Hold owner always has unlimited | + +**Berth Resolution:** +1. If user is captain (owner) → unlimited +2. If crew member has explicit berth → use that berth's limit +3. If crew member has no berth → use `defaults.new_crew_berth` +4. If default berth not found → unlimited + +**Crew Record Example:** +```json +{ + "$type": "io.atcr.hold.crew", + "member": "did:plc:alice123", + "role": "writer", + "permissions": ["blob:write"], + "berth": "bosun", + "addedAt": "2026-01-04T12:00:00Z" +} +``` ### 4. Rate Limiting diff --git a/lexicons/io/atcr/hold/crew.json b/lexicons/io/atcr/hold/crew.json index f37ca5c..5a0912a 100644 --- a/lexicons/io/atcr/hold/crew.json +++ b/lexicons/io/atcr/hold/crew.json @@ -29,6 +29,11 @@ "maxLength": 64 } }, + "berth": { + "type": "string", + "description": "Optional berth (nautical rank) for quota limits (e.g., 'deckhand', 'bosun', 'quartermaster'). If empty, uses defaults.new_crew_berth from quotas.yaml.", + "maxLength": 32 + }, "addedAt": { "type": "string", "format": "datetime", diff --git a/pkg/appview/handlers/storage.go b/pkg/appview/handlers/storage.go index 62490b7..cc2577c 100644 --- a/pkg/appview/handlers/storage.go +++ b/pkg/appview/handlers/storage.go @@ -25,6 +25,8 @@ type QuotaStats struct { UserDID string `json:"userDid"` UniqueBlobs int `json:"uniqueBlobs"` TotalSize int64 `json:"totalSize"` + Limit *int64 `json:"limit,omitempty"` // nil = unlimited + Berth string `json:"berth,omitempty"` // e.g., "deckhand", "bosun", "owner" } func (h *StorageHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) { @@ -87,14 +89,36 @@ func (h *StorageHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) { } func (h *StorageHandler) renderStats(w http.ResponseWriter, stats QuotaStats) { + // Calculate usage percentage if limit exists + var usagePercent int + var hasLimit bool + var humanLimit string + + if stats.Limit != nil && *stats.Limit > 0 { + hasLimit = true + humanLimit = humanizeBytes(*stats.Limit) + usagePercent = int(float64(stats.TotalSize) / float64(*stats.Limit) * 100) + if usagePercent > 100 { + usagePercent = 100 + } + } + data := struct { - UniqueBlobs int - TotalSize int64 - HumanSize string + UniqueBlobs int + TotalSize int64 + HumanSize string + HasLimit bool + HumanLimit string + UsagePercent int + Berth string }{ - UniqueBlobs: stats.UniqueBlobs, - TotalSize: stats.TotalSize, - HumanSize: humanizeBytes(stats.TotalSize), + UniqueBlobs: stats.UniqueBlobs, + TotalSize: stats.TotalSize, + HumanSize: humanizeBytes(stats.TotalSize), + HasLimit: hasLimit, + HumanLimit: humanLimit, + UsagePercent: usagePercent, + Berth: stats.Berth, } w.Header().Set("Content-Type", "text/html") diff --git a/pkg/appview/templates/pages/settings.html b/pkg/appview/templates/pages/settings.html index 059b66e..3eacdbd 100644 --- a/pkg/appview/templates/pages/settings.html +++ b/pkg/appview/templates/pages/settings.html @@ -259,6 +259,71 @@ to { transform: rotate(360deg); } } + /* Quota Progress Bar */ + .storage-section .quota-progress { + display: flex; + align-items: center; + gap: 0.75rem; + padding: 0.75rem 0; + } + .storage-section .progress-bar { + flex: 1; + height: 8px; + background: var(--border); + border-radius: 4px; + overflow: hidden; + } + .storage-section .progress-fill { + height: 100%; + border-radius: 4px; + transition: width 0.3s ease; + } + .storage-section .progress-ok { + background: #22c55e; + } + .storage-section .progress-warning { + background: #eab308; + } + .storage-section .progress-danger { + background: #ef4444; + } + .storage-section .progress-text { + font-size: 0.85rem; + color: var(--fg-muted); + white-space: nowrap; + } + + /* Berth Badge */ + .storage-section .berth-badge { + text-transform: capitalize; + padding: 0.125rem 0.5rem; + border-radius: 4px; + font-size: 0.85rem; + background: var(--accent-bg, #e0f2fe); + color: var(--accent, #0369a1); + } + .storage-section .berth-owner { + background: #fef3c7; + color: #92400e; + } + .storage-section .berth-quartermaster { + background: #dcfce7; + color: #166534; + } + .storage-section .berth-bosun { + background: #e0e7ff; + color: #3730a3; + } + .storage-section .unlimited-badge { + font-size: 0.75rem; + padding: 0.125rem 0.375rem; + background: #22c55e; + color: #fff; + border-radius: 3px; + margin-left: 0.25rem; + font-weight: 500; + } + /* Devices Section Styles */ .devices-section .setup-instructions { margin: 1rem 0; diff --git a/pkg/appview/templates/partials/storage_stats.html b/pkg/appview/templates/partials/storage_stats.html index fb1b554..b2df312 100644 --- a/pkg/appview/templates/partials/storage_stats.html +++ b/pkg/appview/templates/partials/storage_stats.html @@ -1,12 +1,32 @@ {{ define "storage_stats" }}
+ {{ if .Berth }} +
+ Berth: + {{ .Berth }} +
+ {{ end }} +
+ Storage: + + {{ if .HasLimit }} + {{ .HumanSize }} / {{ .HumanLimit }} + {{ else }} + {{ .HumanSize }} Unlimited + {{ end }} + +
+ {{ if .HasLimit }} +
+
+
+
+ {{ .UsagePercent }}% used +
+ {{ end }}
Unique Blobs: {{ .UniqueBlobs }}
-
- Total Storage: - {{ .HumanSize }} -
{{ end }} diff --git a/pkg/atproto/cbor_gen.go b/pkg/atproto/cbor_gen.go index 46dcc8b..9285ce4 100644 --- a/pkg/atproto/cbor_gen.go +++ b/pkg/atproto/cbor_gen.go @@ -25,8 +25,13 @@ func (t *CrewRecord) MarshalCBOR(w io.Writer) error { } cw := cbg.NewCborWriter(w) + fieldCount := 6 - if _, err := cw.Write([]byte{165}); err != nil { + if t.Berth == "" { + fieldCount-- + } + + if _, err := cw.Write(cbg.CborEncodeMajorType(cbg.MajMap, uint64(fieldCount))); err != nil { return err } @@ -76,6 +81,32 @@ func (t *CrewRecord) MarshalCBOR(w io.Writer) error { return err } + // t.Berth (string) (string) + if t.Berth != "" { + + if len("berth") > 8192 { + return xerrors.Errorf("Value in field \"berth\" was too long") + } + + if err := cw.WriteMajorTypeHeader(cbg.MajTextString, uint64(len("berth"))); err != nil { + return err + } + if _, err := cw.WriteString(string("berth")); err != nil { + return err + } + + if len(t.Berth) > 8192 { + return xerrors.Errorf("Value in field t.Berth was too long") + } + + if err := cw.WriteMajorTypeHeader(cbg.MajTextString, uint64(len(t.Berth))); err != nil { + return err + } + if _, err := cw.WriteString(string(t.Berth)); err != nil { + return err + } + } + // t.Member (string) (string) if len("member") > 8192 { return xerrors.Errorf("Value in field \"member\" was too long") @@ -220,6 +251,17 @@ func (t *CrewRecord) UnmarshalCBOR(r io.Reader) (err error) { t.Type = string(sval) } + // t.Berth (string) (string) + case "berth": + + { + sval, err := cbg.ReadStringWithMax(cr, 8192) + if err != nil { + return err + } + + t.Berth = string(sval) + } // t.Member (string) (string) case "member": diff --git a/pkg/atproto/lexicon.go b/pkg/atproto/lexicon.go index 1011a42..5481235 100644 --- a/pkg/atproto/lexicon.go +++ b/pkg/atproto/lexicon.go @@ -594,7 +594,8 @@ type CrewRecord struct { Member string `json:"member" cborgen:"member"` Role string `json:"role" cborgen:"role"` Permissions []string `json:"permissions" cborgen:"permissions"` - AddedAt string `json:"addedAt" cborgen:"addedAt"` // RFC3339 timestamp + Berth string `json:"berth,omitempty" cborgen:"berth,omitempty"` // Optional berth for quota limits (nautical rank) + AddedAt string `json:"addedAt" cborgen:"addedAt"` // RFC3339 timestamp } // LayerRecord represents metadata about a container layer stored in the hold diff --git a/pkg/auth/oauth/client.go b/pkg/auth/oauth/client.go index 289ba69..3d9128b 100644 --- a/pkg/auth/oauth/client.go +++ b/pkg/auth/oauth/client.go @@ -47,14 +47,10 @@ func NewClientApp(baseURL string, store oauth.ClientAuthStore, scopes []string, return nil, fmt.Errorf("failed to configure confidential client: %w", err) } - // Log clock information for debugging timestamp issues - now := time.Now() slog.Info("Configured confidential OAuth client", "key_id", keyID, "key_path", keyPath, - "system_time_unix", now.Unix(), - "system_time_rfc3339", now.Format(time.RFC3339), - "timezone", now.Location().String()) + ) } else { config = oauth.NewLocalhostConfig(redirectURI, scopes) @@ -78,9 +74,8 @@ func RedirectURI(baseURL string) string { func GetDefaultScopes(did string) []string { return []string{ "atproto", - // Permission-set (for future PDS support) + // Permission-set // See lexicons/io/atcr/authFullApp.json for definition - // Uses "include:" prefix per ATProto permission spec "include:io.atcr.authFullApp", // com.atproto scopes must be separate (permission-sets are namespace-limited) "rpc:com.atproto.repo.getRecord?aud=*", diff --git a/pkg/hold/oci/xrpc.go b/pkg/hold/oci/xrpc.go index 216bf8c..130b88c 100644 --- a/pkg/hold/oci/xrpc.go +++ b/pkg/hold/oci/xrpc.go @@ -9,6 +9,7 @@ import ( "atcr.io/pkg/atproto" "atcr.io/pkg/hold/pds" + "atcr.io/pkg/hold/quota" "atcr.io/pkg/s3" storagedriver "github.com/distribution/distribution/v3/registry/storage/driver" "github.com/go-chi/chi/v5" @@ -23,10 +24,11 @@ type XRPCHandler struct { pds *pds.HoldPDS httpClient pds.HTTPClient enableBlueskyPosts bool + quotaMgr *quota.Manager // Quota manager for berth-based limits } // NewXRPCHandler creates a new OCI XRPC handler -func NewXRPCHandler(holdPDS *pds.HoldPDS, s3Service s3.S3Service, driver storagedriver.StorageDriver, disablePresignedURLs bool, enableBlueskyPosts bool, httpClient pds.HTTPClient) *XRPCHandler { +func NewXRPCHandler(holdPDS *pds.HoldPDS, s3Service s3.S3Service, driver storagedriver.StorageDriver, disablePresignedURLs bool, enableBlueskyPosts bool, httpClient pds.HTTPClient, quotaMgr *quota.Manager) *XRPCHandler { return &XRPCHandler{ driver: driver, disablePresignedURLs: disablePresignedURLs, @@ -35,6 +37,7 @@ func NewXRPCHandler(holdPDS *pds.HoldPDS, s3Service s3.S3Service, driver storage pds: holdPDS, httpClient: httpClient, enableBlueskyPosts: enableBlueskyPosts, + quotaMgr: quotaMgr, } } @@ -276,6 +279,24 @@ func (h *XRPCHandler) HandleNotifyManifest(w http.ResponseWriter, r *http.Reques // Only create layer records and Bluesky posts for pushes if operation == "push" { + // Soft limit check: block if ALREADY over quota + // (blobs already uploaded to S3 by this point, no sense rejecting) + stats, err := h.pds.GetQuotaForUserWithBerth(ctx, req.UserDID, h.quotaMgr) + if err == nil && stats.Limit != nil && stats.TotalSize > *stats.Limit { + slog.Warn("Quota exceeded for push", + "userDid", req.UserDID, + "currentUsage", stats.TotalSize, + "limit", *stats.Limit, + "repository", req.Repository, + "tag", req.Tag, + ) + RespondError(w, http.StatusForbidden, fmt.Sprintf( + "quota exceeded: current=%d bytes, limit=%d bytes. Delete images to free space.", + stats.TotalSize, *stats.Limit, + )) + return + } + // Check if manifest posts are enabled // Read from captain record (which is synced with HOLD_BLUESKY_POSTS_ENABLED env var) postsEnabled := false diff --git a/pkg/hold/oci/xrpc_test.go b/pkg/hold/oci/xrpc_test.go index 50b86fe..8aa08b5 100644 --- a/pkg/hold/oci/xrpc_test.go +++ b/pkg/hold/oci/xrpc_test.go @@ -127,7 +127,7 @@ func setupTestOCIHandler(t *testing.T) (*XRPCHandler, context.Context) { // Create OCI handler with buffered mode (no S3) mockS3 := s3.S3Service{} - handler := NewXRPCHandler(holdPDS, mockS3, driver, true, false, mockClient) + handler := NewXRPCHandler(holdPDS, mockS3, driver, true, false, mockClient, nil) return handler, ctx } diff --git a/pkg/hold/pds/layer.go b/pkg/hold/pds/layer.go index b5f5253..c1818cc 100644 --- a/pkg/hold/pds/layer.go +++ b/pkg/hold/pds/layer.go @@ -5,6 +5,7 @@ import ( "fmt" "atcr.io/pkg/atproto" + "atcr.io/pkg/hold/quota" lexutil "github.com/bluesky-social/indigo/lex/util" "github.com/bluesky-social/indigo/repo" ) @@ -65,6 +66,8 @@ type QuotaStats struct { UserDID string `json:"userDid"` UniqueBlobs int `json:"uniqueBlobs"` TotalSize int64 `json:"totalSize"` + Limit *int64 `json:"limit,omitempty"` // nil = unlimited + Berth string `json:"berth,omitempty"` // nautical rank for quota tier } // GetQuotaForUser calculates storage quota for a specific user @@ -160,3 +163,52 @@ func (p *HoldPDS) GetQuotaForUser(ctx context.Context, userDID string) (*QuotaSt TotalSize: totalSize, }, nil } + +// GetQuotaForUserWithBerth calculates quota with berth-aware limits +// It returns the base quota stats plus the berth limit and berth name. +// Captain (owner) always has unlimited quota. +func (p *HoldPDS) GetQuotaForUserWithBerth(ctx context.Context, userDID string, quotaMgr *quota.Manager) (*QuotaStats, error) { + // Get base stats + stats, err := p.GetQuotaForUser(ctx, userDID) + if err != nil { + return nil, err + } + + // If quota manager is nil or disabled, return unlimited + if quotaMgr == nil || !quotaMgr.IsEnabled() { + return stats, nil + } + + // Check if user is captain (owner) - always unlimited + _, captain, err := p.GetCaptainRecord(ctx) + if err == nil && captain.Owner == userDID { + stats.Berth = "owner" + // Limit remains nil (unlimited) + return stats, nil + } + + // Get crew record to find berth + crewBerth := p.getCrewBerth(ctx, userDID) + + // Resolve limit from quota manager + stats.Limit = quotaMgr.GetBerthLimit(crewBerth) + stats.Berth = quotaMgr.GetBerthName(crewBerth) + + return stats, nil +} + +// getCrewBerth returns the berth for a crew member, or empty string if not found +func (p *HoldPDS) getCrewBerth(ctx context.Context, userDID string) string { + crewMembers, err := p.ListCrewMembers(ctx) + if err != nil { + return "" + } + + for _, member := range crewMembers { + if member.Record.Member == userDID { + return member.Record.Berth + } + } + + return "" +} diff --git a/pkg/hold/pds/layer_test.go b/pkg/hold/pds/layer_test.go index 4c6c66a..959616b 100644 --- a/pkg/hold/pds/layer_test.go +++ b/pkg/hold/pds/layer_test.go @@ -1,9 +1,12 @@ package pds import ( + "os" + "path/filepath" "testing" "atcr.io/pkg/atproto" + "atcr.io/pkg/hold/quota" ) func TestCreateLayerRecord(t *testing.T) { @@ -281,3 +284,436 @@ func TestLayerRecord_FieldValidation(t *testing.T) { }) } } + +// setupTestPDSWithIndex creates a PDS with file-based database (enables RecordsIndex) +// and bootstraps it with the given owner. Required for quota tests. +func setupTestPDSWithIndex(t *testing.T, ownerDID string) (*HoldPDS, func()) { + t.Helper() + + ctx := sharedCtx + tmpDir := t.TempDir() + + // Use file-based database to enable RecordsIndex + dbPath := filepath.Join(tmpDir, "pds.db") + keyPath := filepath.Join(tmpDir, "signing-key") + + // Copy shared signing key + if err := os.WriteFile(keyPath, sharedTestKey, 0600); err != nil { + t.Fatalf("Failed to copy shared signing key: %v", err) + } + + pds, err := NewHoldPDS(ctx, "did:web:hold.example.com", "https://hold.example.com", dbPath, keyPath, false) + if err != nil { + t.Fatalf("Failed to create test PDS: %v", err) + } + + // Bootstrap with owner + if err := pds.Bootstrap(ctx, nil, ownerDID, true, false, ""); err != nil { + t.Fatalf("Failed to bootstrap PDS: %v", err) + } + + // Wire up records indexing + indexingHandler := pds.CreateRecordsIndexEventHandler(nil) + pds.RepomgrRef().SetEventHandler(indexingHandler, true) + + // Backfill index from MST + if err := pds.BackfillRecordsIndex(ctx); err != nil { + t.Fatalf("Failed to backfill records index: %v", err) + } + + cleanup := func() { + pds.Close() + } + + return pds, cleanup +} + +// addCrewMemberWithBerth adds a crew member with a specific berth (nautical rank) +func addCrewMemberWithBerth(t *testing.T, pds *HoldPDS, memberDID, role string, permissions []string, berth string) { + t.Helper() + + crewRecord := &atproto.CrewRecord{ + Type: atproto.CrewCollection, + Member: memberDID, + Role: role, + Permissions: permissions, + Berth: berth, + AddedAt: "2026-01-04T12:00:00Z", + } + + _, _, err := pds.repomgr.CreateRecord(sharedCtx, pds.uid, atproto.CrewCollection, crewRecord) + if err != nil { + t.Fatalf("Failed to add crew member with berth: %v", err) + } +} + +func TestGetQuotaForUserWithBerth_OwnerUnlimited(t *testing.T) { + ownerDID := "did:plc:owner123" + pds, cleanup := setupTestPDSWithIndex(t, ownerDID) + defer cleanup() + + ctx := sharedCtx + + // Create quota manager with config + tmpDir := t.TempDir() + configPath := filepath.Join(tmpDir, "quotas.yaml") + configContent := ` +berths: + deckhand: + quota: 5GB + bosun: + quota: 50GB + +defaults: + new_crew_berth: deckhand +` + if err := os.WriteFile(configPath, []byte(configContent), 0644); err != nil { + t.Fatalf("Failed to write quota config: %v", err) + } + + quotaMgr, err := quota.NewManager(configPath) + if err != nil { + t.Fatalf("Failed to create quota manager: %v", err) + } + + // Create layer records for owner + for i := 0; i < 3; i++ { + record := atproto.NewLayerRecord( + "sha256:owner"+string(rune('a'+i)), + 1024*1024*100, // 100MB each + "application/vnd.oci.image.layer.v1.tar+gzip", + ownerDID, + "at://"+ownerDID+"/io.atcr.manifest/test123", + ) + if _, _, err := pds.CreateLayerRecord(ctx, record); err != nil { + t.Fatalf("Failed to create layer record: %v", err) + } + } + + // Get quota for owner + stats, err := pds.GetQuotaForUserWithBerth(ctx, ownerDID, quotaMgr) + if err != nil { + t.Fatalf("GetQuotaForUserWithBerth failed: %v", err) + } + + // Owner should have unlimited quota (nil limit) + if stats.Limit != nil { + t.Errorf("Expected nil limit for owner, got %d", *stats.Limit) + } + + // Berth should be "owner" + if stats.Berth != "owner" { + t.Errorf("Expected berth 'owner', got %q", stats.Berth) + } + + // Should have 3 unique blobs + if stats.UniqueBlobs != 3 { + t.Errorf("Expected 3 unique blobs, got %d", stats.UniqueBlobs) + } + + // Total size should be 300MB + expectedSize := int64(3 * 100 * 1024 * 1024) + if stats.TotalSize != expectedSize { + t.Errorf("Expected total size %d, got %d", expectedSize, stats.TotalSize) + } + + t.Logf("Owner quota stats: %+v", stats) +} + +func TestGetQuotaForUserWithBerth_CrewWithDefaultBerth(t *testing.T) { + ownerDID := "did:plc:owner456" + crewDID := "did:plc:crew123" + pds, cleanup := setupTestPDSWithIndex(t, ownerDID) + defer cleanup() + + ctx := sharedCtx + + // Create quota manager + tmpDir := t.TempDir() + configPath := filepath.Join(tmpDir, "quotas.yaml") + configContent := ` +berths: + deckhand: + quota: 5GB + bosun: + quota: 50GB + +defaults: + new_crew_berth: deckhand +` + if err := os.WriteFile(configPath, []byte(configContent), 0644); err != nil { + t.Fatalf("Failed to write quota config: %v", err) + } + + quotaMgr, err := quota.NewManager(configPath) + if err != nil { + t.Fatalf("Failed to create quota manager: %v", err) + } + + // Add crew member with no berth (should use default) + addCrewMemberWithBerth(t, pds, crewDID, "writer", []string{"blob:write"}, "") + + // Create layer records for crew member + for i := 0; i < 2; i++ { + record := atproto.NewLayerRecord( + "sha256:crew"+string(rune('a'+i)), + 1024*1024*50, // 50MB each + "application/vnd.oci.image.layer.v1.tar+gzip", + crewDID, + "at://"+crewDID+"/io.atcr.manifest/test456", + ) + if _, _, err := pds.CreateLayerRecord(ctx, record); err != nil { + t.Fatalf("Failed to create layer record: %v", err) + } + } + + // Get quota for crew member + stats, err := pds.GetQuotaForUserWithBerth(ctx, crewDID, quotaMgr) + if err != nil { + t.Fatalf("GetQuotaForUserWithBerth failed: %v", err) + } + + // Should have 5GB limit (deckhand berth) + expectedLimit := int64(5 * 1024 * 1024 * 1024) + if stats.Limit == nil { + t.Fatal("Expected non-nil limit for crew member") + } + if *stats.Limit != expectedLimit { + t.Errorf("Expected limit %d, got %d", expectedLimit, *stats.Limit) + } + + // Berth should be "deckhand" + if stats.Berth != "deckhand" { + t.Errorf("Expected berth 'deckhand', got %q", stats.Berth) + } + + // Should have 2 unique blobs + if stats.UniqueBlobs != 2 { + t.Errorf("Expected 2 unique blobs, got %d", stats.UniqueBlobs) + } + + t.Logf("Crew (deckhand berth) quota stats: %+v", stats) +} + +func TestGetQuotaForUserWithBerth_CrewWithExplicitBerth(t *testing.T) { + ownerDID := "did:plc:owner789" + crewDID := "did:plc:bosuncrew456" + pds, cleanup := setupTestPDSWithIndex(t, ownerDID) + defer cleanup() + + ctx := sharedCtx + + // Create quota manager + tmpDir := t.TempDir() + configPath := filepath.Join(tmpDir, "quotas.yaml") + configContent := ` +berths: + deckhand: + quota: 5GB + bosun: + quota: 50GB + +defaults: + new_crew_berth: deckhand +` + if err := os.WriteFile(configPath, []byte(configContent), 0644); err != nil { + t.Fatalf("Failed to write quota config: %v", err) + } + + quotaMgr, err := quota.NewManager(configPath) + if err != nil { + t.Fatalf("Failed to create quota manager: %v", err) + } + + // Add crew member with explicit "bosun" berth + addCrewMemberWithBerth(t, pds, crewDID, "writer", []string{"blob:write"}, "bosun") + + // Create layer records for crew member + record := atproto.NewLayerRecord( + "sha256:bosunlayer1", + 1024*1024*1024, // 1GB + "application/vnd.oci.image.layer.v1.tar+gzip", + crewDID, + "at://"+crewDID+"/io.atcr.manifest/test789", + ) + if _, _, err := pds.CreateLayerRecord(ctx, record); err != nil { + t.Fatalf("Failed to create layer record: %v", err) + } + + // Get quota for crew member + stats, err := pds.GetQuotaForUserWithBerth(ctx, crewDID, quotaMgr) + if err != nil { + t.Fatalf("GetQuotaForUserWithBerth failed: %v", err) + } + + // Should have 50GB limit (bosun berth) + expectedLimit := int64(50 * 1024 * 1024 * 1024) + if stats.Limit == nil { + t.Fatal("Expected non-nil limit for crew member") + } + if *stats.Limit != expectedLimit { + t.Errorf("Expected limit %d, got %d", expectedLimit, *stats.Limit) + } + + // Berth should be "bosun" + if stats.Berth != "bosun" { + t.Errorf("Expected berth 'bosun', got %q", stats.Berth) + } + + t.Logf("Crew (bosun berth) quota stats: %+v", stats) +} + +func TestGetQuotaForUserWithBerth_NoQuotaManager(t *testing.T) { + ownerDID := "did:plc:ownerabc" + crewDID := "did:plc:crewabc" + pds, cleanup := setupTestPDSWithIndex(t, ownerDID) + defer cleanup() + + ctx := sharedCtx + + // Add crew member + addCrewMemberWithBerth(t, pds, crewDID, "writer", []string{"blob:write"}, "deckhand") + + // Create layer record + record := atproto.NewLayerRecord( + "sha256:noquotalayer1", + 1024*1024*100, + "application/vnd.oci.image.layer.v1.tar+gzip", + crewDID, + "at://"+crewDID+"/io.atcr.manifest/testabc", + ) + if _, _, err := pds.CreateLayerRecord(ctx, record); err != nil { + t.Fatalf("Failed to create layer record: %v", err) + } + + // Get quota with nil quota manager (no enforcement) + stats, err := pds.GetQuotaForUserWithBerth(ctx, crewDID, nil) + if err != nil { + t.Fatalf("GetQuotaForUserWithBerth failed: %v", err) + } + + // Should have nil limit (unlimited) + if stats.Limit != nil { + t.Errorf("Expected nil limit when quota manager is nil, got %d", *stats.Limit) + } + + // Berth should be empty + if stats.Berth != "" { + t.Errorf("Expected empty berth, got %q", stats.Berth) + } + + t.Logf("No quota manager stats: %+v", stats) +} + +func TestGetQuotaForUserWithBerth_DisabledQuotas(t *testing.T) { + ownerDID := "did:plc:ownerdef" + crewDID := "did:plc:crewdef" + pds, cleanup := setupTestPDSWithIndex(t, ownerDID) + defer cleanup() + + ctx := sharedCtx + + // Create quota manager with nonexistent config (disabled) + quotaMgr, err := quota.NewManager("/nonexistent/quotas.yaml") + if err != nil { + t.Fatalf("Failed to create quota manager: %v", err) + } + + if quotaMgr.IsEnabled() { + t.Fatal("Expected quotas to be disabled") + } + + // Add crew member + addCrewMemberWithBerth(t, pds, crewDID, "writer", []string{"blob:write"}, "bosun") + + // Create layer record + record := atproto.NewLayerRecord( + "sha256:disabledlayer1", + 1024*1024*100, + "application/vnd.oci.image.layer.v1.tar+gzip", + crewDID, + "at://"+crewDID+"/io.atcr.manifest/testdef", + ) + if _, _, err := pds.CreateLayerRecord(ctx, record); err != nil { + t.Fatalf("Failed to create layer record: %v", err) + } + + // Get quota with disabled quota manager + stats, err := pds.GetQuotaForUserWithBerth(ctx, crewDID, quotaMgr) + if err != nil { + t.Fatalf("GetQuotaForUserWithBerth failed: %v", err) + } + + // Should have nil limit (unlimited when quotas disabled) + if stats.Limit != nil { + t.Errorf("Expected nil limit when quotas disabled, got %d", *stats.Limit) + } + + t.Logf("Disabled quotas stats: %+v", stats) +} + +func TestGetQuotaForUserWithBerth_DeduplicatesBlobs(t *testing.T) { + ownerDID := "did:plc:ownerghi" + crewDID := "did:plc:crewghi" + pds, cleanup := setupTestPDSWithIndex(t, ownerDID) + defer cleanup() + + ctx := sharedCtx + + // Create quota manager + tmpDir := t.TempDir() + configPath := filepath.Join(tmpDir, "quotas.yaml") + configContent := ` +berths: + deckhand: + quota: 5GB + +defaults: + new_crew_berth: deckhand +` + if err := os.WriteFile(configPath, []byte(configContent), 0644); err != nil { + t.Fatalf("Failed to write quota config: %v", err) + } + + quotaMgr, err := quota.NewManager(configPath) + if err != nil { + t.Fatalf("Failed to create quota manager: %v", err) + } + + // Add crew member + addCrewMemberWithBerth(t, pds, crewDID, "writer", []string{"blob:write"}, "") + + // Create multiple layer records with same digest (should be deduplicated) + digest := "sha256:duplicatelayer" + for i := 0; i < 5; i++ { + record := atproto.NewLayerRecord( + digest, + 1024*1024*100, // 100MB + "application/vnd.oci.image.layer.v1.tar+gzip", + crewDID, + "at://"+crewDID+"/io.atcr.manifest/manifest"+string(rune('a'+i)), + ) + if _, _, err := pds.CreateLayerRecord(ctx, record); err != nil { + t.Fatalf("Failed to create layer record %d: %v", i, err) + } + } + + // Get quota + stats, err := pds.GetQuotaForUserWithBerth(ctx, crewDID, quotaMgr) + if err != nil { + t.Fatalf("GetQuotaForUserWithBerth failed: %v", err) + } + + // Should have 1 unique blob (deduplicated) + if stats.UniqueBlobs != 1 { + t.Errorf("Expected 1 unique blob (deduplicated), got %d", stats.UniqueBlobs) + } + + // Total size should be 100MB (not 500MB) + expectedSize := int64(100 * 1024 * 1024) + if stats.TotalSize != expectedSize { + t.Errorf("Expected total size %d, got %d", expectedSize, stats.TotalSize) + } + + t.Logf("Deduplicated quota stats: %+v", stats) +} diff --git a/pkg/hold/pds/status_test.go b/pkg/hold/pds/status_test.go index 5b1b78b..1227096 100644 --- a/pkg/hold/pds/status_test.go +++ b/pkg/hold/pds/status_test.go @@ -55,7 +55,7 @@ func TestStatusPost(t *testing.T) { } // Create handler for XRPC endpoints - handler := NewXRPCHandler(holdPDS, s3.S3Service{}, nil, nil, &mockPDSClient{}) + handler := NewXRPCHandler(holdPDS, s3.S3Service{}, nil, nil, &mockPDSClient{}, nil) // Helper function to list posts via XRPC listPosts := func() ([]map[string]any, error) { @@ -283,7 +283,7 @@ func TestMain(m *testing.M) { } // Create shared handler - sharedHandler = NewXRPCHandler(sharedPDS, s3.S3Service{}, nil, nil, &mockPDSClient{}) + sharedHandler = NewXRPCHandler(sharedPDS, s3.S3Service{}, nil, nil, &mockPDSClient{}, nil) // Run tests code := m.Run() diff --git a/pkg/hold/pds/xrpc.go b/pkg/hold/pds/xrpc.go index 4219723..8645ce5 100644 --- a/pkg/hold/pds/xrpc.go +++ b/pkg/hold/pds/xrpc.go @@ -7,6 +7,7 @@ import ( "fmt" "atcr.io/pkg/atproto" + "atcr.io/pkg/hold/quota" "atcr.io/pkg/s3" "github.com/bluesky-social/indigo/api/bsky" lexutil "github.com/bluesky-social/indigo/lex/util" @@ -46,7 +47,8 @@ type XRPCHandler struct { s3Service s3.S3Service storageDriver driver.StorageDriver broadcaster *EventBroadcaster - httpClient HTTPClient // For testing - allows injecting mock HTTP client + httpClient HTTPClient // For testing - allows injecting mock HTTP client + quotaMgr *quota.Manager // Quota manager for tier-based limits } // PartInfo represents a completed part in a multipart upload @@ -64,13 +66,14 @@ type PartUploadInfo struct { } // NewXRPCHandler creates a new XRPC handler -func NewXRPCHandler(pds *HoldPDS, s3Service s3.S3Service, storageDriver driver.StorageDriver, broadcaster *EventBroadcaster, httpClient HTTPClient) *XRPCHandler { +func NewXRPCHandler(pds *HoldPDS, s3Service s3.S3Service, storageDriver driver.StorageDriver, broadcaster *EventBroadcaster, httpClient HTTPClient, quotaMgr *quota.Manager) *XRPCHandler { return &XRPCHandler{ pds: pds, s3Service: s3Service, storageDriver: storageDriver, broadcaster: broadcaster, httpClient: httpClient, + quotaMgr: quotaMgr, } } @@ -1520,6 +1523,7 @@ func getProxyURL(publicURL string, digest, did string, operation string) string // HandleGetQuota returns storage quota information for a user // This calculates the total unique blob storage used by a specific user // by iterating layer records and deduplicating by digest. +// Also returns tier-aware quota limits if quotas.yaml is configured. func (h *XRPCHandler) HandleGetQuota(w http.ResponseWriter, r *http.Request) { userDID := r.URL.Query().Get("userDid") if userDID == "" { @@ -1533,8 +1537,8 @@ func (h *XRPCHandler) HandleGetQuota(w http.ResponseWriter, r *http.Request) { return } - // Get quota stats - stats, err := h.pds.GetQuotaForUser(r.Context(), userDID) + // Get quota stats with berth-aware limits + stats, err := h.pds.GetQuotaForUserWithBerth(r.Context(), userDID, h.quotaMgr) if err != nil { slog.Error("Failed to get quota", "userDid", userDID, "error", err) http.Error(w, fmt.Sprintf("failed to get quota: %v", err), http.StatusInternalServerError) diff --git a/pkg/hold/pds/xrpc_test.go b/pkg/hold/pds/xrpc_test.go index 225c4ec..a5724b0 100644 --- a/pkg/hold/pds/xrpc_test.go +++ b/pkg/hold/pds/xrpc_test.go @@ -76,7 +76,7 @@ func setupTestXRPCHandler(t *testing.T) (*XRPCHandler, context.Context) { mockS3 := s3.S3Service{} // Create XRPC handler with mock HTTP client - handler := NewXRPCHandler(pds, mockS3, nil, nil, mockClient) + handler := NewXRPCHandler(pds, mockS3, nil, nil, mockClient, nil) return handler, ctx } @@ -143,7 +143,7 @@ func setupTestXRPCHandlerWithIndex(t *testing.T) (*XRPCHandler, context.Context) mockS3 := s3.S3Service{} // Create XRPC handler with mock HTTP client - handler := NewXRPCHandler(pds, mockS3, nil, nil, mockClient) + handler := NewXRPCHandler(pds, mockS3, nil, nil, mockClient, nil) return handler, ctx } @@ -753,7 +753,7 @@ func TestHandleListRecords_EmptyCollection(t *testing.T) { pds, ctx := setupTestPDS(t) // Don't bootstrap - no records created yet mockClient := &mockPDSClient{} mockS3 := s3.S3Service{} - handler := NewXRPCHandler(pds, mockS3, nil, nil, mockClient) + handler := NewXRPCHandler(pds, mockS3, nil, nil, mockClient, nil) // Initialize repo manually (setupTestPDS doesn't call Bootstrap, so no crew members) err := pds.repomgr.InitNewActor(ctx, pds.uid, "", pds.did, "", "", "") @@ -1231,7 +1231,7 @@ func TestHandleListRepos_EmptyRepo(t *testing.T) { pds, ctx := setupTestPDS(t) // Don't bootstrap mockClient := &mockPDSClient{} mockS3 := s3.S3Service{} - handler := NewXRPCHandler(pds, mockS3, nil, nil, mockClient) + handler := NewXRPCHandler(pds, mockS3, nil, nil, mockClient, nil) // setupTestPDS creates the PDS/database but doesn't initialize the repo // Check if implementation returns repos before initialization @@ -1317,7 +1317,7 @@ func TestHandleGetRepoStatus_EmptyRepo(t *testing.T) { pds, ctx := setupTestPDS(t) // Don't bootstrap mockClient := &mockPDSClient{} mockS3 := s3.S3Service{} - handler := NewXRPCHandler(pds, mockS3, nil, nil, mockClient) + handler := NewXRPCHandler(pds, mockS3, nil, nil, mockClient, nil) holdDID := "did:web:hold.example.com" // Initialize repo but don't add any records @@ -2014,7 +2014,7 @@ func setupTestXRPCHandlerWithBlobs(t *testing.T) (*XRPCHandler, *mockS3Service, mockClient := &mockPDSClient{} // Create XRPC handler with mock s3 service and real filesystem driver - handler := NewXRPCHandler(pds, mockS3Svc.toS3Service(), driver, nil, mockClient) + handler := NewXRPCHandler(pds, mockS3Svc.toS3Service(), driver, nil, mockClient, nil) return handler, mockS3Svc, ctx } diff --git a/pkg/hold/quota/config.go b/pkg/hold/quota/config.go new file mode 100644 index 0000000..528068e --- /dev/null +++ b/pkg/hold/quota/config.go @@ -0,0 +1,197 @@ +package quota + +import ( + "fmt" + "os" + "regexp" + "strconv" + "strings" + + "go.yaml.in/yaml/v4" +) + +// Config represents the quotas.yaml configuration +type Config struct { + Berths map[string]BerthConfig `yaml:"berths"` + Defaults DefaultsConfig `yaml:"defaults"` +} + +// BerthConfig represents a single berth's configuration +type BerthConfig struct { + Quota string `yaml:"quota"` // Human-readable size: "5GB", "50GB", etc. +} + +// DefaultsConfig represents default settings +type DefaultsConfig struct { + NewCrewBerth string `yaml:"new_crew_berth"` +} + +// Manager manages quota configuration and berth resolution +type Manager struct { + config *Config + berths map[string]int64 // resolved berth name -> bytes +} + +// NewManager creates a quota manager, loading config from file if present +func NewManager(configPath string) (*Manager, error) { + m := &Manager{ + berths: make(map[string]int64), + } + + // Try to load config file + data, err := os.ReadFile(configPath) + if err != nil { + if os.IsNotExist(err) { + // No config file = no quotas enforced + return m, nil + } + return nil, fmt.Errorf("failed to read quota config: %w", err) + } + + var cfg Config + if err := yaml.Unmarshal(data, &cfg); err != nil { + return nil, fmt.Errorf("failed to parse quota config: %w", err) + } + + m.config = &cfg + + // Parse and resolve all berths + for name, berth := range cfg.Berths { + bytes, err := ParseHumanBytes(berth.Quota) + if err != nil { + return nil, fmt.Errorf("invalid quota for berth %q: %w", name, err) + } + m.berths[name] = bytes + } + + return m, nil +} + +// IsEnabled returns true if quotas are being enforced +func (m *Manager) IsEnabled() bool { + return m.config != nil +} + +// GetBerthLimit resolves the quota limit for a berth key +// Returns nil for unlimited (captain, no config, or berth not found with no default) +// +// Resolution order: +// 1. If quotas disabled → nil (unlimited) +// 2. If berthKey provided and found → return that berth's limit +// 3. If berthKey not found or empty → use defaults.new_crew_berth +// 4. If default berth not found → nil (unlimited) +func (m *Manager) GetBerthLimit(berthKey string) *int64 { + if !m.IsEnabled() { + return nil + } + + // Try the provided berth key first + if berthKey != "" { + if limit, ok := m.berths[berthKey]; ok { + return &limit + } + } + + // Fall back to default berth + if m.config.Defaults.NewCrewBerth != "" { + if limit, ok := m.berths[m.config.Defaults.NewCrewBerth]; ok { + return &limit + } + } + + // No valid berth found - unlimited + return nil +} + +// GetBerthName resolves the berth name for a berth key +// Returns the actual berth name being used (after fallback resolution) +func (m *Manager) GetBerthName(berthKey string) string { + if !m.IsEnabled() { + return "" + } + + // Try the provided berth key first + if berthKey != "" { + if _, ok := m.berths[berthKey]; ok { + return berthKey + } + } + + // Fall back to default berth + if m.config.Defaults.NewCrewBerth != "" { + if _, ok := m.berths[m.config.Defaults.NewCrewBerth]; ok { + return m.config.Defaults.NewCrewBerth + } + } + + return "" +} + +// GetDefaultBerth returns the default berth name for new crew members +func (m *Manager) GetDefaultBerth() string { + if m.config == nil { + return "" + } + return m.config.Defaults.NewCrewBerth +} + +// BerthCount returns the number of configured berths +func (m *Manager) BerthCount() int { + return len(m.berths) +} + +// ParseHumanBytes parses human-readable byte sizes like "5GB", "100MB", "1.5TB" +func ParseHumanBytes(s string) (int64, error) { + s = strings.TrimSpace(strings.ToUpper(s)) + if s == "" { + return 0, fmt.Errorf("empty size string") + } + + // Match number (with optional decimal) followed by optional unit + re := regexp.MustCompile(`^(\d+(?:\.\d+)?)\s*(B|KB|MB|GB|TB|PB)?$`) + matches := re.FindStringSubmatch(s) + if matches == nil { + return 0, fmt.Errorf("invalid size format: %s", s) + } + + value, err := strconv.ParseFloat(matches[1], 64) + if err != nil { + return 0, fmt.Errorf("invalid number: %w", err) + } + + unit := matches[2] + if unit == "" { + unit = "B" + } + + multipliers := map[string]float64{ + "B": 1, + "KB": 1024, + "MB": 1024 * 1024, + "GB": 1024 * 1024 * 1024, + "TB": 1024 * 1024 * 1024 * 1024, + "PB": 1024 * 1024 * 1024 * 1024 * 1024, + } + + mult, ok := multipliers[unit] + if !ok { + return 0, fmt.Errorf("unknown unit: %s", unit) + } + + return int64(value * mult), nil +} + +// FormatHumanBytes formats bytes as a human-readable string +func FormatHumanBytes(bytes int64) string { + const unit = 1024 + if bytes < unit { + return fmt.Sprintf("%d B", bytes) + } + div, exp := int64(unit), 0 + for n := bytes / unit; n >= unit; n /= unit { + div *= unit + exp++ + } + units := []string{"KB", "MB", "GB", "TB", "PB"} + return fmt.Sprintf("%.1f %s", float64(bytes)/float64(div), units[exp]) +} diff --git a/pkg/hold/quota/config_test.go b/pkg/hold/quota/config_test.go new file mode 100644 index 0000000..70aafdf --- /dev/null +++ b/pkg/hold/quota/config_test.go @@ -0,0 +1,275 @@ +package quota + +import ( + "os" + "path/filepath" + "testing" +) + +func TestParseHumanBytes(t *testing.T) { + tests := []struct { + input string + expected int64 + wantErr bool + }{ + // Basic units + {"1024", 1024, false}, + {"1024B", 1024, false}, + {"1KB", 1024, false}, + {"1MB", 1024 * 1024, false}, + {"1GB", 1024 * 1024 * 1024, false}, + {"1TB", 1024 * 1024 * 1024 * 1024, false}, + {"1PB", 1024 * 1024 * 1024 * 1024 * 1024, false}, + + // Common sizes + {"5GB", 5 * 1024 * 1024 * 1024, false}, + {"50GB", 50 * 1024 * 1024 * 1024, false}, + {"100GB", 100 * 1024 * 1024 * 1024, false}, + {"500MB", 500 * 1024 * 1024, false}, + + // Case insensitive + {"5gb", 5 * 1024 * 1024 * 1024, false}, + {"5Gb", 5 * 1024 * 1024 * 1024, false}, + + // With whitespace + {" 5GB ", 5 * 1024 * 1024 * 1024, false}, + + // Decimals + {"1.5GB", int64(1.5 * 1024 * 1024 * 1024), false}, + {"2.5TB", int64(2.5 * 1024 * 1024 * 1024 * 1024), false}, + + // Errors + {"", 0, true}, + {"invalid", 0, true}, + {"GB", 0, true}, + {"-5GB", 0, true}, + } + + for _, tt := range tests { + t.Run(tt.input, func(t *testing.T) { + result, err := ParseHumanBytes(tt.input) + if tt.wantErr { + if err == nil { + t.Errorf("expected error for input %q", tt.input) + } + return + } + if err != nil { + t.Errorf("unexpected error: %v", err) + return + } + if result != tt.expected { + t.Errorf("got %d, want %d", result, tt.expected) + } + }) + } +} + +func TestFormatHumanBytes(t *testing.T) { + tests := []struct { + bytes int64 + expected string + }{ + {0, "0 B"}, + {512, "512 B"}, + {1024, "1.0 KB"}, + {1024 * 1024, "1.0 MB"}, + {1024 * 1024 * 1024, "1.0 GB"}, + {5 * 1024 * 1024 * 1024, "5.0 GB"}, + {1024 * 1024 * 1024 * 1024, "1.0 TB"}, + } + + for _, tt := range tests { + t.Run(tt.expected, func(t *testing.T) { + result := FormatHumanBytes(tt.bytes) + if result != tt.expected { + t.Errorf("got %q, want %q", result, tt.expected) + } + }) + } +} + +func TestNewManager_NoConfigFile(t *testing.T) { + m, err := NewManager("/nonexistent/quotas.yaml") + if err != nil { + t.Fatalf("expected no error for missing file, got: %v", err) + } + if m.IsEnabled() { + t.Error("expected quotas to be disabled when file missing") + } + if m.GetBerthLimit("anything") != nil { + t.Error("expected nil limit when quotas disabled") + } +} + +func TestNewManager_ValidConfig(t *testing.T) { + tmpDir := t.TempDir() + configPath := filepath.Join(tmpDir, "quotas.yaml") + + configContent := ` +berths: + deckhand: + quota: 5GB + bosun: + quota: 50GB + quartermaster: + quota: 100GB + +defaults: + new_crew_berth: deckhand +` + if err := os.WriteFile(configPath, []byte(configContent), 0644); err != nil { + t.Fatalf("failed to write config: %v", err) + } + + m, err := NewManager(configPath) + if err != nil { + t.Fatalf("failed to load config: %v", err) + } + + if !m.IsEnabled() { + t.Error("expected quotas to be enabled") + } + + if m.BerthCount() != 3 { + t.Errorf("expected 3 berths, got %d", m.BerthCount()) + } + + // Test default berth (empty string) + limit := m.GetBerthLimit("") + if limit == nil { + t.Fatal("expected non-nil limit for default berth") + } + if *limit != 5*1024*1024*1024 { + t.Errorf("expected 5GB limit for default, got %d", *limit) + } + + // Test explicit berth + limit = m.GetBerthLimit("bosun") + if limit == nil { + t.Fatal("expected non-nil limit for bosun") + } + if *limit != 50*1024*1024*1024 { + t.Errorf("expected 50GB limit for bosun, got %d", *limit) + } + + // Test berth name resolution + if m.GetBerthName("") != "deckhand" { + t.Errorf("expected berth name 'deckhand' for empty key, got %q", m.GetBerthName("")) + } + if m.GetBerthName("bosun") != "bosun" { + t.Errorf("expected berth name 'bosun', got %q", m.GetBerthName("bosun")) + } +} + +func TestNewManager_FallbackToDefault(t *testing.T) { + tmpDir := t.TempDir() + configPath := filepath.Join(tmpDir, "quotas.yaml") + + configContent := ` +berths: + deckhand: + quota: 5GB + quartermaster: + quota: 50GB + +defaults: + new_crew_berth: deckhand +` + if err := os.WriteFile(configPath, []byte(configContent), 0644); err != nil { + t.Fatalf("failed to write config: %v", err) + } + + m, err := NewManager(configPath) + if err != nil { + t.Fatalf("failed to load config: %v", err) + } + + // Unknown berth should fall back to default + limit := m.GetBerthLimit("unknown_berth") + if limit == nil { + t.Fatal("expected fallback to default berth") + } + if *limit != 5*1024*1024*1024 { + t.Errorf("expected 5GB limit from default fallback, got %d", *limit) + } + + // Berth name should also fall back + if m.GetBerthName("unknown_berth") != "deckhand" { + t.Errorf("expected berth name 'deckhand' for unknown berth, got %q", m.GetBerthName("unknown_berth")) + } +} + +func TestNewManager_InvalidConfig(t *testing.T) { + tmpDir := t.TempDir() + configPath := filepath.Join(tmpDir, "quotas.yaml") + + // Invalid YAML + if err := os.WriteFile(configPath, []byte("invalid: [yaml"), 0644); err != nil { + t.Fatalf("failed to write config: %v", err) + } + + _, err := NewManager(configPath) + if err == nil { + t.Error("expected error for invalid YAML") + } +} + +func TestNewManager_InvalidQuotaSize(t *testing.T) { + tmpDir := t.TempDir() + configPath := filepath.Join(tmpDir, "quotas.yaml") + + configContent := ` +berths: + deckhand: + quota: invalid_size + +defaults: + new_crew_berth: deckhand +` + if err := os.WriteFile(configPath, []byte(configContent), 0644); err != nil { + t.Fatalf("failed to write config: %v", err) + } + + _, err := NewManager(configPath) + if err == nil { + t.Error("expected error for invalid quota size") + } +} + +func TestNewManager_NoDefaultBerth(t *testing.T) { + tmpDir := t.TempDir() + configPath := filepath.Join(tmpDir, "quotas.yaml") + + configContent := ` +berths: + quartermaster: + quota: 50GB + +defaults: + new_crew_berth: nonexistent +` + if err := os.WriteFile(configPath, []byte(configContent), 0644); err != nil { + t.Fatalf("failed to write config: %v", err) + } + + m, err := NewManager(configPath) + if err != nil { + t.Fatalf("failed to load config: %v", err) + } + + // Empty berth key with nonexistent default should return nil (unlimited) + limit := m.GetBerthLimit("") + if limit != nil { + t.Error("expected nil limit when default berth doesn't exist") + } + + // Explicit berth should still work + limit = m.GetBerthLimit("quartermaster") + if limit == nil { + t.Fatal("expected non-nil limit for quartermaster berth") + } + if *limit != 50*1024*1024*1024 { + t.Errorf("expected 50GB limit for quartermaster, got %d", *limit) + } +} diff --git a/quotas.yaml.example b/quotas.yaml.example new file mode 100644 index 0000000..d92bef4 --- /dev/null +++ b/quotas.yaml.example @@ -0,0 +1,35 @@ +# ATCR Hold Service Quota Configuration +# Copy this file to quotas.yaml to enable quota enforcement. +# If quotas.yaml doesn't exist, quotas are disabled (unlimited for all users). + +# Berths define quota tiers using nautical crew ranks. +# Each berth has a quota limit specified in human-readable format. +# Supported units: B, KB, MB, GB, TB, PB (case-insensitive) +berths: + # Entry-level crew - suitable for new or casual users + deckhand: + quota: 5GB + + # Mid-level crew - for regular contributors + bosun: + quota: 50GB + + # Senior crew - for power users or trusted contributors + quartermaster: + quota: 100GB + + # You can add custom berths with any name: + # unlimited_crew: + # quota: 1TB + +defaults: + # Default berth assigned to new crew members who don't have an explicit berth. + # This berth must exist in the berths section above. + new_crew_berth: deckhand + +# Notes: +# - The hold captain (owner) always has unlimited quota regardless of berths. +# - Crew members can be assigned a specific berth in their crew record. +# - If a crew member's berth doesn't exist in config, they fall back to the default. +# - Quota is calculated per-user by summing unique blob sizes (deduplicated). +# - Quota is checked when pushing manifests (after blobs are already uploaded).