1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
|
package controlplane
import (
"context"
"encoding/base64"
"fmt"
"net/url"
"os"
"strings"
"testing"
"time"
"github.com/jackc/pgx/v5/pgxpool"
)
func TestBillingNotificationsUseTenantPreferencesLedgerAndEncryptedOutboxPostgres(t *testing.T) {
databaseURL := os.Getenv("AIGW_TEST_DATABASE_URL")
if databaseURL == "" {
t.Skip("AIGW_TEST_DATABASE_URL is not set")
}
ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
defer cancel()
rootDB, err := pgxpool.New(ctx, databaseURL)
if err != nil {
t.Fatal(err)
}
defer rootDB.Close()
schema := fmt.Sprintf("mail_notifications_%d", time.Now().UnixNano())
if _, err := rootDB.Exec(ctx, "CREATE SCHEMA "+schema); err != nil {
t.Fatal(err)
}
t.Cleanup(func() { _, _ = rootDB.Exec(context.Background(), "DROP SCHEMA "+schema+" CASCADE") })
parsed, err := url.Parse(databaseURL)
if err != nil {
t.Fatal(err)
}
query := parsed.Query()
query.Set("search_path", schema)
parsed.RawQuery = query.Encode()
isolatedURL := parsed.String()
if err := MigrateDatabase(ctx, isolatedURL); err != nil {
t.Fatal(err)
}
credentialKey := base64.StdEncoding.EncodeToString([]byte("01234567890123456789012345678901"))
store, err := NewStore(ctx, Options{DatabaseURL: isolatedURL, CredentialKey: credentialKey})
if err != nil {
t.Fatal(err)
}
t.Cleanup(func() { _ = store.Close() })
tenant, _, err := store.CreateTenant(ctx, CreateTenantInput{Slug: "mail-alert-test", Name: "Mail Alert Test"})
if err != nil {
t.Fatal(err)
}
if _, err := store.db.Exec(ctx, `INSERT INTO tenant_wallets (tenant_id,currency,balance_micros)
VALUES ($1,'usd',1000000)`, tenant.ID); err != nil {
t.Fatal(err)
}
enabled := true
threshold := int64(2_000_000)
anomalyMultiplier := int64(4)
anomalyMinimum := int64(400_000)
if _, err := store.SetBillingPreferences(ctx, SetBillingPreferencesInput{TenantID: tenant.ID,
LowBalanceEnabled: &enabled, LowBalanceThresholdMicros: &threshold, SpendAnomalyEnabled: &enabled,
SpendAnomalyMultiplier: &anomalyMultiplier, SpendAnomalyMinMicros: &anomalyMinimum},
BillingPreferenceDefaults{LowBalanceThresholdMicros: 5_000_000, SpendAnomalyMultiplier: 10, SpendAnomalyMinMicros: 900_000}); err != nil {
t.Fatal(err)
}
if _, err := store.db.Exec(ctx, `INSERT INTO console_users
(tenant_id,email,display_name,role,status,email_verified_at) VALUES
($1,'billing-alert@example.test','Billing Owner','tenant_billing','active',now()),
($1,'developer-no-alert@example.test','Developer','tenant_developer','active',now()),
($1,'unverified-no-alert@example.test','Unverified Billing','tenant_billing','active',NULL)`, tenant.ID); err != nil {
t.Fatal(err)
}
if _, err := store.db.Exec(ctx, `INSERT INTO billing_ledger
(tenant_id,currency,amount_micros,balance_after_micros,kind,source_type,source_id,description,created_at)
SELECT $1,'usd',-100000,1000000,'usage','request','historical-'||day::text,'Historical usage',
date_trunc('day',now())-make_interval(days=>day)
FROM generate_series(1,7) AS day`, tenant.ID); err != nil {
t.Fatal(err)
}
if _, err := store.db.Exec(ctx, `INSERT INTO billing_ledger
(tenant_id,currency,amount_micros,balance_after_micros,kind,source_type,source_id,description,created_at)
VALUES ($1,'usd',-500000,1000000,'usage','request','today-spend','Today usage',now())`, tenant.ID); err != nil {
t.Fatal(err)
}
config := MailNotificationConfig{LowBalanceMicros: 5_000_000, SpendAnomalyMultiplier: 10, SpendAnomalyMinMicros: 900_000}
if err := store.queueBillingNotifications(ctx, config); err != nil {
t.Fatal(err)
}
if err := store.queueBillingNotifications(ctx, config); err != nil {
t.Fatal(err)
}
var billingMessages, otherMessages, notificationEvents int
if err := store.db.QueryRow(ctx, `SELECT
count(*) FILTER (WHERE recipient='billing-alert@example.test'),
count(*) FILTER (WHERE recipient<>'billing-alert@example.test')
FROM console_mail_outbox`).Scan(&billingMessages, &otherMessages); err != nil {
t.Fatal(err)
}
if err := store.db.QueryRow(ctx, `SELECT count(*) FROM mail_notification_events`).Scan(¬ificationEvents); err != nil {
t.Fatal(err)
}
if billingMessages != 2 || otherMessages != 0 || notificationEvents != 2 {
t.Fatalf("notification dedupe or recipient filtering failed: billing=%d other=%d events=%d", billingMessages, otherMessages, notificationEvents)
}
var plaintextLeaks int
if err := store.db.QueryRow(ctx, `SELECT count(*) FROM console_mail_outbox
WHERE convert_from(body_ciphertext,'UTF8') LIKE '%1.000000 USD%'`).Scan(&plaintextLeaks); err == nil {
if plaintextLeaks != 0 {
t.Fatal("notification body was stored as plaintext")
}
} else {
// Authenticated encryption output is arbitrary bytes and usually is not valid UTF-8.
var containsPlaintext bool
if scanErr := store.db.QueryRow(ctx, `SELECT bool_or(position(convert_to('1.000000 USD','UTF8') in body_ciphertext)>0)
FROM console_mail_outbox`).Scan(&containsPlaintext); scanErr != nil {
t.Fatal(scanErr)
}
if containsPlaintext {
t.Fatal("notification body was stored as plaintext")
}
}
bodies := make([]string, 0, 2)
for range 2 {
message, ok, err := store.ClaimMail(ctx)
if err != nil {
t.Fatal(err)
}
if !ok || message.Recipient != "billing-alert@example.test" {
t.Fatalf("unexpected claimed notification: ok=%v message=%+v", ok, message)
}
bodies = append(bodies, message.Body)
}
joined := strings.Join(bodies, "\n")
for _, expected := range []string{"1.000000 USD", "0.500000 USD", "0.100000 USD"} {
if !strings.Contains(joined, expected) {
t.Fatalf("decrypted notifications do not contain %q: %s", expected, joined)
}
}
}
|
