mirrors/at-container-registry
1
0
Fork 0
mirror of https://tangled.org/evan.jarrett.net/at-container-registry synced 2026-08-31 05:07:09 +00:00
Code Issues Packages Projects Releases Wiki Activity

basic implementation of quotas

Browse Source
This commit is contained in:
Evan Jarrett
2026-01-04 20:09:41 -06:00
parent 53e196a261
commit aa4b32bbd6
21 changed files with 1297 additions and 41 deletions
Download Patch File Download Diff File Expand all files Collapse all files
cmd/hold/main.go
+15 -2
View File
@@ -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
deploy/docker-compose.prod.yml
+1
View File
@@ -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:
deploy/quotas.yaml
+35
View File
@@ -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).
docs/QUOTAS.md
+41 -6
View File
@@ -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
lexicons/io/atcr/hold/crew.json
+5
View File
@@ -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",
pkg/appview/handlers/storage.go
+30 -6
View File
@@ -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")
pkg/appview/templates/pages/settings.html
+65
View File
@@ -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;
pkg/appview/templates/partials/storage_stats.html
+24 -4
View File
@@ -1,12 +1,32 @@
{{ define "storage_stats" }}
<div class="storage-stats">
{{ if .Berth }}
<div class="stat-row">
<span class="stat-label">Berth:</span>
<span class="stat-value berth-badge berth-{{ .Berth }}">{{ .Berth }}</span>
</div>
{{ end }}
<div class="stat-row">
<span class="stat-label">Storage:</span>
<span class="stat-value">
{{ if .HasLimit }}
{{ .HumanSize }} / {{ .HumanLimit }}
{{ else }}
{{ .HumanSize }} <span class="unlimited-badge">Unlimited</span>
{{ end }}
</span>
</div>
{{ if .HasLimit }}
<div class="quota-progress">
<div class="progress-bar">
<div class="progress-fill {{ if ge .UsagePercent 95 }}progress-danger{{ else if ge .UsagePercent 80 }}progress-warning{{ else }}progress-ok{{ end }}" style="width: {{ .UsagePercent }}%"></div>
</div>
<span class="progress-text">{{ .UsagePercent }}% used</span>
</div>
{{ end }}
<div class="stat-row">
<span class="stat-label">Unique Blobs:</span>
<span class="stat-value">{{ .UniqueBlobs }}</span>
</div>
<div class="stat-row">
<span class="stat-label">Total Storage:</span>
<span class="stat-value">{{ .HumanSize }}</span>
</div>
</div>
{{ end }}
pkg/atproto/cbor_gen.go
+43 -1
View File
@@ -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":
pkg/atproto/lexicon.go
+2 -1
View File
@@ -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
pkg/auth/oauth/client.go
+2 -7
View File
@@ -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=*",
pkg/hold/oci/xrpc.go
+22 -1
View File
@@ -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
pkg/hold/oci/xrpc_test.go
+1 -1
View File
@@ -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
}
pkg/hold/pds/layer.go
+52
View File
@@ -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 ""
}
pkg/hold/pds/layer_test.go
+436
View File
@@ -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)
}
pkg/hold/pds/status_test.go
+2 -2
View File
@@ -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()
pkg/hold/pds/xrpc.go
+8 -4
View File
@@ -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)
pkg/hold/pds/xrpc_test.go
+6 -6
View File
@@ -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
}
pkg/hold/quota/config.go
+197
View File
@@ -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])
}
pkg/hold/quota/config_test.go
+275
View File
@@ -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)
}
}
quotas.yaml.example
+35
View File
@@ -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).