diff --git a/admin/billing/orb.go b/admin/billing/orb.go index fad0bc621b0c..dad84c7c8c45 100644 --- a/admin/billing/orb.go +++ b/admin/billing/orb.go @@ -35,10 +35,11 @@ var ErrCustomerIDRequired = errors.New("customer id is required") var _ Biller = &Orb{} type Orb struct { - client *orb.Client - logger *zap.Logger - webhookSecret string - taxProvider string + client *orb.Client + logger *zap.Logger + webhookSecret string + taxProvider string + usageRetryBackoff []time.Duration } func NewOrb(logger *zap.Logger, orbKey, webhookSecret, taxProvider string) Biller { @@ -672,7 +673,11 @@ func (o *Orb) getAllPlans(ctx context.Context) ([]*Plan, error) { } func (o *Orb) pushUsage(ctx context.Context, usage *[]orb.EventIngestParamsEvent) error { - re := retrier.New(retrier.ExponentialBackoff(5, 500*time.Millisecond), retryErrClassifier{}) + backoff := o.usageRetryBackoff + if backoff == nil { + backoff = retrier.ExponentialBackoff(5, 500*time.Millisecond) + } + re := retrier.New(backoff, retryErrClassifier{}) err := re.RunCtx(ctx, func(ctx context.Context) error { resp, err := o.client.Events.Ingest(ctx, orb.EventIngestParams{ Events: orb.F(*usage), diff --git a/admin/billing/orb_usage_test.go b/admin/billing/orb_usage_test.go new file mode 100644 index 000000000000..5af6661aa5b3 --- /dev/null +++ b/admin/billing/orb_usage_test.go @@ -0,0 +1,230 @@ +package billing + +import ( + "encoding/json" + "fmt" + "net/http" + "net/http/httptest" + "sync" + "testing" + "time" + + "github.com/orbcorp/orb-go" + "github.com/orbcorp/orb-go/option" + "github.com/stretchr/testify/require" + "go.uber.org/zap" +) + +func TestOrbReportUsageBatchBoundaries(t *testing.T) { + // Orb accepts at most 500 events per ingestion request. These exact boundaries + // guard against empty requests, dropped remainders, and accidental 501-event batches. + tests := []struct { + count int + wantBatches []int + }{ + {count: 0, wantBatches: nil}, + {count: 1, wantBatches: []int{1}}, + {count: 500, wantBatches: []int{500}}, + {count: 501, wantBatches: []int{500, 1}}, + } + for _, tt := range tests { + t.Run(fmt.Sprintf("%d events", tt.count), func(t *testing.T) { + server := newOrbUsageServer(t, func(_ int, _ orbUsageRequest) (int, string) { + return http.StatusOK, `{"validation_failed":[]}` + }) + biller := newTestOrbUsageBiller(server.server.URL) + + err := biller.ReportUsage(t.Context(), makeOrbUsageFixtures(tt.count)) + require.NoError(t, err) + require.Equal(t, tt.wantBatches, server.batchSizes()) + }) + } +} + +func TestOrbReportUsageSerializesStableLogicalEvent(t *testing.T) { + // Retries and repeated checkpoints must serialize the same logical bucket to + // the same key, regardless of metadata map insertion order. + server := newOrbUsageServer(t, func(_ int, _ orbUsageRequest) (int, string) { + return http.StatusOK, `{"validation_failed":[]}` + }) + biller := newTestOrbUsageBiller(server.server.URL) + end := time.Date(2026, time.January, 2, 4, 0, 0, 0, time.UTC) + first := &Usage{ + CustomerID: "org-1", MetricName: "api_calls", Value: 42, + ReportingGrain: UsageReportingGranularityHour, EndTime: end, + Metadata: map[string]interface{}{"region": "us-east", "tier": "pro"}, + } + second := &Usage{ + CustomerID: "org-1", MetricName: "api_calls", Value: 42, + ReportingGrain: UsageReportingGranularityHour, EndTime: end, + Metadata: map[string]interface{}{"tier": "pro", "region": "us-east"}, + } + + require.NoError(t, biller.ReportUsage(t.Context(), []*Usage{first})) + require.NoError(t, biller.ReportUsage(t.Context(), []*Usage{second})) + requests := server.requestsSnapshot() + require.Len(t, requests, 2) + require.Len(t, requests[0].Events, 1) + require.Equal(t, requests[0].Events[0].IdempotencyKey, requests[1].Events[0].IdempotencyKey) + require.Equal(t, "api_calls_hour", requests[0].Events[0].EventName) + require.Equal(t, "org-1", requests[0].Events[0].ExternalCustomerID) + require.Equal(t, end.Add(-time.Second), requests[0].Events[0].Timestamp) + require.EqualValues(t, 42, requests[0].Events[0].Properties["amount"]) + require.Equal(t, "us-east", requests[0].Events[0].Properties["region"]) +} + +func TestOrbReportUsageCorrectionKeepsLogicalIdempotencyKey(t *testing.T) { + // A corrected value represents the same customer/metric/time bucket. Keeping + // its identity stable prevents a correction retry from becoming an additive event. + server := newOrbUsageServer(t, func(_ int, _ orbUsageRequest) (int, string) { + return http.StatusOK, `{"validation_failed":[]}` + }) + biller := newTestOrbUsageBiller(server.server.URL) + fixtures := makeOrbUsageFixtures(1) + require.NoError(t, biller.ReportUsage(t.Context(), fixtures)) + fixtures[0].Value = 99 + require.NoError(t, biller.ReportUsage(t.Context(), fixtures)) + + requests := server.requestsSnapshot() + require.Equal(t, requests[0].Events[0].IdempotencyKey, requests[1].Events[0].IdempotencyKey) + require.EqualValues(t, 1, requests[0].Events[0].Properties["amount"]) + require.EqualValues(t, 99, requests[1].Events[0].Properties["amount"]) +} + +func TestOrbReportUsageRetriesOnlyTransientProviderFailures(t *testing.T) { + // The generated Orb client is configured without its own retry loop here so + // the application retry classifier and its exact request count are observable. + tests := []struct { + name string + firstStatus int + wantRequests int + wantErr bool + }{ + {name: "500 retries", firstStatus: http.StatusInternalServerError, wantRequests: 2}, + {name: "429 retries", firstStatus: http.StatusTooManyRequests, wantRequests: 2}, + {name: "400 fails immediately", firstStatus: http.StatusBadRequest, wantRequests: 1, wantErr: true}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + server := newOrbUsageServer(t, func(attempt int, _ orbUsageRequest) (int, string) { + if attempt == 1 { + return tt.firstStatus, orbUsageErrorBody(tt.firstStatus) + } + return http.StatusOK, `{"validation_failed":[]}` + }) + biller := newTestOrbUsageBiller(server.server.URL) + + err := biller.ReportUsage(t.Context(), makeOrbUsageFixtures(1)) + if tt.wantErr { + require.Error(t, err) + } else { + require.NoError(t, err) + } + require.Equal(t, tt.wantRequests, len(server.requestsSnapshot())) + }) + } +} + +func orbUsageErrorBody(status int) string { + return fmt.Sprintf(`{"status":%d,"title":"provider failure","type":"test_error","validation_errors":[],"detail":"provider failure"}`, status) +} + +func TestOrbReportUsageReturnsRecordValidationDetails(t *testing.T) { + // A 200 response may still reject individual records; the error must identify + // their idempotency keys so operators can repair the correct usage checkpoint. + server := newOrbUsageServer(t, func(_ int, request orbUsageRequest) (int, string) { + return http.StatusOK, fmt.Sprintf(`{"validation_failed":[{"idempotency_key":%q,"validation_errors":["amount must be positive"]}]}`, request.Events[0].IdempotencyKey) + }) + biller := newTestOrbUsageBiller(server.server.URL) + + err := biller.ReportUsage(t.Context(), makeOrbUsageFixtures(1)) + require.ErrorContains(t, err, "validation failure for 1 events") + require.ErrorContains(t, err, "amount must be positive") + require.ErrorContains(t, err, "org-0") + require.Len(t, server.requestsSnapshot(), 1) +} + +type orbUsageRequest struct { + Events []struct { + EventName string `json:"event_name"` + IdempotencyKey string `json:"idempotency_key"` + ExternalCustomerID string `json:"external_customer_id"` + Timestamp time.Time `json:"timestamp"` + Properties map[string]interface{} `json:"properties"` + } `json:"events"` +} + +type orbUsageServer struct { + server *httptest.Server + + mu sync.Mutex + requests []orbUsageRequest +} + +func newOrbUsageServer(t *testing.T, response func(attempt int, request orbUsageRequest) (int, string)) *orbUsageServer { + t.Helper() + fixture := &orbUsageServer{} + fixture.server = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + require.Equal(t, http.MethodPost, r.Method) + require.Equal(t, "/ingest", r.URL.Path) + var request orbUsageRequest + require.NoError(t, json.NewDecoder(r.Body).Decode(&request)) + fixture.mu.Lock() + fixture.requests = append(fixture.requests, request) + attempt := len(fixture.requests) + fixture.mu.Unlock() + status, body := response(attempt, request) + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(status) + _, _ = w.Write([]byte(body)) + })) + t.Cleanup(fixture.server.Close) + return fixture +} + +func (s *orbUsageServer) requestsSnapshot() []orbUsageRequest { + s.mu.Lock() + defer s.mu.Unlock() + return append([]orbUsageRequest(nil), s.requests...) +} + +func (s *orbUsageServer) batchSizes() []int { + requests := s.requestsSnapshot() + if len(requests) == 0 { + return nil + } + res := make([]int, len(requests)) + for i, request := range requests { + res[i] = len(request.Events) + } + return res +} + +func newTestOrbUsageBiller(baseURL string) *Orb { + client := orb.NewClient( + option.WithAPIKey("test-key"), + option.WithBaseURL(baseURL), + option.WithMaxRetries(0), + ) + return &Orb{ + client: client, + logger: zap.NewNop(), + usageRetryBackoff: []time.Duration{0, 0, 0, 0, 0}, + } +} + +func makeOrbUsageFixtures(count int) []*Usage { + usage := make([]*Usage, count) + for i := range usage { + usage[i] = &Usage{ + CustomerID: fmt.Sprintf("org-%d", i), + MetricName: "api_calls", + Value: float64(i + 1), + ReportingGrain: UsageReportingGranularityHour, + StartTime: time.Date(2026, time.January, 2, 3, 0, 0, 0, time.UTC), + EndTime: time.Date(2026, time.January, 2, 4, 0, 0, 0, time.UTC), + Metadata: map[string]interface{}{"region": "us-east"}, + } + } + return usage +} diff --git a/admin/billing/orb_webhook.go b/admin/billing/orb_webhook.go index 7da960aec7ff..3316066095d5 100644 --- a/admin/billing/orb_webhook.go +++ b/admin/billing/orb_webhook.go @@ -46,6 +46,10 @@ func (o *orbWebhook) handleWebhook(w http.ResponseWriter, r *http.Request) error r.Body = http.MaxBytesReader(w, r.Body, maxBodyBytes) payload, err := io.ReadAll(r.Body) if err != nil { + var maxBytesErr *http.MaxBytesError + if errors.As(err, &maxBytesErr) { + return httputil.Errorf(http.StatusRequestEntityTooLarge, "webhook request body exceeds %d bytes", maxBodyBytes) + } return httputil.Errorf(http.StatusServiceUnavailable, "error reading request body: %w", err) } @@ -61,12 +65,17 @@ func (o *orbWebhook) handleWebhook(w http.ResponseWriter, r *http.Request) error return nil } + // The generic type above is only used to discard events we do not consume. + // Before a recognized event can cause a side effect, verify the signature + // against the exact bytes Orb sent; re-marshaling would change those bytes. now := time.Now().UTC() err = o.verifySignature(payload, r.Header, now) if err != nil { return httputil.Errorf(http.StatusBadRequest, "error verifying webhook signature: %w", err) } + // Do not acknowledge work-bearing events until their durable job is queued. + // A queue error becomes a 5xx so Orb retries instead of silently losing work. switch e.Type { case "invoice.payment_succeeded": var ie invoiceEvent @@ -154,6 +163,9 @@ func (o *orbWebhook) handleWebhook(w http.ResponseWriter, r *http.Request) error return nil } +// Orb delivers webhooks at least once. These handlers deliberately leave replay +// detection to the durable jobs layer; Duplicate means the work is already +// recorded, so it is safe to acknowledge the delivery. func (o *orbWebhook) handleInvoicePaymentSucceeded(ctx context.Context, ie invoiceEvent) error { res, err := o.jobs.PaymentSuccess(ctx, ie.OrbInvoice.Customer.ExternalCustomerID, ie.OrbInvoice.ID) if err != nil { @@ -269,20 +281,22 @@ func (o *orbWebhook) verifySignature(payload []byte, headers http.Header, now ti mac.Write(payload) expected := mac.Sum(nil) - for _, part := range msgSignature { - parts := strings.Split(part, "=") - if len(parts) != 2 { - continue - } - if parts[0] != "v1" { - continue - } - signature, err := hex.DecodeString(parts[1]) - if err != nil { - continue - } - if hmac.Equal(signature, expected) { - return nil + // Orb can include more than one signature during secret rotation. HTTP + // intermediaries may preserve them as repeated fields or combine them with + // commas, so accept either representation when any v1 signature matches. + for _, value := range msgSignature { + for _, part := range strings.Split(value, ",") { + version, encoded, ok := strings.Cut(strings.TrimSpace(part), "=") + if !ok || version != "v1" { + continue + } + signature, err := hex.DecodeString(encoded) + if err != nil { + continue + } + if hmac.Equal(signature, expected) { + return nil + } } } diff --git a/admin/billing/orb_webhook_test.go b/admin/billing/orb_webhook_test.go new file mode 100644 index 000000000000..1d542ec727d0 --- /dev/null +++ b/admin/billing/orb_webhook_test.go @@ -0,0 +1,511 @@ +package billing + +import ( + "bytes" + "context" + "crypto/hmac" + "crypto/sha256" + "encoding/hex" + "encoding/json" + "errors" + "net/http" + "net/http/httptest" + "testing" + "time" + + "github.com/rilldata/rill/admin/jobs" + "github.com/stretchr/testify/require" + "go.uber.org/zap" +) + +const testOrbWebhookSecret = "orb_webhook_secret_bill_001" + +var ( + testOrbDueDate = time.Date(2026, time.July, 29, 12, 30, 0, 0, time.UTC) + testOrbPaymentFailed = time.Date(2026, time.July, 22, 8, 15, 30, 123000000, time.UTC) +) + +type orbPaymentSuccessCall struct { + billingCustomerID string + invoiceID string +} + +type orbPaymentFailedCall struct { + billingCustomerID string + invoiceID string + invoiceNumber string + invoiceURL string + amount string + currency string + dueDate time.Time + failedAt time.Time +} + +// orbWebhookJobs deliberately implements only the queue calls reachable from +// this webhook. The embedded interface makes any unexpected queue dependency +// fail loudly instead of turning the test fake into a second jobs.Client. +type orbWebhookJobs struct { + jobs.Client + + paymentSuccessCalls []orbPaymentSuccessCall + paymentFailedCalls []orbPaymentFailedCall + planChangedCalls []string + creditBalanceDroppedCalls []string + creditBalanceDepletedCalls []string + queueErr error + duplicateOnRepeat bool +} + +func (j *orbWebhookJobs) PaymentSuccess(_ context.Context, billingCustomerID, invoiceID string) (*jobs.InsertResult, error) { + j.paymentSuccessCalls = append(j.paymentSuccessCalls, orbPaymentSuccessCall{ + billingCustomerID: billingCustomerID, + invoiceID: invoiceID, + }) + return j.insertResult() +} + +func (j *orbWebhookJobs) PaymentFailed(_ context.Context, billingCustomerID, invoiceID, invoiceNumber, invoiceURL, amount, currency string, dueDate, failedAt time.Time) (*jobs.InsertResult, error) { + j.paymentFailedCalls = append(j.paymentFailedCalls, orbPaymentFailedCall{ + billingCustomerID: billingCustomerID, + invoiceID: invoiceID, + invoiceNumber: invoiceNumber, + invoiceURL: invoiceURL, + amount: amount, + currency: currency, + dueDate: dueDate, + failedAt: failedAt, + }) + return j.insertResult() +} + +func (j *orbWebhookJobs) PlanChanged(_ context.Context, billingCustomerID string) (*jobs.InsertResult, error) { + j.planChangedCalls = append(j.planChangedCalls, billingCustomerID) + return j.insertResult() +} + +func (j *orbWebhookJobs) CreditBalanceDropped(_ context.Context, billingCustomerID string) (*jobs.InsertResult, error) { + j.creditBalanceDroppedCalls = append(j.creditBalanceDroppedCalls, billingCustomerID) + return j.insertResult() +} + +func (j *orbWebhookJobs) CreditBalanceDepleted(_ context.Context, billingCustomerID string) (*jobs.InsertResult, error) { + j.creditBalanceDepletedCalls = append(j.creditBalanceDepletedCalls, billingCustomerID) + return j.insertResult() +} + +func (j *orbWebhookJobs) insertResult() (*jobs.InsertResult, error) { + if j.queueErr != nil { + return nil, j.queueErr + } + return &jobs.InsertResult{Duplicate: j.duplicateOnRepeat && j.callCount() > 1}, nil +} + +func (j *orbWebhookJobs) callCount() int { + return len(j.paymentSuccessCalls) + + len(j.paymentFailedCalls) + + len(j.planChangedCalls) + + len(j.creditBalanceDroppedCalls) + + len(j.creditBalanceDepletedCalls) +} + +func TestOrbWebhookSignedInterestingEvents(t *testing.T) { + // Map every supported signed event to its durable job payload and characterize intentionally ignored events. + tests := []struct { + name string + eventType string + payload func(*testing.T, string) []byte + assertJobs func(*testing.T, *orbWebhookJobs) + }{ + { + name: "invoice payment succeeded", + eventType: "invoice.payment_succeeded", + payload: orbInvoicePayload, + assertJobs: func(t *testing.T, got *orbWebhookJobs) { + require.Equal(t, []orbPaymentSuccessCall{{ + billingCustomerID: "org_invoice_customer", + invoiceID: "inv_bill_001", + }}, got.paymentSuccessCalls) + }, + }, + { + name: "invoice payment failed", + eventType: "invoice.payment_failed", + payload: orbInvoicePayload, + assertJobs: func(t *testing.T, got *orbWebhookJobs) { + require.Equal(t, []orbPaymentFailedCall{{ + billingCustomerID: "org_invoice_customer", + invoiceID: "inv_bill_001", + invoiceNumber: "RILL-000042", + invoiceURL: "https://billing.example.test/invoices/inv_bill_001", + amount: "42.75", + currency: "USD", + dueDate: testOrbDueDate, + failedAt: testOrbPaymentFailed, + }}, got.paymentFailedCalls) + }, + }, + { + name: "invoice issue failed", + eventType: "invoice.issue_failed", + payload: orbInvoicePayload, + assertJobs: func(t *testing.T, got *orbWebhookJobs) { + require.Zero(t, got.callCount()) + }, + }, + { + name: "subscription started", + eventType: "subscription.started", + payload: orbSubscriptionPayload, + assertJobs: func(t *testing.T, got *orbWebhookJobs) { + require.Equal(t, []string{"org_subscription_customer"}, got.planChangedCalls) + }, + }, + { + name: "subscription ended", + eventType: "subscription.ended", + payload: orbSubscriptionPayload, + assertJobs: func(t *testing.T, got *orbWebhookJobs) { + require.Equal(t, []string{"org_subscription_customer"}, got.planChangedCalls) + }, + }, + { + name: "subscription plan changed", + eventType: "subscription.plan_changed", + payload: orbSubscriptionPayload, + assertJobs: func(t *testing.T, got *orbWebhookJobs) { + require.Equal(t, []string{"org_subscription_customer"}, got.planChangedCalls) + }, + }, + { + name: "credit balance dropped", + eventType: "customer.credit_balance_dropped", + payload: orbCreditBalancePayload, + assertJobs: func(t *testing.T, got *orbWebhookJobs) { + require.Equal(t, []string{"org_credit_customer"}, got.creditBalanceDroppedCalls) + }, + }, + { + name: "credit balance depleted", + eventType: "customer.credit_balance_depleted", + payload: orbCreditBalancePayload, + assertJobs: func(t *testing.T, got *orbWebhookJobs) { + require.Equal(t, []string{"org_credit_customer"}, got.creditBalanceDepletedCalls) + }, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + // Every state-changing path is exercised with an HMAC over the exact + // transmitted JSON, not a stubbed verifier or a made-up digest. + payload := tt.payload(t, tt.eventType) + timestamp := time.Now().UTC().Format(webhookHeaderTimestampFormat) + jobClient := &orbWebhookJobs{} + response := serveOrbWebhook(t, jobClient, payload, timestamp, orbSignedHeader(payload, timestamp, testOrbWebhookSecret)) + + require.Equal(t, http.StatusOK, response.Code) + tt.assertJobs(t, jobClient) + if tt.eventType != "invoice.issue_failed" { + require.Equal(t, 1, jobClient.callCount()) + } + }) + } +} + +func TestOrbWebhookSignatureHeaders(t *testing.T) { + // Signature parsing must reject incomplete candidates while accepting one valid key-rotation candidate. + payload := orbInvoicePayload(t, "invoice.payment_succeeded") + timestamp := time.Now().UTC().Format(webhookHeaderTimestampFormat) + valid := orbSignedHeader(payload, timestamp, testOrbWebhookSecret) + wrong := orbSignedHeader(payload, timestamp, "wrong-secret") + + tests := []struct { + name string + timestamp string + signatures []string + wantStatus int + wantCalls int + }{ + {name: "missing signature", timestamp: timestamp, wantStatus: http.StatusBadRequest}, + {name: "missing timestamp", signatures: []string{valid}, wantStatus: http.StatusBadRequest}, + {name: "invalid timestamp", timestamp: "not-a-timestamp", signatures: []string{valid}, wantStatus: http.StatusBadRequest}, + {name: "invalid hex", timestamp: timestamp, signatures: []string{"v1=not-hex"}, wantStatus: http.StatusBadRequest}, + {name: "wrong digest", timestamp: timestamp, signatures: []string{wrong}, wantStatus: http.StatusBadRequest}, + {name: "unsupported version", timestamp: timestamp, signatures: []string{"v2=" + valid[3:]}, wantStatus: http.StatusBadRequest}, + {name: "all repeated signatures invalid", timestamp: timestamp, signatures: []string{wrong, "v1=00"}, wantStatus: http.StatusBadRequest}, + {name: "valid repeated signature", timestamp: timestamp, signatures: []string{wrong, valid}, wantStatus: http.StatusOK, wantCalls: 1}, + {name: "valid comma separated signature", timestamp: timestamp, signatures: []string{wrong + ", " + valid}, wantStatus: http.StatusOK, wantCalls: 1}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + // During key rotation Orb may send several candidates. Authentication + // succeeds only if at least one complete v1 HMAC matches. + jobClient := &orbWebhookJobs{} + response := serveOrbWebhook(t, jobClient, payload, tt.timestamp, tt.signatures...) + + require.Equal(t, tt.wantStatus, response.Code) + require.Equal(t, tt.wantCalls, jobClient.callCount()) + if tt.wantCalls == 1 { + require.Equal(t, []orbPaymentSuccessCall{{ + billingCustomerID: "org_invoice_customer", + invoiceID: "inv_bill_001", + }}, jobClient.paymentSuccessCalls) + } + }) + } +} + +func TestOrbWebhookSignatureTimestampWindow(t *testing.T) { + // Pin the inclusive replay window at both five-minute boundaries without relying on wall-clock timing. + now := time.Date(2026, time.July, 23, 9, 45, 0, 123456789, time.UTC) + payload := orbInvoicePayload(t, "invoice.payment_succeeded") + verifier := &orbWebhook{orb: &Orb{webhookSecret: testOrbWebhookSecret}} + + tests := []struct { + name string + timestamp time.Time + wantErrText string + }{ + {name: "exactly five minutes old", timestamp: now.Add(-5 * time.Minute)}, + {name: "one nanosecond too old", timestamp: now.Add(-5*time.Minute - time.Nanosecond), wantErrText: "too old"}, + {name: "current", timestamp: now}, + {name: "exactly five minutes new", timestamp: now.Add(5 * time.Minute)}, + {name: "one nanosecond too new", timestamp: now.Add(5*time.Minute + time.Nanosecond), wantErrText: "too new"}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + // The freshness window is inclusive at both five-minute boundaries; + // anything outside it is rejected even when its HMAC is otherwise valid. + timestamp := tt.timestamp.Format(webhookHeaderTimestampFormat) + headers := make(http.Header) + headers.Set("X-Orb-Timestamp", timestamp) + headers.Set("X-Orb-Signature", orbSignedHeader(payload, timestamp, testOrbWebhookSecret)) + + err := verifier.verifySignature(payload, headers, now) + if tt.wantErrText == "" { + require.NoError(t, err) + } else { + require.ErrorContains(t, err, tt.wantErrText) + } + }) + } +} + +func TestOrbWebhookReplayUsesDurableQueueDeduplication(t *testing.T) { + // Webhook delivery is at-least-once. Both deliveries must reach the durable + // queue with identical keys; its Duplicate result makes the replay a safe 200. + payload := orbInvoicePayload(t, "invoice.payment_succeeded") + timestamp := time.Now().UTC().Format(webhookHeaderTimestampFormat) + signature := orbSignedHeader(payload, timestamp, testOrbWebhookSecret) + jobClient := &orbWebhookJobs{duplicateOnRepeat: true} + + first := serveOrbWebhook(t, jobClient, payload, timestamp, signature) + second := serveOrbWebhook(t, jobClient, payload, timestamp, signature) + + require.Equal(t, http.StatusOK, first.Code) + require.Equal(t, http.StatusOK, second.Code) + require.Equal(t, []orbPaymentSuccessCall{ + {billingCustomerID: "org_invoice_customer", invoiceID: "inv_bill_001"}, + {billingCustomerID: "org_invoice_customer", invoiceID: "inv_bill_001"}, + }, jobClient.paymentSuccessCalls) + require.Equal(t, 2, jobClient.callCount()) +} + +func TestOrbWebhookBodyValidation(t *testing.T) { + // Signed input is still untrusted: malformed events and bodies beyond the explicit limit must fail before enqueue. + t.Run("malformed JSON", func(t *testing.T) { + payload := []byte(`{"type":`) + timestamp := time.Now().UTC().Format(webhookHeaderTimestampFormat) + jobClient := &orbWebhookJobs{} + response := serveOrbWebhook(t, jobClient, payload, timestamp, orbSignedHeader(payload, timestamp, testOrbWebhookSecret)) + + require.Equal(t, http.StatusBadRequest, response.Code) + require.Zero(t, jobClient.callCount()) + }) + + t.Run("malformed recognized event", func(t *testing.T) { + payload := []byte(`{"type":"invoice.payment_succeeded","created_at":42}`) + timestamp := time.Now().UTC().Format(webhookHeaderTimestampFormat) + jobClient := &orbWebhookJobs{} + response := serveOrbWebhook(t, jobClient, payload, timestamp, orbSignedHeader(payload, timestamp, testOrbWebhookSecret)) + + require.Equal(t, http.StatusBadRequest, response.Code) + require.Zero(t, jobClient.callCount()) + }) + + t.Run("exactly 64 KiB", func(t *testing.T) { + payload := []byte(`{"type":"event.not_used"}`) + payload = append(payload, bytes.Repeat([]byte(" "), int(maxBodyBytes)-len(payload))...) + jobClient := &orbWebhookJobs{} + response := serveOrbWebhook(t, jobClient, payload, "") + + require.Equal(t, int(maxBodyBytes), len(payload)) + require.Equal(t, http.StatusOK, response.Code) + require.Zero(t, jobClient.callCount()) + }) + + t.Run("larger than 64 KiB", func(t *testing.T) { + payload := bytes.Repeat([]byte("x"), int(maxBodyBytes)+1) + timestamp := time.Now().UTC().Format(webhookHeaderTimestampFormat) + jobClient := &orbWebhookJobs{} + response := serveOrbWebhook(t, jobClient, payload, timestamp, orbSignedHeader(payload, timestamp, testOrbWebhookSecret)) + + require.Equal(t, http.StatusRequestEntityTooLarge, response.Code) + require.Zero(t, jobClient.callCount()) + }) +} + +func TestOrbWebhookQueueFailuresReturnRetryableStatus(t *testing.T) { + // Every supported event returns a retryable server error when its handoff to the durable queue fails. + queueErr := errors.New("queue unavailable") + tests := []struct { + name string + eventType string + payload func(*testing.T, string) []byte + assertJobs func(*testing.T, *orbWebhookJobs) + }{ + { + name: "payment success", eventType: "invoice.payment_succeeded", payload: orbInvoicePayload, + assertJobs: func(t *testing.T, got *orbWebhookJobs) { + require.Equal(t, []orbPaymentSuccessCall{{billingCustomerID: "org_invoice_customer", invoiceID: "inv_bill_001"}}, got.paymentSuccessCalls) + }, + }, + { + name: "payment failed", eventType: "invoice.payment_failed", payload: orbInvoicePayload, + assertJobs: func(t *testing.T, got *orbWebhookJobs) { + require.Equal(t, []orbPaymentFailedCall{{ + billingCustomerID: "org_invoice_customer", + invoiceID: "inv_bill_001", + invoiceNumber: "RILL-000042", + invoiceURL: "https://billing.example.test/invoices/inv_bill_001", + amount: "42.75", + currency: "USD", + dueDate: testOrbDueDate, + failedAt: testOrbPaymentFailed, + }}, got.paymentFailedCalls) + }, + }, + { + name: "plan changed", eventType: "subscription.plan_changed", payload: orbSubscriptionPayload, + assertJobs: func(t *testing.T, got *orbWebhookJobs) { + require.Equal(t, []string{"org_subscription_customer"}, got.planChangedCalls) + }, + }, + { + name: "credit dropped", eventType: "customer.credit_balance_dropped", payload: orbCreditBalancePayload, + assertJobs: func(t *testing.T, got *orbWebhookJobs) { + require.Equal(t, []string{"org_credit_customer"}, got.creditBalanceDroppedCalls) + }, + }, + { + name: "credit depleted", eventType: "customer.credit_balance_depleted", payload: orbCreditBalancePayload, + assertJobs: func(t *testing.T, got *orbWebhookJobs) { + require.Equal(t, []string{"org_credit_customer"}, got.creditBalanceDepletedCalls) + }, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + // A failed insert must never be acknowledged: the 5xx is what asks Orb + // to retry, while the recorded call proves no identifiers were dropped. + payload := tt.payload(t, tt.eventType) + timestamp := time.Now().UTC().Format(webhookHeaderTimestampFormat) + jobClient := &orbWebhookJobs{queueErr: queueErr} + response := serveOrbWebhook(t, jobClient, payload, timestamp, orbSignedHeader(payload, timestamp, testOrbWebhookSecret)) + + require.Equal(t, http.StatusInternalServerError, response.Code) + require.Equal(t, 1, jobClient.callCount()) + tt.assertJobs(t, jobClient) + }) + } +} + +func orbInvoicePayload(t *testing.T, eventType string) []byte { + t.Helper() + return orbJSONPayload(t, map[string]any{ + "id": "evt_invoice_bill_001", + "type": eventType, + "created_at": "2026-07-23T08:00:00Z", + "properties": map[string]any{"reason": "test"}, + "invoice": map[string]any{ + "id": "inv_bill_001", + "invoice_number": "RILL-000042", + "hosted_invoice_url": "https://billing.example.test/invoices/inv_bill_001", + "amount_due": "42.75", + "currency": "USD", + "due_date": testOrbDueDate.Format(time.RFC3339Nano), + "payment_failed_at": testOrbPaymentFailed.Format(time.RFC3339Nano), + "customer": map[string]any{ + "id": "orb_customer_invoice", + "external_customer_id": "org_invoice_customer", + }, + }, + }) +} + +func orbSubscriptionPayload(t *testing.T, eventType string) []byte { + t.Helper() + return orbJSONPayload(t, map[string]any{ + "id": "evt_subscription_bill_001", + "type": eventType, + "created_at": "2026-07-23T08:00:00Z", + "subscription": map[string]any{ + "id": "sub_bill_001", + "customer": map[string]any{ + "id": "orb_customer_subscription", + "external_customer_id": "org_subscription_customer", + }, + }, + }) +} + +func orbCreditBalancePayload(t *testing.T, eventType string) []byte { + t.Helper() + return orbJSONPayload(t, map[string]any{ + "id": "evt_credit_bill_001", + "type": eventType, + "created_at": "2026-07-23T08:00:00Z", + "customer": map[string]any{ + "id": "orb_customer_credit", + "external_customer_id": "org_credit_customer", + }, + }) +} + +func orbJSONPayload(t *testing.T, event map[string]any) []byte { + t.Helper() + payload, err := json.Marshal(event) + require.NoError(t, err) + return payload +} + +func orbSignedHeader(payload []byte, timestamp, secret string) string { + mac := hmac.New(sha256.New, []byte(secret)) + mac.Write([]byte("v1:")) + mac.Write([]byte(timestamp)) + mac.Write([]byte(":")) + mac.Write(payload) + return "v1=" + hex.EncodeToString(mac.Sum(nil)) +} + +func serveOrbWebhook(t *testing.T, jobClient jobs.Client, payload []byte, timestamp string, signatures ...string) *httptest.ResponseRecorder { + t.Helper() + handler := NewOrb(zap.NewNop(), "", testOrbWebhookSecret, "").WebhookHandlerFunc(context.Background(), jobClient) + require.NotNil(t, handler) + + request := httptest.NewRequest(http.MethodPost, "/orb/webhook", bytes.NewReader(payload)) + if timestamp != "" { + request.Header.Set("X-Orb-Timestamp", timestamp) + } + for _, signature := range signatures { + request.Header.Add("X-Orb-Signature", signature) + } + response := httptest.NewRecorder() + require.NotPanics(t, func() { + handler.ServeHTTP(response, request) + }) + return response +} diff --git a/admin/billing/payment/stripe_webhook.go b/admin/billing/payment/stripe_webhook.go index cff72b0b38a4..c1893df0f2ea 100644 --- a/admin/billing/payment/stripe_webhook.go +++ b/admin/billing/payment/stripe_webhook.go @@ -58,8 +58,8 @@ func (s *stripeWebhook) handleWebhook(w http.ResponseWriter, r *http.Request) er if err := json.Unmarshal(event.Data.Raw, &paymentMethod); err != nil { return httputil.Errorf(http.StatusBadRequest, "error parsing payment method data: %w", err) } - if cust, ok := event.Data.PreviousAttributes["customer"]; ok && cust != nil { - err = s.handlePaymentMethodRemoved(r.Context(), event.ID, cust.(string), &paymentMethod) + if customerID, ok := event.Data.PreviousAttributes["customer"].(string); ok { + err = s.handlePaymentMethodRemoved(r.Context(), event.ID, customerID, &paymentMethod) if err != nil { return httputil.Errorf(http.StatusInternalServerError, "error handling payment_method.detached event: %w", err) } diff --git a/admin/billing/payment/stripe_webhook_test.go b/admin/billing/payment/stripe_webhook_test.go new file mode 100644 index 000000000000..096b092395ee --- /dev/null +++ b/admin/billing/payment/stripe_webhook_test.go @@ -0,0 +1,336 @@ +package payment + +import ( + "bytes" + "context" + "encoding/json" + "errors" + "net/http" + "net/http/httptest" + "testing" + "time" + + "github.com/rilldata/rill/admin/jobs" + "github.com/stretchr/testify/require" + "github.com/stripe/stripe-go/v79/webhook" + "go.uber.org/zap" +) + +const testStripeWebhookSecret = "whsec_bill_002" + +type paymentMethodAddedCall struct { + methodID string + paymentCustomerID string + typ string + eventTime time.Time +} + +type paymentMethodRemovedCall struct { + methodID string + paymentCustomerID string + eventTime time.Time +} + +type customerAddressUpdatedCall struct { + paymentCustomerID string + eventTime time.Time +} + +type stripeWebhookJobs struct { + jobs.Client + + paymentMethodAddedCalls []paymentMethodAddedCall + paymentMethodRemovedCalls []paymentMethodRemovedCall + customerAddressUpdatedCalls []customerAddressUpdatedCall + paymentMethodAddedErr error + paymentMethodRemovedErr error + customerAddressUpdatedErr error +} + +func (j *stripeWebhookJobs) PaymentMethodAdded(_ context.Context, methodID, paymentCustomerID, typ string, eventTime time.Time) (*jobs.InsertResult, error) { + j.paymentMethodAddedCalls = append(j.paymentMethodAddedCalls, paymentMethodAddedCall{ + methodID: methodID, + paymentCustomerID: paymentCustomerID, + typ: typ, + eventTime: eventTime, + }) + return &jobs.InsertResult{}, j.paymentMethodAddedErr +} + +func (j *stripeWebhookJobs) PaymentMethodRemoved(_ context.Context, methodID, paymentCustomerID string, eventTime time.Time) (*jobs.InsertResult, error) { + j.paymentMethodRemovedCalls = append(j.paymentMethodRemovedCalls, paymentMethodRemovedCall{ + methodID: methodID, + paymentCustomerID: paymentCustomerID, + eventTime: eventTime, + }) + return &jobs.InsertResult{}, j.paymentMethodRemovedErr +} + +func (j *stripeWebhookJobs) CustomerAddressUpdated(_ context.Context, paymentCustomerID string, eventTime time.Time) (*jobs.InsertResult, error) { + j.customerAddressUpdatedCalls = append(j.customerAddressUpdatedCalls, customerAddressUpdatedCall{ + paymentCustomerID: paymentCustomerID, + eventTime: eventTime, + }) + return &jobs.InsertResult{}, j.customerAddressUpdatedErr +} + +func (j *stripeWebhookJobs) callCount() int { + return len(j.paymentMethodAddedCalls) + len(j.paymentMethodRemovedCalls) + len(j.customerAddressUpdatedCalls) +} + +func TestStripeWebhookSignedEventsEnqueueParsedJobs(t *testing.T) { + // Stripe's signing helper keeps the tests on the same verification boundary as production requests. + tests := []struct { + name string + payload []byte + assertCall func(*testing.T, *stripeWebhookJobs) + }{ + { + name: "payment method attached", + payload: stripeEventPayload(t, map[string]any{ + "id": "evt_attached", + "type": "payment_method.attached", + "created": int64(1_717_171_717), + "data": map[string]any{"object": map[string]any{ + "id": "pm_attached", + "object": "payment_method", + "created": int64(1_700_000_111), + "customer": "cus_attached", + "type": "card", + }}, + }), + assertCall: func(t *testing.T, got *stripeWebhookJobs) { + require.Equal(t, []paymentMethodAddedCall{{ + methodID: "pm_attached", + paymentCustomerID: "cus_attached", + typ: "card", + eventTime: time.Unix(1_700_000_111, 0), + }}, got.paymentMethodAddedCalls) + }, + }, + { + name: "payment method detached", + payload: stripeEventPayload(t, map[string]any{ + "id": "evt_detached", + "type": "payment_method.detached", + "created": int64(1_717_171_718), + "data": map[string]any{ + "object": map[string]any{ + "id": "pm_detached", + "object": "payment_method", + "created": int64(1_700_000_222), + "type": "us_bank_account", + }, + "previous_attributes": map[string]any{"customer": "cus_detached"}, + }, + }), + assertCall: func(t *testing.T, got *stripeWebhookJobs) { + require.Equal(t, []paymentMethodRemovedCall{{ + methodID: "pm_detached", + paymentCustomerID: "cus_detached", + eventTime: time.Unix(1_700_000_222, 0), + }}, got.paymentMethodRemovedCalls) + }, + }, + { + name: "customer address updated", + payload: stripeEventPayload(t, map[string]any{ + "id": "evt_customer", + "type": "customer.updated", + "created": int64(1_700_000_333), + "data": map[string]any{ + "object": map[string]any{ + "id": "cus_updated", + "object": "customer", + "address": map[string]any{ + "country": "US", + "line1": "123 Main Street", + }, + }, + "previous_attributes": map[string]any{"address": nil}, + }, + }), + assertCall: func(t *testing.T, got *stripeWebhookJobs) { + require.Equal(t, []customerAddressUpdatedCall{{ + paymentCustomerID: "cus_updated", + eventTime: time.Unix(1_700_000_333, 0), + }}, got.customerAddressUpdatedCalls) + }, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + // Every recognized, authenticated event must cross into the jobs system exactly once with its parsed identifiers and timestamp. + jobClient := &stripeWebhookJobs{} + response := serveStripeWebhook(t, jobClient, tt.payload, signedStripeHeader(tt.payload, testStripeWebhookSecret)) + + require.Equal(t, http.StatusOK, response.Code) + require.Equal(t, 1, jobClient.callCount()) + tt.assertCall(t, jobClient) + }) + } +} + +func TestStripeWebhookDetachedPreviousCustomerShapes(t *testing.T) { + // Detached payment methods omit their current customer, making the signed previous value the only usable association. + tests := []struct { + name string + previousCustomer any + wantCustomerID string + wantCalls int + }{ + {name: "string", previousCustomer: "cus_previous", wantCustomerID: "cus_previous", wantCalls: 1}, + {name: "null", previousCustomer: nil}, + {name: "object", previousCustomer: map[string]any{"id": "cus_object"}}, + {name: "number", previousCustomer: 42}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + // Webhook metadata is untrusted input even after authentication; schema drift must not panic or invent a customer ID. + payload := stripeEventPayload(t, map[string]any{ + "id": "evt_previous_customer_" + tt.name, + "type": "payment_method.detached", + "created": int64(1_700_000_400), + "data": map[string]any{ + "object": map[string]any{ + "id": "pm_previous_customer", + "object": "payment_method", + "created": int64(1_700_000_401), + }, + "previous_attributes": map[string]any{"customer": tt.previousCustomer}, + }, + }) + jobClient := &stripeWebhookJobs{} + response := serveStripeWebhook(t, jobClient, payload, signedStripeHeader(payload, testStripeWebhookSecret)) + + require.Equal(t, http.StatusOK, response.Code) + require.Equal(t, tt.wantCalls, jobClient.callCount()) + if tt.wantCalls == 1 { + require.Equal(t, tt.wantCustomerID, jobClient.paymentMethodRemovedCalls[0].paymentCustomerID) + } + }) + } +} + +func TestStripeWebhookRejectsUntrustedOrInvalidBodies(t *testing.T) { + // Rejected requests must stop at the HTTP boundary and never reach any asynchronous billing work. + t.Run("bad signature", func(t *testing.T) { + // A syntactically valid event signed with another secret must be indistinguishable from unauthenticated input. + payload := stripeEventPayload(t, map[string]any{ + "id": "evt_bad_signature", "type": "payment_method.attached", "created": int64(1_700_000_500), + "data": map[string]any{"object": map[string]any{ + "id": "pm_bad_signature", "object": "payment_method", "customer": "cus_bad_signature", "type": "card", + }}, + }) + jobClient := &stripeWebhookJobs{} + response := serveStripeWebhook(t, jobClient, payload, signedStripeHeader(payload, "whsec_wrong")) + + require.Equal(t, http.StatusBadRequest, response.Code) + require.Zero(t, jobClient.callCount()) + }) + + t.Run("signed malformed JSON", func(t *testing.T) { + // Signing malformed bytes proves authentication alone cannot bypass event decoding. + payload := []byte(`{"id":"evt_malformed"`) + jobClient := &stripeWebhookJobs{} + response := serveStripeWebhook(t, jobClient, payload, signedStripeHeader(payload, testStripeWebhookSecret)) + + require.Equal(t, http.StatusBadRequest, response.Code) + require.Zero(t, jobClient.callCount()) + }) + + t.Run("body exceeds limit", func(t *testing.T) { + // One byte beyond the explicit cap verifies oversized bodies fail before signature parsing or enqueueing. + payload := bytes.Repeat([]byte("x"), int(maxBodyBytes)+1) + jobClient := &stripeWebhookJobs{} + response := serveStripeWebhook(t, jobClient, payload, signedStripeHeader(payload, testStripeWebhookSecret)) + + require.Equal(t, http.StatusServiceUnavailable, response.Code) + require.Zero(t, jobClient.callCount()) + }) +} + +func TestStripeWebhookReturnsServerErrorWhenEnqueueFails(t *testing.T) { + // A failed handoff must not be acknowledged, allowing Stripe to retry without a duplicate enqueue in this request. + tests := []struct { + name string + payload []byte + configure func(*stripeWebhookJobs) + }{ + { + name: "payment method attached", + payload: stripeEventPayload(t, map[string]any{ + "id": "evt_failed_attached", "type": "payment_method.attached", "created": int64(1_700_000_600), + "data": map[string]any{"object": map[string]any{ + "id": "pm_failed_attached", "object": "payment_method", "created": int64(1_700_000_601), "customer": "cus_failed_attached", "type": "card", + }}, + }), + configure: func(j *stripeWebhookJobs) { j.paymentMethodAddedErr = errors.New("queue unavailable") }, + }, + { + name: "payment method detached", + payload: stripeEventPayload(t, map[string]any{ + "id": "evt_failed_detached", "type": "payment_method.detached", "created": int64(1_700_000_610), + "data": map[string]any{ + "object": map[string]any{"id": "pm_failed_detached", "object": "payment_method", "created": int64(1_700_000_611)}, + "previous_attributes": map[string]any{"customer": "cus_failed_detached"}, + }, + }), + configure: func(j *stripeWebhookJobs) { j.paymentMethodRemovedErr = errors.New("queue unavailable") }, + }, + { + name: "customer address updated", + payload: stripeEventPayload(t, map[string]any{ + "id": "evt_failed_customer", "type": "customer.updated", "created": int64(1_700_000_620), + "data": map[string]any{ + "object": map[string]any{"id": "cus_failed", "object": "customer", "address": map[string]any{"country": "US"}}, + "previous_attributes": map[string]any{"address": nil}, + }, + }), + configure: func(j *stripeWebhookJobs) { j.customerAddressUpdatedErr = errors.New("queue unavailable") }, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + // Each job path gets one handoff attempt; returning 500 preserves Stripe's retry signal when that attempt fails. + jobClient := &stripeWebhookJobs{} + tt.configure(jobClient) + response := serveStripeWebhook(t, jobClient, tt.payload, signedStripeHeader(tt.payload, testStripeWebhookSecret)) + + require.Equal(t, http.StatusInternalServerError, response.Code) + require.Equal(t, 1, jobClient.callCount()) + }) + } +} + +func stripeEventPayload(t *testing.T, event map[string]any) []byte { + t.Helper() + event["object"] = "event" + payload, err := json.Marshal(event) + require.NoError(t, err) + return payload +} + +func signedStripeHeader(payload []byte, secret string) string { + return webhook.GenerateTestSignedPayload(&webhook.UnsignedPayload{ + Payload: payload, + Secret: secret, + }).Header +} + +func serveStripeWebhook(t *testing.T, jobClient jobs.Client, payload []byte, signature string) *httptest.ResponseRecorder { + t.Helper() + handler := NewStripe(zap.NewNop(), "", testStripeWebhookSecret).WebhookHandlerFunc(context.Background(), jobClient) + require.NotNil(t, handler) + + request := httptest.NewRequest(http.MethodPost, "/stripe/webhook", bytes.NewReader(payload)) + request.Header.Set("Stripe-Signature", signature) + response := httptest.NewRecorder() + require.NotPanics(t, func() { + handler.ServeHTTP(response, request) + }) + return response +} diff --git a/admin/jobs/river/billing_reporter.go b/admin/jobs/river/billing_reporter.go index 16893e96daf0..354871a633a8 100644 --- a/admin/jobs/river/billing_reporter.go +++ b/admin/jobs/river/billing_reporter.go @@ -2,11 +2,13 @@ package river import ( "context" + "errors" "fmt" "time" "github.com/rilldata/rill/admin" "github.com/rilldata/rill/admin/billing" + "github.com/rilldata/rill/admin/metrics" "github.com/riverqueue/river" "go.uber.org/zap" ) @@ -28,15 +30,15 @@ var counterMetrics = map[string]bool{ // reporter alongside the runtime-derived metrics for every org that reported usage (see Work). type orgUsageMetric struct { name string - collect func(ctx context.Context, adm *admin.Service, orgID string) (float64, error) + collect func(ctx context.Context, db billingReporterDatabase, orgID string) (float64, error) } var orgUsageMetrics = []orgUsageMetric{ { name: "seats", - collect: func(ctx context.Context, adm *admin.Service, orgID string) (float64, error) { + collect: func(ctx context.Context, db billingReporterDatabase, orgID string) (float64, error) { // Exclude internal Rill users from billable seat counts. - n, err := adm.DB.CountOrganizationMemberUsers(ctx, orgID, "", "%@"+billing.InternalEmailDomain, true) + n, err := db.CountOrganizationMemberUsers(ctx, orgID, "", "%@"+billing.InternalEmailDomain, true) if err != nil { return 0, err } @@ -51,12 +53,50 @@ func (BillingReporterArgs) Kind() string { return "billing_reporter" } type BillingReporterWorker struct { river.WorkerDefaults[BillingReporterArgs] - admin *admin.Service - logger *zap.Logger + admin *admin.Service + logger *zap.Logger + database billingReporterDatabase + openMetricsProject func(context.Context) (billingUsageMetricsClient, bool, error) + now func() time.Time +} + +type billingUsageMetricsClient interface { + GetUsageMetrics(ctx context.Context, startTime, endTime, afterTime time.Time, afterOrgID, afterProjectID, afterInstanceID, afterBillingService, afterEventName, grain string, limit int) ([]*metrics.Usage, error) +} + +type billingReporterDatabase interface { + FindBillingUsageReportedOn(ctx context.Context) (time.Time, error) + UpdateBillingUsageReportedOn(ctx context.Context, usageReportedOn time.Time) error + CountOrganizationMemberUsers(ctx context.Context, orgID, filterRoleID, searchPattern string, negateSearch bool) (int, error) + FindOrganizationIDsWithBilling(ctx context.Context) ([]string, error) + CountBillingProjectsForOrganization(ctx context.Context, orgID string, createdBefore time.Time) (int, error) +} + +func (w *BillingReporterWorker) db() billingReporterDatabase { + if w.database != nil { + return w.database + } + return w.admin.DB +} + +func (w *BillingReporterWorker) metricsClient(ctx context.Context) (billingUsageMetricsClient, bool, error) { + if w.openMetricsProject != nil { + return w.openMetricsProject(ctx) + } + client, ok, err := w.admin.OpenMetricsProject(ctx) + return client, ok, err +} + +func (w *BillingReporterWorker) currentTime() time.Time { + if w.now != nil { + return w.now() + } + return time.Now() } // NewBillingReporterWorker creates a new worker that reports billing information. func (w *BillingReporterWorker) Work(ctx context.Context, job *river.Job[BillingReporterArgs]) error { + db := w.db() // Get reporting granularity var granularity time.Duration var sqlGrainIdentifier string @@ -73,13 +113,13 @@ func (w *BillingReporterWorker) Work(ctx context.Context, job *river.Job[Billing return fmt.Errorf("unsupported reporting granularity: %s", w.admin.Biller.GetReportingGranularity()) } - t, err := w.admin.DB.FindBillingUsageReportedOn(ctx) + t, err := db.FindBillingUsageReportedOn(ctx) if err != nil { return fmt.Errorf("failed to get last usage reporting time: %w", err) } // after going back by grace period, report until the "end" of the last grain period - endTime := time.Now().UTC().Add(-gracePeriod).Truncate(granularity) + endTime := w.currentTime().UTC().Add(-gracePeriod).Truncate(granularity) // start reporting from the last reported time or from the "start" of the last grain period for first time reporting var startTime time.Time @@ -95,7 +135,7 @@ func (w *BillingReporterWorker) Work(ctx context.Context, job *river.Job[Billing } // Get metrics client - client, ok, err := w.admin.OpenMetricsProject(ctx) + client, ok, err := w.metricsClient(ctx) if err != nil { w.logger.Error("failed to report usage: unable to get metrics client", zap.Error(err)) return err @@ -173,7 +213,7 @@ func (w *BillingReporterWorker) Work(ctx context.Context, job *river.Job[Billing if afterTime.After(checkPoint) { checkPoint = afterTime - err = w.admin.DB.UpdateBillingUsageReportedOn(ctx, checkPoint) + err = db.UpdateBillingUsageReportedOn(ctx, checkPoint) if err != nil { return fmt.Errorf("failed to update last usage reporting time: %w", err) } @@ -190,18 +230,20 @@ func (w *BillingReporterWorker) Work(ctx context.Context, job *river.Job[Billing return fmt.Errorf("failed to update last usage reporting time: max end time not updated after reporting usage data") } - err = w.admin.DB.UpdateBillingUsageReportedOn(ctx, maxEndTime) + // Report admin-database-derived billable metrics (e.g. seats) for every org that reported usage in this run. + // These are current gauges, so we report them for the last completed grain period. They must be durable before the final checkpoint or a retry can skip them permanently. + if err := w.reportOrgUsageMetrics(ctx, db, reportedOrgs, endTime.Add(-granularity), endTime); err != nil { + return fmt.Errorf("failed to report admin usage metrics: %w", err) + } + + err = db.UpdateBillingUsageReportedOn(ctx, maxEndTime) if err != nil { return fmt.Errorf("failed to update last usage reporting time: %w", err) } - // Report admin-database-derived billable metrics (e.g. seats) for every org that reported usage in this run. - // These are current gauges, so we report them for the last completed grain period; the biller aggregates over the billing period. - w.reportOrgUsageMetrics(ctx, reportedOrgs, endTime.Add(-granularity), endTime) - // TODO move the validation to background job // get orgs which have billing customer id - orgs, err := w.admin.DB.FindOrganizationIDsWithBilling(ctx) + orgs, err := db.FindOrganizationIDsWithBilling(ctx) if err != nil { return fmt.Errorf("failed to report usage: unable to fetch orgs: %w", err) } @@ -210,7 +252,7 @@ func (w *BillingReporterWorker) Work(ctx context.Context, job *river.Job[Billing for _, org := range orgs { if _, ok := reportedOrgs[org]; !ok { // count the projects which are not hibernated and created before the given time - count, err := w.admin.DB.CountBillingProjectsForOrganization(ctx, org, endTime) + count, err := db.CountBillingProjectsForOrganization(ctx, org, endTime) if err != nil { w.logger.Warn("failed to validate active projects for org", zap.String("org_id", org), zap.Error(err)) continue @@ -224,15 +266,17 @@ func (w *BillingReporterWorker) Work(ctx context.Context, job *river.Job[Billing } // reportOrgUsageMetrics reports the admin-database-derived billable metrics (orgUsageMetrics) for the given orgs. -// It is best-effort: a failure for one org or metric is logged and skipped so the rest still get reported. -func (w *BillingReporterWorker) reportOrgUsageMetrics(ctx context.Context, reportedOrgs map[string]string, startTime, endTime time.Time) { +// It attempts every org and returns all failures so the caller can withhold the final checkpoint and safely retry. +func (w *BillingReporterWorker) reportOrgUsageMetrics(ctx context.Context, db billingReporterDatabase, reportedOrgs map[string]string, startTime, endTime time.Time) error { grain := w.admin.Biller.GetReportingGranularity() + var reportErr error for orgID, customerID := range reportedOrgs { var usage []*billing.Usage for _, m := range orgUsageMetrics { - val, err := m.collect(ctx, w.admin, orgID) + val, err := m.collect(ctx, db, orgID) if err != nil { w.logger.Error("failed to collect admin usage metric", zap.String("metric", m.name), zap.String("org_id", orgID), zap.Error(err)) + reportErr = errors.Join(reportErr, fmt.Errorf("collect %s for org %s: %w", m.name, orgID, err)) continue } usage = append(usage, &billing.Usage{ @@ -250,6 +294,8 @@ func (w *BillingReporterWorker) reportOrgUsageMetrics(ctx context.Context, repor } if err := w.admin.Biller.ReportUsage(ctx, usage); err != nil { w.logger.Error("failed to report admin usage metrics", zap.String("org_id", orgID), zap.Error(err)) + reportErr = errors.Join(reportErr, fmt.Errorf("report org %s: %w", orgID, err)) } } + return reportErr } diff --git a/admin/jobs/river/billing_reporter_test.go b/admin/jobs/river/billing_reporter_test.go new file mode 100644 index 000000000000..ecce445d4d4e --- /dev/null +++ b/admin/jobs/river/billing_reporter_test.go @@ -0,0 +1,318 @@ +package river + +import ( + "context" + "errors" + "fmt" + "testing" + "time" + + "github.com/rilldata/rill/admin" + "github.com/rilldata/rill/admin/billing" + "github.com/rilldata/rill/admin/metrics" + "github.com/stretchr/testify/require" + "go.uber.org/zap" +) + +func TestBillingReporterPaginationBoundariesAndTupleCursor(t *testing.T) { + // The metrics API uses a six-column cursor. Exact page boundaries must neither + // skip the first row of the next page nor loop when timestamps are identical. + tests := []struct { + rows int + wantCalls int + }{ + {rows: 0, wantCalls: 1}, + {rows: 9_999, wantCalls: 1}, + {rows: 10_000, wantCalls: 2}, + {rows: 10_001, wantCalls: 2}, + } + for _, tt := range tests { + t.Run(fmt.Sprintf("%d rows", tt.rows), func(t *testing.T) { + worker, db, biller, client := newBillingReporterFixture(makeBillingMetricRows(tt.rows)) + + err := worker.Work(t.Context(), nil) + require.NoError(t, err) + require.Len(t, client.calls, tt.wantCalls) + require.Equal(t, tt.rows, countReportedMetric(biller.delivered, "query")) + + if tt.rows >= 10_000 { + last := client.rows[9_999] + second := client.calls[1] + require.Equal(t, last.StartTime, second.afterTime) + require.Equal(t, last.OrgID, second.afterOrgID) + require.Equal(t, last.ProjectID, second.afterProjectID) + require.Equal(t, last.InstanceID, second.afterInstanceID) + require.Equal(t, last.BillingService, second.afterBillingService) + require.Equal(t, last.EventName, second.afterEventName) + } + if tt.rows == 0 { + require.True(t, db.checkpoint.IsZero()) + } else { + require.Equal(t, client.rows[len(client.rows)-1].EndTime, db.checkpoint) + require.Equal(t, 1, countReportedMetric(biller.delivered, "seats")) + } + }) + } +} + +func TestBillingReporterUsesSumForCountersAndMaxForGauges(t *testing.T) { + // Counter metrics conserve all events in a grain, while gauges represent the + // peak. Swapping these silently over- or under-bills customers. + rows := makeBillingMetricRows(2) + rows[0].EventName = "query" + rows[0].SumValue = 17 + rows[0].MaxValue = 99 + rows[1].EventName = "storage_bytes" + rows[1].SumValue = 1_000 + rows[1].MaxValue = 250 + worker, _, biller, _ := newBillingReporterFixture(rows) + + require.NoError(t, worker.Work(t.Context(), nil)) + runtimeUsage := findBillingUsageBatch(t, biller.delivered, "query") + require.Len(t, runtimeUsage, 2) + require.EqualValues(t, 17, runtimeUsage[0].Value) + require.EqualValues(t, 250, runtimeUsage[1].Value) + require.Equal(t, "org-1", runtimeUsage[0].Metadata["org_id"]) +} + +func TestBillingReporterFailurePathsDoNotAdvanceFinalCheckpoint(t *testing.T) { + // A checkpoint is a durability claim. Fetch, provider, seat collection, seat + // delivery, and checkpoint failures must all leave the grain retryable. + tests := []struct { + name string + configure func(*fakeBillingReporterDB, *recordingBillingBiller, *fakeBillingMetricsClient) + wantError string + }{ + { + name: "metrics fetch failure", + configure: func(_ *fakeBillingReporterDB, _ *recordingBillingBiller, client *fakeBillingMetricsClient) { + client.fetchErr = errors.New("metrics unavailable") + }, + wantError: "failed to get usage metrics", + }, + { + name: "provider delivery failure", + configure: func(_ *fakeBillingReporterDB, biller *recordingBillingBiller, _ *fakeBillingMetricsClient) { + biller.failRuntime = errors.New("provider unavailable") + }, + wantError: "failed to report usage", + }, + { + name: "seat collection failure", + configure: func(db *fakeBillingReporterDB, _ *recordingBillingBiller, _ *fakeBillingMetricsClient) { + db.seatErr = errors.New("member query failed") + }, + wantError: "failed to report admin usage metrics", + }, + { + name: "seat delivery failure", + configure: func(_ *fakeBillingReporterDB, biller *recordingBillingBiller, _ *fakeBillingMetricsClient) { + biller.failSeats = errors.New("seat delivery failed") + }, + wantError: "failed to report admin usage metrics", + }, + { + name: "checkpoint persistence failure", + configure: func(db *fakeBillingReporterDB, _ *recordingBillingBiller, _ *fakeBillingMetricsClient) { + db.updateFailures = 1 + }, + wantError: "failed to update last usage reporting time", + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + worker, db, biller, client := newBillingReporterFixture(makeBillingMetricRows(1)) + tt.configure(db, biller, client) + + err := worker.Work(t.Context(), nil) + require.ErrorContains(t, err, tt.wantError) + require.True(t, db.checkpoint.IsZero()) + }) + } +} + +func TestBillingReporterRetryAfterCheckpointFailureReplaysStableUsage(t *testing.T) { + // Provider delivery can succeed immediately before checkpoint persistence + // fails. A retry must replay the same logical records so provider idempotency can deduplicate them. + worker, db, biller, _ := newBillingReporterFixture(makeBillingMetricRows(1)) + db.updateFailures = 1 + + require.ErrorContains(t, worker.Work(t.Context(), nil), "failed to update last usage reporting time") + require.NoError(t, worker.Work(t.Context(), nil)) + require.Len(t, biller.delivered, 4) // runtime + seats on each attempt + require.Equal(t, biller.delivered[0], biller.delivered[2]) + require.Equal(t, biller.delivered[1], biller.delivered[3]) + require.False(t, db.checkpoint.IsZero()) +} + +type billingMetricsCall struct { + startTime time.Time + endTime time.Time + afterTime time.Time + afterOrgID string + afterProjectID string + afterInstanceID string + afterBillingService string + afterEventName string + limit int +} + +type fakeBillingMetricsClient struct { + rows []*metrics.Usage + calls []billingMetricsCall + fetchErr error +} + +func (c *fakeBillingMetricsClient) GetUsageMetrics(_ context.Context, startTime, endTime, afterTime time.Time, afterOrgID, afterProjectID, afterInstanceID, afterBillingService, afterEventName, _ string, limit int) ([]*metrics.Usage, error) { + c.calls = append(c.calls, billingMetricsCall{ + startTime: startTime, endTime: endTime, afterTime: afterTime, + afterOrgID: afterOrgID, afterProjectID: afterProjectID, afterInstanceID: afterInstanceID, + afterBillingService: afterBillingService, afterEventName: afterEventName, limit: limit, + }) + if c.fetchErr != nil { + return nil, c.fetchErr + } + offset := 0 + if !afterTime.IsZero() { + offset = len(c.rows) + for i, row := range c.rows { + if row.StartTime.Equal(afterTime) && row.OrgID == afterOrgID && row.ProjectID == afterProjectID && row.InstanceID == afterInstanceID && row.BillingService == afterBillingService && row.EventName == afterEventName { + offset = i + 1 + break + } + } + } + if offset >= len(c.rows) { + return nil, nil + } + end := min(offset+limit, len(c.rows)) + return c.rows[offset:end], nil +} + +type fakeBillingReporterDB struct { + checkpoint time.Time + updates []time.Time + updateFailures int + seatCount int + seatErr error + organizationIDs []string +} + +func (d *fakeBillingReporterDB) FindBillingUsageReportedOn(context.Context) (time.Time, error) { + return d.checkpoint, nil +} + +func (d *fakeBillingReporterDB) UpdateBillingUsageReportedOn(_ context.Context, checkpoint time.Time) error { + d.updates = append(d.updates, checkpoint) + if d.updateFailures > 0 { + d.updateFailures-- + return errors.New("checkpoint unavailable") + } + d.checkpoint = checkpoint + return nil +} + +func (d *fakeBillingReporterDB) CountOrganizationMemberUsers(context.Context, string, string, string, bool) (int, error) { + return d.seatCount, d.seatErr +} + +func (d *fakeBillingReporterDB) FindOrganizationIDsWithBilling(context.Context) ([]string, error) { + return d.organizationIDs, nil +} + +func (d *fakeBillingReporterDB) CountBillingProjectsForOrganization(context.Context, string, time.Time) (int, error) { + return 0, nil +} + +type recordingBillingBiller struct { + billing.Biller + delivered [][]*billing.Usage + failRuntime error + failSeats error +} + +func (b *recordingBillingBiller) GetReportingGranularity() billing.UsageReportingGranularity { + return billing.UsageReportingGranularityHour +} + +func (b *recordingBillingBiller) ReportUsage(_ context.Context, usage []*billing.Usage) error { + if len(usage) > 0 && usage[0].MetricName == "seats" && b.failSeats != nil { + return b.failSeats + } + if len(usage) > 0 && usage[0].MetricName != "seats" && b.failRuntime != nil { + return b.failRuntime + } + cloned := make([]*billing.Usage, len(usage)) + for i, item := range usage { + copy := *item + copy.Metadata = make(map[string]interface{}, len(item.Metadata)) + for key, value := range item.Metadata { + copy.Metadata[key] = value + } + cloned[i] = © + } + b.delivered = append(b.delivered, cloned) + return nil +} + +func newBillingReporterFixture(rows []*metrics.Usage) (*BillingReporterWorker, *fakeBillingReporterDB, *recordingBillingBiller, *fakeBillingMetricsClient) { + db := &fakeBillingReporterDB{seatCount: 3} + biller := &recordingBillingBiller{} + client := &fakeBillingMetricsClient{rows: rows} + worker := &BillingReporterWorker{ + admin: &admin.Service{Biller: biller}, + logger: zap.NewNop(), + database: db, + openMetricsProject: func(context.Context) (billingUsageMetricsClient, bool, error) { + return client, true, nil + }, + now: func() time.Time { + return time.Date(2026, time.January, 2, 5, 55, 0, 0, time.UTC) + }, + } + return worker, db, biller, client +} + +func makeBillingMetricRows(count int) []*metrics.Usage { + rows := make([]*metrics.Usage, count) + billingCustomerID := "customer-1" + for i := range rows { + rows[i] = &metrics.Usage{ + OrgID: "org-1", + ProjectID: fmt.Sprintf("project-%05d", i), + ProjectName: "Project", + BillingCustomerID: &billingCustomerID, + StartTime: time.Date(2026, time.January, 2, 3, 0, 0, 0, time.UTC), + EndTime: time.Date(2026, time.January, 2, 4, 0, 0, 0, time.UTC), + EventName: "query", + MaxValue: 1, + SumValue: 1, + BillingService: "runtime", + InstanceID: fmt.Sprintf("instance-%05d", i), + } + } + return rows +} + +func countReportedMetric(batches [][]*billing.Usage, metric string) int { + count := 0 + for _, batch := range batches { + for _, usage := range batch { + if usage.MetricName == metric { + count++ + } + } + } + return count +} + +func findBillingUsageBatch(t *testing.T, batches [][]*billing.Usage, metric string) []*billing.Usage { + t.Helper() + for _, batch := range batches { + if len(batch) > 0 && batch[0].MetricName == metric { + return batch + } + } + require.FailNow(t, "usage batch not found", metric) + return nil +}