From c28ef3e82aaa0afd6f1adea8fa86f334cf55cb93 Mon Sep 17 00:00:00 2001 From: Matthew Meszaros Date: Fri, 28 Aug 2026 21:26:57 -0700 Subject: [PATCH] feat: keep a restricted or suspended workspace out of the premium warmup pool by checking the organization's risk posture before the stored warmup_pool_type in resolveWarmupPoolType (tasks and email services), which is never empty so the restricted branch below it was dead code, make the scheduler's recipient-capacity count risk-aware the same way, wire the org risk repository into the email service, and refuse the premium membership move in UpdateEmailAccountWarmupPoolType while the owning organization is restricted so a worker rebalance cannot readmit a risky tenant --- cmd/backend/main.go | 4 + internal/app/email/handler.go | 24 +++- internal/app/email/service.go | 17 +++ internal/app/email/warmup_pool_type_test.go | 59 ++++++++++ internal/repository/pg_worker.go | 12 +- internal/scheduler/warmup_pool_type_test.go | 66 +++++++++++ internal/scheduler/warmup_scheduler.go | 15 ++- internal/tasks/email_task.go | 10 +- internal/tasks/warmup_pool_type_test.go | 116 ++++++++++++++++++++ 9 files changed, 312 insertions(+), 11 deletions(-) create mode 100644 internal/app/email/warmup_pool_type_test.go create mode 100644 internal/scheduler/warmup_pool_type_test.go create mode 100644 internal/tasks/warmup_pool_type_test.go diff --git a/cmd/backend/main.go b/cmd/backend/main.go index e99dfa78..de2237d7 100644 --- a/cmd/backend/main.go +++ b/cmd/backend/main.go @@ -1148,6 +1148,10 @@ func main() { if instanceSettings != nil { emailService.WireSyncBudget(instanceSettings) } + // A restricted organization's mailboxes join the free pool on connect. + if aware, ok := emailService.(email.OrgRiskAware); ok { + aware.WireOrgRisk(orgRiskRepository) + } analyticsRepository := repository.NewAnalyticsRepository(primaryDB) emailAccountErrorRepository := repository.NewEmailAccountErrorRepository(primaryDB) analyticsService = analytics.NewService(analyticsRepository, emailRepostory, campaignRepostory, emailAccountErrorRepository, warmupRepository) diff --git a/internal/app/email/handler.go b/internal/app/email/handler.go index 962414c3..adc0fd70 100644 --- a/internal/app/email/handler.go +++ b/internal/app/email/handler.go @@ -359,18 +359,36 @@ func (s *emailService) canUseWarmupPool(ctx context.Context, account *models.Ema return err == nil && canWarmup } +// orgSuspendedOrRestricted reports whether the workspace's posture bars the +// paid warmup pool. Fails open, like every other risk read. +func (s *emailService) orgSuspendedOrRestricted(ctx context.Context, orgID uuid.UUID) bool { + if s.orgRiskRepo == nil { + return false + } + states, err := s.orgRiskRepo.GetOrgRiskStates(ctx, []uuid.UUID{orgID}) + if err != nil { + return false + } + return states[orgID].ForcesFreeWarmupPool() +} + func (s *emailService) resolveWarmupPoolType(ctx context.Context, account *models.Email) string { if account == nil { return "premium" } - if account.WarmupPoolType != "" { - return account.WarmupPoolType - } // No organization means no entitlement to check, so the mailbox gets the // lower-trust pool rather than defaulting into the paid one. if account.OrganizationID == nil { return "free" } + // A restricted organization warms in the free pool whatever it pays; the + // stored tier is checked after, since it is always set (issue #242). + if s.orgSuspendedOrRestricted(ctx, *account.OrganizationID) { + return "free" + } + if account.WarmupPoolType != "" { + return account.WarmupPoolType + } if s.featureGate != nil { isPaid, err := s.featureGate.IsPaidOrganization(ctx, *account.OrganizationID) if err == nil && !isPaid { diff --git a/internal/app/email/service.go b/internal/app/email/service.go index 13f74976..1b41df63 100644 --- a/internal/app/email/service.go +++ b/internal/app/email/service.go @@ -108,8 +108,25 @@ type emailService struct { // (email_account.connected, email_account.removed) are dispatched to // subscribed customer webhooks. webhookService webhook.Service + // orgRiskRepo bars a restricted organization from the paid warmup pool. + // Optional/nil-safe. + orgRiskRepo repository.OrgRiskRepository } +// WireOrgRisk attaches the organization risk posture. +func (s *emailService) WireOrgRisk(r repository.OrgRiskRepository) { + s.orgRiskRepo = r +} + +// OrgRiskAware is the optional capability the caller uses to attach org risk. +type OrgRiskAware interface { + WireOrgRisk(r repository.OrgRiskRepository) +} + +// The backend attaches the posture through a type assertion, so a silently +// unsatisfied interface would leave restricted workspaces in the paid pool. +var _ OrgRiskAware = (*emailService)(nil) + // SyncBudgetSource is the operator-editable sync fair-use section, satisfied // by instancesettings.Service. Injected post-construction; when unset the // loader ships compiled defaults. diff --git a/internal/app/email/warmup_pool_type_test.go b/internal/app/email/warmup_pool_type_test.go new file mode 100644 index 00000000..6128daa5 --- /dev/null +++ b/internal/app/email/warmup_pool_type_test.go @@ -0,0 +1,59 @@ +package email + +import ( + "context" + "errors" + "testing" + + "github.com/google/uuid" + + "github.com/warmbly/warmbly/internal/models" + "github.com/warmbly/warmbly/internal/repository" +) + +type poolTypeRiskRepo struct { + repository.OrgRiskRepository + states map[uuid.UUID]models.OrgRiskState + err error +} + +func (r *poolTypeRiskRepo) GetOrgRiskStates(_ context.Context, _ []uuid.UUID) (map[uuid.UUID]models.OrgRiskState, error) { + return r.states, r.err +} + +// A mailbox connected under a restricted workspace joins the free pool, not the +// tier its worker assignment wrote (issue #242). +func TestResolveWarmupPoolTypeDemotesARestrictedOrganization(t *testing.T) { + org := uuid.New() + account := &models.Email{ID: uuid.New(), OrganizationID: &org, Status: "active", WarmupPoolType: "premium"} + + for _, tc := range []struct { + name string + risk *poolTypeRiskRepo + wants string + }{ + {"restricted", &poolTypeRiskRepo{states: map[uuid.UUID]models.OrgRiskState{org: models.OrgRiskRestricted}}, "free"}, + {"suspended", &poolTypeRiskRepo{states: map[uuid.UUID]models.OrgRiskState{org: models.OrgRiskSuspended}}, "free"}, + {"trusted", &poolTypeRiskRepo{states: map[uuid.UUID]models.OrgRiskState{org: models.OrgRiskTrusted}}, "premium"}, + {"unknown org", &poolTypeRiskRepo{states: map[uuid.UUID]models.OrgRiskState{}}, "premium"}, + {"lookup failure fails open", &poolTypeRiskRepo{err: errors.New("connection reset by peer")}, "premium"}, + } { + t.Run(tc.name, func(t *testing.T) { + svc := &emailService{orgRiskRepo: tc.risk} + if got := svc.resolveWarmupPoolType(context.Background(), account); got != tc.wants { + t.Fatalf("resolved %q, want %q", got, tc.wants) + } + }) + } +} + +// Without the risk repository wired (jobs, tests) the stored tier stands. +func TestResolveWarmupPoolTypeWithoutRiskUsesTheStoredTier(t *testing.T) { + org := uuid.New() + svc := &emailService{} + account := &models.Email{ID: uuid.New(), OrganizationID: &org, Status: "active", WarmupPoolType: "premium"} + + if got := svc.resolveWarmupPoolType(context.Background(), account); got != "premium" { + t.Fatalf("resolved %q, want premium", got) + } +} diff --git a/internal/repository/pg_worker.go b/internal/repository/pg_worker.go index 81ae7d55..535cdfbc 100644 --- a/internal/repository/pg_worker.go +++ b/internal/repository/pg_worker.go @@ -496,6 +496,9 @@ func (r *workerRepository) ClearEmailAccountWorker(ctx context.Context, emailAcc // UpdateEmailAccountWarmupPoolType writes the tier and moves the mailbox's pool membership to // match in one transaction: they record the same fact, and updating only the column left // downgraded mailboxes in the premium pool (issue #211). A mailbox in no pool stays in none. +// The membership move into premium is refused while the owning organization is restricted or +// suspended, so a worker rebalance cannot readmit a risky tenant to the paid pool (issue #242). +// The tier column is still written: it records what the workspace pays for, not where it warms. func (r *workerRepository) UpdateEmailAccountWarmupPoolType(ctx context.Context, emailAccountID uuid.UUID, poolType string) error { tx, err := r.db.Begin(ctx) if err != nil { @@ -515,7 +518,14 @@ func (r *workerRepository) UpdateEmailAccountWarmupPoolType(ctx context.Context, FROM warmup_pools wp WHERE wp.pool_type = $1::warmup_pool_type AND wpp.email_account_id = $2::uuid - AND wpp.pool_id <> wp.id`, + AND wpp.pool_id <> wp.id + AND ($1::text <> 'premium' OR NOT EXISTS ( + SELECT 1 + FROM email_accounts ea + JOIN organizations o ON o.id = ea.organization_id + WHERE ea.id = $2::uuid + AND o.risk_state IN ('restricted', 'suspended') + ))`, poolType, emailAccountID); err != nil { return err } diff --git a/internal/scheduler/warmup_pool_type_test.go b/internal/scheduler/warmup_pool_type_test.go new file mode 100644 index 00000000..2cfabcda --- /dev/null +++ b/internal/scheduler/warmup_pool_type_test.go @@ -0,0 +1,66 @@ +package scheduler + +import ( + "context" + "errors" + "testing" + + "github.com/google/uuid" + + "github.com/warmbly/warmbly/internal/models" + "github.com/warmbly/warmbly/internal/repository" +) + +type poolTypeRiskRepo struct { + repository.OrgRiskRepository + states map[uuid.UUID]models.OrgRiskState + err error +} + +func (r *poolTypeRiskRepo) GetOrgRiskStates(_ context.Context, _ []uuid.UUID) (map[uuid.UUID]models.OrgRiskState, error) { + return r.states, r.err +} + +// The recipient-capacity count has to be taken against the pool the mailbox +// will actually send into, or a restricted mailbox sizes its day off the +// premium pool it is no longer in (issue #242). +func TestWarmupPoolTypeForAccountFollowsTheOrganizationsPosture(t *testing.T) { + org := uuid.New() + account := &models.Email{ID: uuid.New(), OrganizationID: &org, WarmupPoolType: "premium"} + + for _, tc := range []struct { + name string + risk *poolTypeRiskRepo + want string + }{ + {"restricted", &poolTypeRiskRepo{states: map[uuid.UUID]models.OrgRiskState{org: models.OrgRiskRestricted}}, "free"}, + {"suspended", &poolTypeRiskRepo{states: map[uuid.UUID]models.OrgRiskState{org: models.OrgRiskSuspended}}, "free"}, + {"watch", &poolTypeRiskRepo{states: map[uuid.UUID]models.OrgRiskState{org: models.OrgRiskWatch}}, "premium"}, + {"trusted", &poolTypeRiskRepo{states: map[uuid.UUID]models.OrgRiskState{org: models.OrgRiskTrusted}}, "premium"}, + {"lookup failure fails open", &poolTypeRiskRepo{err: errors.New("connection reset by peer")}, "premium"}, + } { + t.Run(tc.name, func(t *testing.T) { + s := &schedulerService{orgRiskRepo: tc.risk} + if got := s.warmupPoolTypeForAccount(context.Background(), account); got != tc.want { + t.Fatalf("resolved %q, want %q", got, tc.want) + } + }) + } +} + +// Without the risk repository wired the stored tier stands, and a mailbox with +// no tier recorded keeps the historical premium default. +func TestWarmupPoolTypeForAccountWithoutRiskWired(t *testing.T) { + s := &schedulerService{} + org := uuid.New() + + if got := s.warmupPoolTypeForAccount(context.Background(), &models.Email{OrganizationID: &org, WarmupPoolType: "free"}); got != "free" { + t.Fatalf("resolved %q, want free", got) + } + if got := s.warmupPoolTypeForAccount(context.Background(), &models.Email{OrganizationID: &org}); got != "premium" { + t.Fatalf("resolved %q, want premium", got) + } + if got := s.warmupPoolTypeForAccount(context.Background(), nil); got != "premium" { + t.Fatalf("nil account resolved %q, want premium", got) + } +} diff --git a/internal/scheduler/warmup_scheduler.go b/internal/scheduler/warmup_scheduler.go index 20a6b8fc..24c815db 100644 --- a/internal/scheduler/warmup_scheduler.go +++ b/internal/scheduler/warmup_scheduler.go @@ -47,8 +47,17 @@ func adjustmentFor(state models.WarmupHealthState) healthAdjustment { } } -func warmupPoolTypeForAccount(account *models.Email) string { - if account != nil && account.WarmupPoolType != "" { +// warmupPoolTypeForAccount is the pool the mailbox actually warms in. A +// restricted organization is held in the free pool whatever tier its mailbox +// carries, matching the task service's resolution (issue #242). +func (s *schedulerService) warmupPoolTypeForAccount(ctx context.Context, account *models.Email) string { + if account == nil { + return "premium" + } + if s.orgRiskState(ctx, account.OrganizationID).ForcesFreeWarmupPool() { + return "free" + } + if account.WarmupPoolType != "" { return account.WarmupPoolType } return "premium" @@ -232,7 +241,7 @@ func (s *schedulerService) CalculateNextWarmupTime(ctx context.Context, accountI // here, so operators can add inbound capacity without making those // mailboxes warmup senders. if s.warmupRepo != nil { - eligibleRecipients, err := s.warmupRepo.CountEligibleRecipients(ctx, warmupPoolTypeForAccount(account), accountID) + eligibleRecipients, err := s.warmupRepo.CountEligibleRecipients(ctx, s.warmupPoolTypeForAccount(ctx, account), accountID) if err == nil { if eligibleRecipients <= 0 { return recipientRecheckTime(), nil diff --git a/internal/tasks/email_task.go b/internal/tasks/email_task.go index dd4e788c..aa699873 100644 --- a/internal/tasks/email_task.go +++ b/internal/tasks/email_task.go @@ -739,9 +739,6 @@ func (s *tasksService) resolveWarmupPoolType(ctx context.Context, account *Email if account == nil { return "premium" } - if account.WarmupPoolType != "" { - return account.WarmupPoolType - } // No organization means no entitlement to check, so the mailbox gets the // lower-trust pool rather than defaulting into the paid one. if account.OrganizationID == nil { @@ -749,10 +746,15 @@ func (s *tasksService) resolveWarmupPoolType(ctx context.Context, account *Email } // A restricted organization leaves the paid pool whatever it pays: the // shared reputation paying customers depend on is not for spending on a - // risky tenant. + // risky tenant. Checked before the stored tier, which is always set + // (schema default 'free', worker assignment writes it) and so would + // otherwise short-circuit this gate (issue #242). if s.orgSuspendedOrRestricted(ctx, *account.OrganizationID) { return "free" } + if account.WarmupPoolType != "" { + return account.WarmupPoolType + } if s.featureGate != nil { isPaid, xerr := s.featureGate.IsPaidOrganization(ctx, *account.OrganizationID) if xerr == nil && !isPaid { diff --git a/internal/tasks/warmup_pool_type_test.go b/internal/tasks/warmup_pool_type_test.go new file mode 100644 index 00000000..791a3d82 --- /dev/null +++ b/internal/tasks/warmup_pool_type_test.go @@ -0,0 +1,116 @@ +package tasks + +import ( + "context" + "errors" + "testing" + + "github.com/google/uuid" + + "github.com/warmbly/warmbly/internal/errx" + "github.com/warmbly/warmbly/internal/models" + "github.com/warmbly/warmbly/internal/repository" +) + +type poolTypeRiskRepo struct { + repository.OrgRiskRepository + states map[uuid.UUID]models.OrgRiskState + err error +} + +func (r *poolTypeRiskRepo) GetOrgRiskStates(_ context.Context, _ []uuid.UUID) (map[uuid.UUID]models.OrgRiskState, error) { + return r.states, r.err +} + +type poolTypeGate struct { + poolReconcileGate + paid map[uuid.UUID]bool +} + +func (g *poolTypeGate) IsPaidOrganization(_ context.Context, orgID uuid.UUID) (bool, *errx.Error) { + return g.paid[orgID], nil +} + +func poolTypeAccount(orgID uuid.UUID, tier string) *Email { + return &Email{ID: uuid.New(), OrganizationID: &orgID, Status: "active", WarmupPoolType: tier} +} + +// The headline defect (issue #242): the stored tier is always set, so checking +// it first meant a restricted organization never left the premium pool. +func TestResolveWarmupPoolTypeDemotesARestrictedOrganizationWhateverItsTier(t *testing.T) { + for _, state := range []models.OrgRiskState{models.OrgRiskRestricted, models.OrgRiskSuspended} { + t.Run(string(state), func(t *testing.T) { + org := uuid.New() + svc := &tasksService{orgRiskRepo: &poolTypeRiskRepo{states: map[uuid.UUID]models.OrgRiskState{org: state}}} + + if got := svc.resolveWarmupPoolType(context.Background(), poolTypeAccount(org, "premium")); got != "free" { + t.Fatalf("resolved %q for a %s organization, want free", got, state) + } + }) + } +} + +// Lifting the restriction hands the mailbox its paid tier back on the next resolution. +func TestResolveWarmupPoolTypeKeepsTheStoredTierForATrustedOrganization(t *testing.T) { + org := uuid.New() + svc := &tasksService{orgRiskRepo: &poolTypeRiskRepo{states: map[uuid.UUID]models.OrgRiskState{org: models.OrgRiskWatch}}} + + if got := svc.resolveWarmupPoolType(context.Background(), poolTypeAccount(org, "premium")); got != "premium" { + t.Fatalf("resolved %q, want the stored premium tier", got) + } + if got := svc.resolveWarmupPoolType(context.Background(), poolTypeAccount(org, "free")); got != "free" { + t.Fatalf("resolved %q, want the stored free tier", got) + } +} + +// An unreadable posture must not demote: the risk read fails open everywhere else too. +func TestResolveWarmupPoolTypeFailsOpenWhenRiskCannotBeRead(t *testing.T) { + org := uuid.New() + svc := &tasksService{orgRiskRepo: &poolTypeRiskRepo{err: errors.New("connection reset by peer")}} + + if got := svc.resolveWarmupPoolType(context.Background(), poolTypeAccount(org, "premium")); got != "premium" { + t.Fatalf("resolved %q on a risk lookup failure, want the stored tier", got) + } +} + +// With no stored tier the subscription decides, and no organization means free. +func TestResolveWarmupPoolTypeFallsBackToEntitlementWithoutAStoredTier(t *testing.T) { + paid, trial := uuid.New(), uuid.New() + svc := &tasksService{featureGate: &poolTypeGate{paid: map[uuid.UUID]bool{paid: true, trial: false}}} + + if got := svc.resolveWarmupPoolType(context.Background(), poolTypeAccount(paid, "")); got != "premium" { + t.Fatalf("paid organization resolved %q, want premium", got) + } + if got := svc.resolveWarmupPoolType(context.Background(), poolTypeAccount(trial, "")); got != "free" { + t.Fatalf("trial organization resolved %q, want free", got) + } + if got := svc.resolveWarmupPoolType(context.Background(), &Email{ID: uuid.New(), WarmupPoolType: "premium"}); got != "free" { + t.Fatalf("mailbox with no organization resolved %q, want free", got) + } +} + +// The reconciler is what moves a mailbox that is not actively warming (a +// recipient-only member, say), so restriction has to reach it there too. +func TestPoolReconcileMovesARestrictedOrganizationToTheFreePool(t *testing.T) { + f := newPoolReconcileFixture() + restricted := f.mailbox(f.orgPaid, "premium", "premium") + participants := []uuid.UUID{restricted} + f.svc = &tasksService{ + warmupRepo: &poolReconcileWarmupRepo{participants: participants}, + emailRepo: f.emails, + featureGate: f.gate, + warmupHealth: f.warmup, + orgRiskRepo: &poolTypeRiskRepo{states: map[uuid.UUID]models.OrgRiskState{f.orgPaid: models.OrgRiskRestricted}}, + } + + moved, removed, err := f.svc.ReconcileWarmupPoolMembership(context.Background()) + if err != nil { + t.Fatalf("reconcile: %v", err) + } + if moved != 1 || removed != 0 { + t.Fatalf("moved %d removed %d, want 1 and 0", moved, removed) + } + if f.warmup.moved[restricted] != "free" { + t.Fatalf("moved to %q, want free", f.warmup.moved[restricted]) + } +}