From 765fa46efd8180dea7706dc2ea9a69705d9be2d5 Mon Sep 17 00:00:00 2001 From: Matthew Meszaros Date: Thu, 8 Oct 2026 04:08:09 +0000 Subject: [PATCH] feat: authorize relay-rewritten warmup receipts with recipient-bound send proof and serialize token engagement Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- docs/content/docs/guides/warmup.mdx | 2 + internal/app/consumer/event_new_email.go | 28 +++- .../consumer/warmup_verification_live_test.go | 157 ++++++++++++++++++ internal/repository/diagnostic_auth.go | 8 +- .../repository/diagnostic_auth_live_test.go | 46 ++++- internal/repository/pg_warmup.go | 69 +++++++- .../warmup_analytics_scope_live_test.go | 7 +- internal/repository/warmup_lineage.go | 45 +++-- .../repository/warmup_lineage_live_test.go | 125 +++++++++++++- 9 files changed, 458 insertions(+), 29 deletions(-) diff --git a/docs/content/docs/guides/warmup.mdx b/docs/content/docs/guides/warmup.mdx index 5a1e5f8cb..1a0aa22c5 100644 --- a/docs/content/docs/guides/warmup.mdx +++ b/docs/content/docs/guides/warmup.mdx @@ -21,6 +21,8 @@ Diagnostics carry task/sender/recipient-bound application tokens and disclose au Providers may strip headers or replace identifiers. Exact authorized lineage and send-result identifiers correlate diagnostics; missing or ambiguous lineage stays closed. Tokens identify application-owned traffic, not cryptographic authentication. +When an SMTP relay rewrites a Message-ID, a valid recipient-bound verification token, the expected sender address and a confirmed send can authorize one received copy. Replies and message-specific DKIM verification use that copy's received identifier. Headerless arrivals still require an exact confirmed identifier for new diagnostic lineage; a matching subject or mailbox pair alone does not authorize a reply. Arrivals held before send confirmation or during a verification outage are retained and retried automatically while their token remains valid. + New mailboxes default off. Independent send/receive controls and shared daily/rolling limits are on the Warmup tab. Start/resume enables diagnostic sending, Pause preserves separately authorized receiving, and Stop disables both while retaining ramp history and holds. Recipient-only cannot reply. Existing NULL stays legacy until explicit action; unrelated edits do not relabel it. New controls cannot re-enable legacy synthetic behavior. Cloud controls require updated Cloud replicas. See [send safety](/development/send-safety/). ### Where the content comes from diff --git a/internal/app/consumer/event_new_email.go b/internal/app/consumer/event_new_email.go index e8f856da3..674a49dfb 100644 --- a/internal/app/consumer/event_new_email.go +++ b/internal/app/consumer/event_new_email.go @@ -245,12 +245,7 @@ func (s *JobsService) handleWarmupEmail(ctx context.Context, e *models.JobEventN return token.SenderAccountID == e.Message.EmailID, nil } - if token.ConsumedAt == nil && token.ExpiresAt.After(time.Now()) { - if err := s.acceptWarmupEmail(ctx, e, token); err != nil { - return false, err - } - } - return true, nil + return s.processWarmupReceipt(ctx, e, token) } // handleUnmarkedWarmupEmail verifies warmup mail that arrived without its @@ -279,7 +274,24 @@ func (s *JobsService) handleUnmarkedWarmupEmail(ctx context.Context, e *models.J Str("token", token.Token.String()). Str("email_account_id", e.Message.EmailID.String()). Msg("verified warmup mail that arrived without its verify header") - return true, s.acceptWarmupEmail(ctx, e, token) + return s.processWarmupReceipt(ctx, e, token) +} + +func warmupReceiptProof(message *models.EmailMessageStoreData) repository.WarmupReceiptProof { + marker, _ := uuid.Parse(warmupTokenFromMessage(message)) + return repository.WarmupReceiptProof{Token: marker, SenderAddress: firstSenderAddress(message.FromAddr), MessageID: message.MessageID} +} + +func (s *JobsService) processWarmupReceipt(ctx context.Context, e *models.JobEventNewEmail, token *models.WarmupToken) (bool, error) { + if processor, ok := s.WarmupRepo.(repository.WarmupReceiptProcessor); ok { + return processor.ProcessWarmupReceipt(ctx, token.Token, e.Message.EmailID, warmupReceiptProof(e.Message), func(current *models.WarmupToken) error { + return s.acceptWarmupEmail(ctx, e, current) + }) + } + if token.ConsumedAt == nil && token.ExpiresAt.After(time.Now()) { + return true, s.acceptWarmupEmail(ctx, e, token) + } + return true, nil } const cloudWarmupCheckTimeout = 5 * time.Second @@ -497,7 +509,7 @@ func (s *JobsService) acceptWarmupEmail(ctx context.Context, e *models.JobEventN } } if lineage, ok := s.TaskRepo.(repository.WarmupLineageRepository); ok { - if err := lineage.RecordVerifiedWarmupParent(ctx, token.TaskID, e.Message.EmailID, e.Message.ID, e.Message.ThreadID); err != nil { + if err := lineage.RecordVerifiedWarmupParent(ctx, token.TaskID, e.Message.EmailID, e.Message.ID, e.Message.ThreadID, warmupReceiptProof(e.Message)); err != nil { return err } } diff --git a/internal/app/consumer/warmup_verification_live_test.go b/internal/app/consumer/warmup_verification_live_test.go index 7a8569412..ddc49da4b 100644 --- a/internal/app/consumer/warmup_verification_live_test.go +++ b/internal/app/consumer/warmup_verification_live_test.go @@ -91,6 +91,7 @@ func newWarmupFixture(t *testing.T, handle *db.DB) *warmupFixture { {`DELETE FROM warmup_received WHERE sender_account_id = $1`, f.sender}, {`DELETE FROM warmup_tokens WHERE sender_account_id = $1`, f.sender}, {`DELETE FROM tasks WHERE email_account_id = $1`, f.sender}, + {`DELETE FROM tasks WHERE email_account_id = $1`, f.partner}, {`DELETE FROM email_accounts WHERE id = $1`, f.sender}, {`DELETE FROM email_accounts WHERE id = $1`, f.partner}, {`DELETE FROM organizations WHERE id = $1`, f.senderOrg}, @@ -162,6 +163,162 @@ func liveWarmupService(t *testing.T) (*JobsService, *db.DB) { }, handle } +func TestLiveWarmupLineageRelayReceiptRecoveryAndConcurrentReplay(t *testing.T) { + s, handle := liveWarmupService(t) + ctx := t.Context() + f := newWarmupFixture(t, handle) + token := f.mintToken(t, s.WarmupRepo) + exec := func(query string, args ...any) { + t.Helper() + if _, err := handle.Exec(ctx, query, args...); err != nil { + t.Fatal(err) + } + } + exec(`INSERT INTO warmup_tasks(task_id,lineage_version,subject,max_turns)VALUES($1,1,$2,3)`, f.task, f.subject) + exec(`UPDATE warmup_tokens SET conversation_id=$2 WHERE token=$1`, token, uuid.New()) + exec(`UPDATE email_accounts SET warmup=NOW(),warmup_reply_rate=100,warmup_days=127 WHERE id=$1`, f.partner) + e := f.arrival("", []string{config.WarmupVerifyHeader + ":" + token.String()}) + e.Message.ThreadID = "recipient-provider-thread" + if err := s.HandleNewEmail(ctx, e); err != nil { + t.Fatal(err) + } + var pending, receipts int + if err := handle.QueryRow(ctx, `SELECT COUNT(*) FROM unibox_pending_emails WHERE id=$1`, e.Message.ID).Scan(&pending); err != nil || pending != 1 { + t.Fatalf("unstamped arrival not retained: %d %v", pending, err) + } + if err := handle.QueryRow(ctx, `SELECT COUNT(*) FROM warmup_received WHERE sender_account_id=$1`, f.sender).Scan(&receipts); err != nil || receipts != 0 { + t.Fatalf("unconfirmed send created receipt evidence: %d %v", receipts, err) + } + if err := s.WarmupRepo.RecordWarmupTokenDelivery(ctx, f.task, f.sentMsgID); err != nil { + t.Fatal(err) + } + exec(`UPDATE unibox_pending_emails SET retry_at=NOW() WHERE id=$1`, e.Message.ID) + if err := s.retryPendingWarmupVerification(ctx); err != nil { + t.Fatal("existing queued relay arrival did not recover", err) + } + var parent string + if err := handle.QueryRow(ctx, `SELECT parent_message_id FROM warmup_tasks WHERE parent_task_id=$1`, f.task).Scan(&parent); err != nil || parent != e.Message.MessageID { + t.Fatalf("reply parent=%q err=%v", parent, err) + } + stored, err := s.WarmupRepo.FindWarmupToken(ctx, token) + if err != nil || stored == nil || stored.ConsumedAt == nil { + t.Fatal("recovered token not consumed", stored, err) + } + consumed := *stored.ConsumedAt + var wg sync.WaitGroup + for range 12 { + wg.Add(1) + go func() { + defer wg.Done() + duplicate := f.arrival("", slices.Clone(e.Message.Flags)) + if err := s.HandleNewEmail(ctx, duplicate); err != nil { + t.Error(err) + } + }() + } + wg.Wait() + exec(`UPDATE unibox_pending_emails SET retry_at=NOW() WHERE email_account_id=$1`, f.partner) + if err := s.retryPendingWarmupVerification(ctx); err != nil { + t.Fatal("concurrent redeliveries did not drain", err) + } + if err := handle.QueryRow(ctx, `SELECT COUNT(*) FROM warmup_received WHERE sender_account_id=$1`, f.sender).Scan(&receipts); err != nil || receipts != 1 { + t.Fatalf("replay created additional receipt evidence: %d %v", receipts, err) + } + stored, err = s.WarmupRepo.FindWarmupToken(ctx, token) + if err != nil || !stored.ConsumedAt.Equal(consumed) { + t.Fatal("replay changed consumption timestamp", stored, err) + } + if err := handle.QueryRow(ctx, `SELECT COUNT(*) FROM unibox_pending_emails WHERE email_account_id=$1`, f.partner).Scan(&pending); err != nil || pending != 0 { + t.Fatalf("recovered arrival still pending: %d %v", pending, err) + } +} + +func TestLiveWarmupLineageConcurrentInitialRelayArrivalsEngageOnce(t *testing.T) { + s, handle := liveWarmupService(t) + ctx := t.Context() + f := newWarmupFixture(t, handle) + token := f.mintToken(t, s.WarmupRepo) + if _, err := handle.Exec(ctx, `INSERT INTO warmup_tasks(task_id,lineage_version,subject,max_turns)VALUES($1,1,$2,1)`, f.task, f.subject); err != nil { + t.Fatal(err) + } + if err := s.WarmupRepo.RecordWarmupTokenDelivery(ctx, f.task, f.sentMsgID); err != nil { + t.Fatal(err) + } + var wg sync.WaitGroup + for range 16 { + wg.Add(1) + go func() { + defer wg.Done() + e := f.arrival("<"+uuid.NewString()+"@mailjet.com>", []string{config.WarmupVerifyHeader + ":" + token.String()}) + if err := s.HandleNewEmail(ctx, e); err != nil { + t.Error(err) + } + }() + } + wg.Wait() + if _, err := handle.Exec(ctx, `UPDATE unibox_pending_emails SET retry_at=NOW() WHERE email_account_id=$1`, f.partner); err != nil { + t.Fatal(err) + } + if err := s.retryPendingWarmupVerification(ctx); err != nil { + t.Fatal(err) + } + var receipts int + if err := handle.QueryRow(ctx, `SELECT COUNT(*) FROM warmup_received WHERE sender_account_id=$1`, f.sender).Scan(&receipts); err != nil || receipts != 1 { + t.Fatalf("concurrent token claim created %d receipts: %v", receipts, err) + } +} + +func TestLiveWarmupLineageConsumerReceiptGuards(t *testing.T) { + for _, name := range []string{"exact headerless", "headerless rewrite", "forwarded marker", "foreign mailbox", "expired", "retired"} { + t.Run(name, func(t *testing.T) { + s, handle := liveWarmupService(t) + ctx := t.Context() + f := newWarmupFixture(t, handle) + token := f.mintToken(t, s.WarmupRepo) + if _, err := handle.Exec(ctx, `INSERT INTO warmup_tasks(task_id,lineage_version,subject,max_turns)VALUES($1,1,$2,1)`, f.task, f.subject); err != nil { + t.Fatal(err) + } + if err := s.WarmupRepo.RecordWarmupTokenDelivery(ctx, f.task, f.sentMsgID); err != nil { + t.Fatal(err) + } + e := f.arrival("", []string{config.WarmupVerifyHeader + ":" + token.String()}) + switch name { + case "exact headerless": + e.Message.MessageID, e.Message.Flags = f.sentMsgID, nil + case "headerless rewrite": + e.Message.Flags = nil + case "forwarded marker": + e.Message.FromAddr = []string{"forwarder@example.test"} + case "foreign mailbox": + e.Message.EmailID, e.UserID = f.sender, f.senderUser + case "expired": + if _, err := handle.Exec(ctx, `UPDATE warmup_tokens SET expires_at=NOW()-INTERVAL '1 second' WHERE token=$1`, token); err != nil { + t.Fatal(err) + } + case "retired": + if _, err := handle.Exec(ctx, `UPDATE warmup_tokens SET sent_retired_at=NOW() WHERE token=$1`, token); err != nil { + t.Fatal(err) + } + } + if err := s.HandleNewEmail(ctx, e); err != nil { + t.Fatal(err) + } + want := 0 + if name == "exact headerless" { + want = 1 + } + var receipts int + if err := handle.QueryRow(ctx, `SELECT COUNT(*) FROM warmup_received WHERE sender_account_id=$1`, f.sender).Scan(&receipts); err != nil || receipts != want { + t.Fatalf("receipt evidence=%d want=%d err=%v", receipts, want, err) + } + stored, err := s.WarmupRepo.FindWarmupToken(ctx, token) + if err != nil || stored == nil || (stored.ConsumedAt != nil) != (want == 1) { + t.Fatalf("unexpected token consumption: %+v %v", stored, err) + } + }) + } +} + // The regression: a warmup email that arrives with no verify header must still // be verified, consumed and kept out of the recipient's unibox. func TestLiveWarmupVerifiesByDeliveredMessageIDWhenTheHeaderIsStripped(t *testing.T) { diff --git a/internal/repository/diagnostic_auth.go b/internal/repository/diagnostic_auth.go index 00d374cae..732177afe 100644 --- a/internal/repository/diagnostic_auth.go +++ b/internal/repository/diagnostic_auth.go @@ -29,7 +29,7 @@ func (r *taskRepository) DiagnosticAuth(ctx context.Context, req models.Diagnost defer tx.Rollback(ctx) grant := &models.DiagnosticAuthGrant{} var task uuid.UUID - err = tx.QueryRow(ctx, `SELECT wt.task_id,sender.email,ea.email,wt.sent_message_id + err = tx.QueryRow(ctx, `SELECT wt.task_id,sender.email,ea.email,$4::text FROM warmup_tokens wt JOIN warmup_tasks w ON w.task_id=wt.task_id JOIN email_accounts sender ON sender.id=wt.sender_account_id JOIN email_accounts ea ON ea.id=wt.recipient_account_id @@ -37,7 +37,11 @@ func (r *taskRepository) DiagnosticAuth(ctx context.Context, req models.Diagnost JOIN warmup_pools wp ON wp.id=wpp.pool_id JOIN fleet_nodes n ON n.id=ea.worker_id WHERE wt.token=$1 AND ea.id=$2 AND ea.worker_id=$3 AND n.role='worker' AND n.active AND n.warmup_send_protocol>=2 AND n.last_seen_at>NOW()-INTERVAL '10 minutes' AND w.lineage_version=1 AND w.scenario_version='diagnostic-v1' AND w.rendering_version='canonical-v1' - AND wt.expires_at>NOW() AND wt.sent_message_id<>'' AND btrim(wt.sent_message_id,'<>')=$4 + AND wt.expires_at>NOW() AND wt.sent_retired_at IS NULL AND wt.sent_message_id<>'' + AND EXISTS(SELECT 1 FROM tasks sent WHERE sent.id=wt.task_id AND sent.status='completed' AND sent.email_account_id=wt.sender_account_id) + AND (btrim(wt.sent_message_id,'<>')=$4 OR EXISTS(SELECT 1 FROM warmup_received wr + WHERE wr.task_id=wt.task_id AND wr.email_account_id=ea.id AND wr.sender_account_id=wt.sender_account_id + AND wr.retired_at IS NULL AND btrim(wr.message_id,'<>')=$4)) AND sender.status='active' AND (sender.test_mode IS NULL OR sender.test_mode='legacy' OR sender.test_mode='diagnostic' AND sender.test_send_enabled) AND NOT EXISTS(SELECT 1 FROM cloud_link_mailboxes clm WHERE clm.email_account_id=ea.id) AND `+partnerEligibleSQL, req.Token, req.MailboxID, req.WorkerID, strings.Trim(req.MessageID, "<> \t")).Scan(&task, &grant.From, &grant.To, &grant.MessageID) diff --git a/internal/repository/diagnostic_auth_live_test.go b/internal/repository/diagnostic_auth_live_test.go index 0d8c3434f..0ea82bece 100644 --- a/internal/repository/diagnostic_auth_live_test.go +++ b/internal/repository/diagnostic_auth_live_test.go @@ -93,7 +93,7 @@ func TestLiveDiagnosticAuthBindsWorkerContextAndRetainsOnlyMinimalProof(t *testi } receipt := uuid.New() exec(`INSERT INTO warmup_received(email_account_id,internal_id,sender_account_id,message_id,evidence)VALUES($1,$2,$3,'probe@example.test','{"spf":"unknown","dkim":"unknown","dmarc":"unknown","tls":"unknown","alignment":"unknown","trust":"unverified_headers"}')`, f.recipient, receipt, f.sender) - if err = r.RecordVerifiedWarmupParent(ctx, f.task, f.recipient, receipt, ""); err != nil { + if err = r.RecordVerifiedWarmupParent(ctx, f.task, f.recipient, receipt, "", WarmupReceiptProof{SenderAddress: f.senderTo}); err != nil { t.Fatal(err) } var evidence models.ReceivedEvidence @@ -108,3 +108,47 @@ func TestLiveDiagnosticAuthBindsWorkerContextAndRetainsOnlyMinimalProof(t *testi t.Fatal("late exact receipt lost crypto proof or fabricated other auth", evidence) } } + +func TestLiveDiagnosticAuthRelayRewriteUsesOnlyBoundReceivedID(t *testing.T) { + f, r := lineageFixture(t) + ctx := t.Context() + worker, token, receipt := uuid.New(), uuid.New(), uuid.New() + exec := func(query string, args ...any) { + t.Helper() + if _, err := f.pool.Exec(ctx, query, args...); err != nil { + t.Fatal(err) + } + } + exec(`INSERT INTO fleet_nodes(id,role,active,last_seen_at,warmup_send_protocol)VALUES($1,'worker',true,NOW(),2)`, worker) + exec(`INSERT INTO workers(id)VALUES($1)`, worker) + t.Cleanup(func() { _, _ = f.pool.Exec(context.Background(), `DELETE FROM fleet_nodes WHERE id=$1`, worker) }) + exec(`UPDATE email_accounts SET worker_id=$1,provider=$3 WHERE id=$2`, worker, f.recipient, models.InboxProviderGoogle) + exec(`INSERT INTO warmup_pool_participants(pool_id,email_account_id)SELECT id,$1 FROM warmup_pools WHERE pool_type='free'`, f.recipient) + exec(`INSERT INTO warmup_tasks(task_id,lineage_version,subject,scenario_version,rendering_version,max_turns)VALUES($1,1,'Diagnostic','diagnostic-v1','canonical-v1',1)`, f.task) + exec(`INSERT INTO warmup_tokens(token,task_id,sender_account_id,recipient_account_id,sent_message_id)VALUES($1,$2,$3,$4,'')`, token, f.task, f.sender, f.recipient) + exec(`INSERT INTO warmup_received(email_account_id,internal_id,sender_account_id,message_id)VALUES($1,$2,$3,'received@mailjet.com')`, f.recipient, receipt, f.sender) + request := models.DiagnosticAuthRequest{Token: token, MailboxID: f.recipient, WorkerID: worker, MessageID: ""} + if grant, err := r.DiagnosticAuth(ctx, request); err == nil || grant != nil { + t.Fatal("unbound rewritten ID authorized raw retrieval", grant, err) + } + if err := r.RecordVerifiedWarmupParent(ctx, f.task, f.recipient, receipt, "", WarmupReceiptProof{Token: token, SenderAddress: f.senderTo}); err != nil { + t.Fatal(err) + } + grant, err := r.DiagnosticAuth(ctx, request) + if err != nil || grant == nil || grant.MessageID != "received@mailjet.com" { + t.Fatalf("rewritten grant=%+v err=%v", grant, err) + } + request.Nonce = grant.Nonce + request.Result = &models.DiagnosticDKIMResult{DKIM: "pass", SigningDomain: "test.local", Alignment: "pass", Verifier: models.DiagnosticDKIMVerifier, ObservedAt: time.Now()} + if _, err := r.DiagnosticAuth(ctx, request); err != nil { + t.Fatal(err) + } + var dkim string + if err := f.pool.QueryRow(ctx, `SELECT evidence->>'dkim' FROM warmup_received WHERE email_account_id=$1 AND internal_id=$2`, f.recipient, receipt).Scan(&dkim); err != nil || dkim != "pass" { + t.Fatalf("received-ID proof not attached: %q %v", dkim, err) + } + request.Result, request.MessageID = nil, "unrelated@smtp-pulse.com" + if grant, err := r.DiagnosticAuth(ctx, request); err == nil || grant != nil { + t.Fatal("bound rewrite authorized a different ID", grant, err) + } +} diff --git a/internal/repository/pg_warmup.go b/internal/repository/pg_warmup.go index d40ef3670..61b082c73 100644 --- a/internal/repository/pg_warmup.go +++ b/internal/repository/pg_warmup.go @@ -6,6 +6,7 @@ import ( "errors" "strconv" "strings" + "sync" "time" "github.com/google/uuid" @@ -309,7 +310,9 @@ type WarmupRepository interface { } type warmupRepository struct { - db *pgxpool.Pool + db *pgxpool.Pool + receiptOnce sync.Once + receiptSlots chan struct{} } // NewWarmupRepository creates a new warmup repository @@ -1655,9 +1658,71 @@ func (r *warmupRepository) FindWarmupToken(ctx context.Context, tokenID uuid.UUI return scanWarmupToken(resultDB(ctx, r.db).QueryRow(ctx, query, tokenID)) } +// WarmupReceiptProcessor serializes recipient engagement across delivery retries. +type WarmupReceiptProcessor interface { + ProcessWarmupReceipt(context.Context, uuid.UUID, uuid.UUID, WarmupReceiptProof, func(*models.WarmupToken) error) (bool, error) +} + +var ErrWarmupReceiptBusy = errors.New("warmup receipt is already being processed") + +func (r *warmupRepository) ProcessWarmupReceipt(ctx context.Context, tokenID, recipient uuid.UUID, proof WarmupReceiptProof, accept func(*models.WarmupToken) error) (bool, error) { + if r.db.Config().MaxConns < 2 { + return false, errors.New("warmup receipt processing requires at least two database connections") + } + // Leave capacity for acceptance's repository calls while the receipt lock is held. + r.receiptOnce.Do(func() { r.receiptSlots = make(chan struct{}, r.db.Config().MaxConns/2) }) + select { + case r.receiptSlots <- struct{}{}: + defer func() { <-r.receiptSlots }() + default: + return false, ErrWarmupReceiptBusy + } + tx, err := r.db.Begin(ctx) + if err != nil { + return false, err + } + defer tx.Rollback(context.WithoutCancel(ctx)) + // The lock spans acceptance, including consumption, but not the callback's writes. + var locked bool + if err = tx.QueryRow(ctx, `SELECT pg_try_advisory_xact_lock(hashtext('warmup_receipt_' || $1::text))`, tokenID).Scan(&locked); err != nil { + return false, err + } + if !locked { + return false, ErrWarmupReceiptBusy + } + token, err := scanWarmupToken(tx.QueryRow(ctx, `SELECT `+warmupTokenColumns+` FROM warmup_tokens wt + WHERE wt.token=$1 AND wt.recipient_account_id=$2 + AND EXISTS(SELECT 1 FROM email_accounts sender WHERE sender.id=wt.sender_account_id AND lower(sender.email)=lower($3))`, + tokenID, recipient, strings.TrimSpace(proof.SenderAddress))) + if err != nil || token == nil { + return false, err + } + if token.ConsumedAt != nil || !token.ExpiresAt.After(time.Now()) { + return true, nil + } + var lineage, awaitingSend, retired bool + if err = tx.QueryRow(ctx, `SELECT EXISTS(SELECT 1 FROM warmup_tasks WHERE task_id=$1 AND lineage_version=1), + EXISTS(SELECT 1 FROM warmup_tasks w JOIN tasks sent ON sent.id=w.task_id + WHERE w.task_id=$1 AND w.lineage_version=1 AND (sent.status<>'completed' OR $2='')), + EXISTS(SELECT 1 FROM warmup_tokens WHERE token=$3 AND sent_retired_at IS NOT NULL)`, token.TaskID, token.SentMessageID, token.Token).Scan(&lineage, &awaitingSend, &retired); err != nil { + return false, err + } + if retired { + return true, nil + } + if awaitingSend { + return false, ErrWarmupDeliveryPending + } + if lineage && (strings.Trim(proof.MessageID, "<> \t") == "" || + proof.Token != token.Token && strings.Trim(token.SentMessageID, "<> \t") != strings.Trim(proof.MessageID, "<> \t")) { + return false, nil + } + return true, accept(token) +} + // ConsumeWarmupToken marks a warmup token as consumed func (r *warmupRepository) ConsumeWarmupToken(ctx context.Context, tokenID uuid.UUID) error { - query := `UPDATE warmup_tokens SET consumed_at = NOW() WHERE token = $1` + query := `UPDATE warmup_tokens SET consumed_at = NOW() WHERE token = $1 AND consumed_at IS NULL` _, err := resultDB(ctx, r.db).Exec(ctx, query, tokenID) return err } diff --git a/internal/repository/warmup_analytics_scope_live_test.go b/internal/repository/warmup_analytics_scope_live_test.go index 0942038ae..ea1ca17a2 100644 --- a/internal/repository/warmup_analytics_scope_live_test.go +++ b/internal/repository/warmup_analytics_scope_live_test.go @@ -49,6 +49,7 @@ func TestLiveWarmupAnalyticsAreOrganizationScoped(t *testing.T) { arg uuid.UUID }{ {`DELETE FROM warmup_received WHERE email_account_id IN (SELECT id FROM email_accounts WHERE organization_id = $1)`, orgID}, + {`DELETE FROM cloud_link WHERE organization_id = $1`, orgID}, {`DELETE FROM email_accounts WHERE organization_id = $1`, orgID}, {`DELETE FROM organizations WHERE id = $1`, orgID}, {`DELETE FROM users WHERE id = $1`, userID}, @@ -65,8 +66,10 @@ func TestLiveWarmupAnalyticsAreOrganizationScoped(t *testing.T) { instanceID, otherInstanceID, orgID, uuid.NewString(), uuid.NewString()) exec(`INSERT INTO pool_link_mailboxes (instance_id, remote_id, email_account_id) VALUES ($1, $3, $5), ($2, $4, $6)`, instanceID, otherInstanceID, remoteID, otherRemoteID, accountID, secondAccountID) - exec(`INSERT INTO cloud_link_mailboxes (email_account_id, remote_id, managed) VALUES ($1, $3, true), ($2, $4, false)`, - accountID, secondAccountID, remoteID, otherRemoteID) + cloudInstanceID := uuid.New() + exec(`INSERT INTO cloud_link(instance_id,organization_id,cloud_url,token) VALUES($1,$2,'https://cloud.test','fixture')`, cloudInstanceID, orgID) + exec(`INSERT INTO cloud_link_mailboxes (email_account_id, remote_id, managed, instance_id) VALUES ($1, $3, true, $5), ($2, $4, false, $5)`, + accountID, secondAccountID, remoteID, otherRemoteID, cloudInstanceID) t.Cleanup(func() { if _, err := pool.Exec(context.Background(), `DELETE FROM pool_link_instances WHERE id = ANY($1::uuid[])`, []uuid.UUID{instanceID, otherInstanceID}); err != nil { t.Errorf("instance cleanup: %v", err) diff --git a/internal/repository/warmup_lineage.go b/internal/repository/warmup_lineage.go index 2ce960354..08cb5ca52 100644 --- a/internal/repository/warmup_lineage.go +++ b/internal/repository/warmup_lineage.go @@ -3,6 +3,7 @@ package repository import ( "context" "errors" + "strings" "time" "github.com/google/uuid" @@ -12,7 +13,7 @@ import ( // WarmupLineageRepository keeps receipt authority separate from pair selection. type WarmupLineageRepository interface { - RecordVerifiedWarmupParent(context.Context, uuid.UUID, uuid.UUID, uuid.UUID, string) error + RecordVerifiedWarmupParent(context.Context, uuid.UUID, uuid.UUID, uuid.UUID, string, WarmupReceiptProof) error BindWarmupSuccessor(context.Context, uuid.UUID, uuid.UUID, uuid.UUID, time.Time) (bool, error) ExactWarmupParent(context.Context, uuid.UUID) (*WarmupReplyCandidate, error) WarmupExecutionContext(context.Context, uuid.UUID) (uuid.UUID, *WarmupReplyCandidate, error) @@ -39,24 +40,45 @@ type WarmupQueueRevision struct { PreviousHandle *string } -func (r *taskRepository) RecordVerifiedWarmupParent(ctx context.Context, parent, recipient, internalID uuid.UUID, threadID string) error { - tag, err := r.db.Exec(ctx, `UPDATE warmup_received wr SET task_id=$1, provider_thread_id=NULLIF($4,'') +type WarmupReceiptProof struct { + Token uuid.UUID + SenderAddress string + MessageID string +} + +func (r *taskRepository) RecordVerifiedWarmupParent(ctx context.Context, parent, recipient, internalID uuid.UUID, threadID string, proof WarmupReceiptProof) error { + tx, err := r.db.Begin(ctx) + if err != nil { + return err + } + defer tx.Rollback(ctx) + if _, err = tx.Exec(ctx, `SELECT pg_advisory_xact_lock(hashtext('warmup_parent_' || $1::text))`, parent); err != nil { + return err + } + tag, err := tx.Exec(ctx, `UPDATE warmup_received wr SET task_id=$1, provider_thread_id=NULLIF($4,'') WHERE wr.email_account_id=$2 AND wr.internal_id=$3 AND wr.task_id IS NULL - AND EXISTS (SELECT 1 FROM warmup_tokens wt JOIN warmup_tasks original ON original.task_id=wt.task_id WHERE wt.task_id=$1 + AND wr.message_id<>'' + AND NOT EXISTS(SELECT 1 FROM warmup_received prior WHERE prior.task_id=$1) + AND EXISTS (SELECT 1 FROM warmup_tokens wt JOIN warmup_tasks original ON original.task_id=wt.task_id + JOIN tasks sent ON sent.id=wt.task_id JOIN email_accounts sender ON sender.id=wt.sender_account_id WHERE wt.task_id=$1 AND original.lineage_version=1 AND wt.recipient_account_id=$2 AND wt.sender_account_id=wr.sender_account_id - AND wt.sent_message_id<>'' AND btrim(wt.sent_message_id,'<>')=btrim(wr.message_id,'<>'))`, parent, recipient, internalID, threadID) + AND sent.status='completed' AND sent.email_account_id=wt.sender_account_id + AND wt.expires_at>NOW() AND wt.consumed_at IS NULL AND wt.sent_retired_at IS NULL + AND lower(sender.email)=lower($6) AND wt.sent_message_id<>'' + AND (btrim(wt.sent_message_id,'<>')=btrim(wr.message_id,'<>') OR wt.token=$5))`, + parent, recipient, internalID, threadID, proof.Token, strings.TrimSpace(proof.SenderAddress)) if err != nil { return err } if tag.RowsAffected() == 0 { var same bool - err = r.db.QueryRow(ctx, `SELECT EXISTS(SELECT 1 FROM warmup_received WHERE email_account_id=$2 AND internal_id=$3 AND task_id=$1)`, parent, recipient, internalID).Scan(&same) + err = tx.QueryRow(ctx, `SELECT EXISTS(SELECT 1 FROM warmup_received WHERE email_account_id=$2 AND internal_id=$3 AND task_id=$1)`, parent, recipient, internalID).Scan(&same) if err != nil { return err } if !same { var legacy bool - if err := r.db.QueryRow(ctx, `SELECT NOT EXISTS(SELECT 1 FROM warmup_tasks WHERE task_id=$1 AND lineage_version=1)`, parent).Scan(&legacy); err != nil { + if err := tx.QueryRow(ctx, `SELECT NOT EXISTS(SELECT 1 FROM warmup_tasks WHERE task_id=$1 AND lineage_version=1)`, parent).Scan(&legacy); err != nil { return err } if legacy { @@ -65,7 +87,10 @@ func (r *taskRepository) RecordVerifiedWarmupParent(ctx context.Context, parent, return errors.New("verified warmup receipt context unavailable") } } - return attachDiagnosticAuth(ctx, r.db, parent, recipient) + if err := attachDiagnosticAuth(ctx, tx, parent, recipient); err != nil { + return err + } + return tx.Commit(ctx) } func (r *taskRepository) BindWarmupSuccessor(ctx context.Context, parent, recipient, internalID uuid.UUID, at time.Time) (bool, error) { @@ -110,7 +135,7 @@ func (r *taskRepository) BindWarmupSuccessor(ctx context.Context, parent, recipi AND original.lineage_version=1 AND original.subject<>'' AND wt.conversation_id IS NOT NULL AND wt.expires_at>NOW() AND wt.sent_retired_at IS NULL AND wt.conversation_turn+1'' - AND wt.sent_message_id<>'' AND btrim(wt.sent_message_id,'<>')=btrim(wr.message_id,'<>') + AND wt.sent_message_id<>'' AND wr.sender_account_id=wt.sender_account_id AND NOT EXISTS(SELECT 1 FROM warmup_tasks prior WHERE prior.parent_task_id=$2)`, id, parent, internalID, recipient) if err != nil { return false, err @@ -137,7 +162,7 @@ func (r *taskRepository) ExactWarmupParent(ctx context.Context, successor uuid.U AND wr.created_at>NOW()-INTERVAL '7 days' AND original.lineage_version=1 AND wt.expires_at>NOW() AND wt.sent_retired_at IS NULL AND wt.conversation_turn+1'' AND btrim(wt.sent_message_id,'<>')=btrim(wr.message_id,'<>')`, successor).Scan(&c.MessageID, &c.Subject, &c.ThreadID, &c.ConversationTheme, &c.ContentSource, &c.ConversationID, &c.ConversationTurn, &c.References, &c.ScenarioVersion, &c.RenderingVersion, &c.MaxTurns) + AND wt.sent_message_id<>'' AND wr.sender_account_id=wt.sender_account_id`, successor).Scan(&c.MessageID, &c.Subject, &c.ThreadID, &c.ConversationTheme, &c.ContentSource, &c.ConversationID, &c.ConversationTurn, &c.References, &c.ScenarioVersion, &c.RenderingVersion, &c.MaxTurns) if errors.Is(err, pgx.ErrNoRows) { return nil, nil } diff --git a/internal/repository/warmup_lineage_live_test.go b/internal/repository/warmup_lineage_live_test.go index e76d9aa63..7831930be 100644 --- a/internal/repository/warmup_lineage_live_test.go +++ b/internal/repository/warmup_lineage_live_test.go @@ -2,6 +2,7 @@ package repository import ( "context" + "errors" "net/url" "os" "sync" @@ -10,6 +11,7 @@ import ( "time" "github.com/google/uuid" + "github.com/jackc/pgx/v5/pgxpool" "github.com/warmbly/warmbly/internal/infrastructure/db" "github.com/warmbly/warmbly/internal/models" ) @@ -52,14 +54,14 @@ func TestLiveLineageBindsOnlyExactVerifiedParentAndOneConcurrentSuccessor(t *tes exec(`INSERT INTO warmup_tokens(token,task_id,sender_account_id,recipient_account_id,subject,content_source,conversation_id,conversation_turn)VALUES($1,$2,$3,$4,'Pinned diagnostic','static',$5,0)`, uuid.New(), f.task, f.sender, f.recipient, uuid.New()) receipt := uuid.New() exec(`INSERT INTO warmup_received(email_account_id,internal_id,sender_account_id,message_id)VALUES($1,$2,$3,'parent@example.test')`, f.recipient, receipt, f.sender) - if err := r.RecordVerifiedWarmupParent(ctx, f.task, f.recipient, receipt, "receiving-mailbox-thread"); err == nil { + if err := r.RecordVerifiedWarmupParent(ctx, f.task, f.recipient, receipt, "receiving-mailbox-thread", WarmupReceiptProof{SenderAddress: f.senderTo}); err == nil { t.Fatal("unstamped header claim became send authority") } exec(`UPDATE warmup_tokens SET sent_message_id='' WHERE task_id=$1`, f.task) - if err := r.RecordVerifiedWarmupParent(ctx, f.task, f.recipient, receipt, "receiving-mailbox-thread"); err != nil { + if err := r.RecordVerifiedWarmupParent(ctx, f.task, f.recipient, receipt, "receiving-mailbox-thread", WarmupReceiptProof{SenderAddress: f.senderTo}); err != nil { t.Fatal(err) } - if err := r.RecordVerifiedWarmupParent(ctx, f.task, f.recipient, receipt, "receiving-mailbox-thread"); err != nil { + if err := r.RecordVerifiedWarmupParent(ctx, f.task, f.recipient, receipt, "receiving-mailbox-thread", WarmupReceiptProof{SenderAddress: f.senderTo}); err != nil { t.Fatal(err) } var won atomic.Int32 @@ -114,7 +116,7 @@ func TestLiveLineageBindsOnlyExactVerifiedParentAndOneConcurrentSuccessor(t *tes t.Fatalf("unsafe parent continued: %+v %v", got, err) } } - if err := r.RecordVerifiedWarmupParent(ctx, uuid.New(), f.recipient, receipt, "foreign"); err == nil { + if err := r.RecordVerifiedWarmupParent(ctx, uuid.New(), f.recipient, receipt, "foreign", WarmupReceiptProof{SenderAddress: f.senderTo}); err == nil { var bound uuid.UUID if err := f.pool.QueryRow(ctx, `SELECT task_id FROM warmup_received WHERE internal_id=$1 AND email_account_id=$2`, receipt, f.recipient).Scan(&bound); err != nil || bound != f.task { t.Fatal("legacy or foreign task claimed receipt") @@ -188,6 +190,121 @@ func TestLiveLineageQueueRevisionsRejectStaleAcknowledgements(t *testing.T) { } } +func TestLiveLineageRelayRewriteRequiresExplicitReceiptAuthority(t *testing.T) { + for _, name := range []string{"marked rewrite", "exact headerless", "headerless rewrite", "wrong token", "wrong sender", "foreign recipient", "unstamped", "unfinished", "expired", "consumed", "retired"} { + t.Run(name, func(t *testing.T) { + f, r := lineageFixture(t) + ctx := t.Context() + exec := func(query string, args ...any) { + t.Helper() + if _, err := f.pool.Exec(ctx, query, args...); err != nil { + t.Fatal(err) + } + } + token, receipt := uuid.New(), uuid.New() + exec(`INSERT INTO warmup_tasks(task_id,lineage_version,subject,max_turns,scenario_version,rendering_version)VALUES($1,1,'Pinned',3,'static-v1','canonical-v1')`, f.task) + exec(`INSERT INTO warmup_tokens(token,task_id,sender_account_id,recipient_account_id,sent_message_id,conversation_id)VALUES($1,$2,$3,$4,'',$5)`, token, f.task, f.sender, f.recipient, uuid.New()) + proof := WarmupReceiptProof{Token: token, SenderAddress: f.senderTo} + recipient, messageID := f.recipient, "received@mailjet.com" + switch name { + case "exact headerless": + proof.Token, messageID = uuid.Nil, "sent@example.test" + case "headerless rewrite": + proof.Token = uuid.Nil + case "wrong token": + proof.Token = uuid.New() + case "wrong sender": + proof.SenderAddress = "forwarder@example.test" + case "foreign recipient": + recipient = f.sender + case "unstamped": + exec(`UPDATE warmup_tokens SET sent_message_id='' WHERE token=$1`, token) + case "unfinished": + exec(`UPDATE tasks SET status='active' WHERE id=$1`, f.task) + case "expired": + exec(`UPDATE warmup_tokens SET expires_at=NOW()-INTERVAL '1 second' WHERE token=$1`, token) + case "consumed": + exec(`UPDATE warmup_tokens SET consumed_at=NOW() WHERE token=$1`, token) + case "retired": + exec(`UPDATE warmup_tokens SET sent_retired_at=NOW() WHERE token=$1`, token) + } + exec(`INSERT INTO warmup_received(email_account_id,internal_id,sender_account_id,message_id)VALUES($1,$2,$3,$4)`, recipient, receipt, f.sender, messageID) + err := r.RecordVerifiedWarmupParent(ctx, f.task, recipient, receipt, "recipient-thread", proof) + allowed := name == "marked rewrite" || name == "exact headerless" + if (err == nil) != allowed { + t.Fatalf("allowed=%v err=%v", allowed, err) + } + if !allowed { + return + } + if err := r.RecordVerifiedWarmupParent(ctx, f.task, recipient, receipt, "recipient-thread", proof); err != nil { + t.Fatal("retry of bound receipt failed", err) + } + duplicate := uuid.New() + exec(`INSERT INTO warmup_received(email_account_id,internal_id,sender_account_id,message_id)VALUES($1,$2,$3,'replayed@smtp-pulse.com')`, recipient, duplicate, f.sender) + if err := r.RecordVerifiedWarmupParent(ctx, f.task, recipient, duplicate, "", proof); err == nil { + t.Fatal("repeated marker authorized another receipt") + } + ok, err := r.BindWarmupSuccessor(ctx, f.task, recipient, receipt, time.Now().Add(time.Hour)) + if err != nil || !ok { + t.Fatalf("successor bound=%v err=%v", ok, err) + } + var successor uuid.UUID + if err := f.pool.QueryRow(ctx, `SELECT task_id FROM warmup_tasks WHERE parent_task_id=$1`, f.task).Scan(&successor); err != nil { + t.Fatal(err) + } + parent, err := r.ExactWarmupParent(ctx, successor) + if err != nil || parent == nil || parent.MessageID != messageID || parent.ThreadID == nil || *parent.ThreadID != "recipient-thread" { + t.Fatalf("reply did not retain actual received ID and thread: %+v %v", parent, err) + } + }) + } +} + +func TestLiveWarmupReceiptSerializationLeavesSmallPoolCapacity(t *testing.T) { + f, _ := lineageFixture(t) + ctx, cancel := context.WithTimeout(t.Context(), 5*time.Second) + defer cancel() + cfg, err := pgxpool.ParseConfig(os.Getenv("WARMBLY_TEST_DB")) + if err != nil { + t.Fatal(err) + } + cfg.MaxConns, cfg.MinConns = 2, 0 + pool, err := pgxpool.NewWithConfig(ctx, cfg) + if err != nil { + t.Fatal(err) + } + t.Cleanup(pool.Close) + r := NewWarmupRepository(pool).(*warmupRepository) + var wg sync.WaitGroup + var accepted atomic.Int32 + for range 16 { + token := uuid.New() + if _, err := f.pool.Exec(ctx, `INSERT INTO warmup_tokens(token,task_id,sender_account_id,recipient_account_id)VALUES($1,$2,$3,$4)`, token, f.task, f.sender, f.recipient); err != nil { + t.Fatal(err) + } + wg.Add(1) + go func() { + defer wg.Done() + _, err := r.ProcessWarmupReceipt(ctx, token, f.recipient, WarmupReceiptProof{SenderAddress: f.senderTo}, func(current *models.WarmupToken) error { + var value int + if err := pool.QueryRow(ctx, `SELECT 1`).Scan(&value); err != nil { + return err + } + accepted.Add(1) + return r.ConsumeWarmupToken(ctx, current.Token) + }) + if err != nil && !errors.Is(err, ErrWarmupReceiptBusy) { + t.Error(err) + } + }() + } + wg.Wait() + if accepted.Load() == 0 { + t.Fatal("receipt lock exhausted capacity for the acceptance callback") + } +} + func TestLiveLineageDispatchClaimsNativeExecutionOnceAndReplaysDurableResult(t *testing.T) { f, r := lineageFixture(t) ctx := t.Context()