| 1 | package main |
| 2 | |
| 3 | import ( |
| 4 | "encoding/json" |
| 5 | "net/http" |
| 6 | "os" |
| 7 | "path/filepath" |
| 8 | "sync" |
| 9 | "testing" |
| 10 | "time" |
| 11 | ) |
| 12 | |
| 13 | func TestStorageBytesForDB(t *testing.T) { |
| 14 | dir := t.TempDir() |
| 15 | db := filepath.Join(dir, "space.11ch.at.db") |
| 16 | write := func(path string, n int) { |
| 17 | t.Helper() |
| 18 | if err := os.MkdirAll(filepath.Dir(path), 0o700); err != nil { |
| 19 | t.Fatal(err) |
| 20 | } |
| 21 | if err := os.WriteFile(path, make([]byte, n), 0o600); err != nil { |
| 22 | t.Fatal(err) |
| 23 | } |
| 24 | } |
| 25 | write(db, 1000) |
| 26 | write(db+"-wal", 200) |
| 27 | write(filepath.Join(dir, "space.11ch.at-blobs", "aa"), 5000) |
| 28 | write(filepath.Join(dir, "space.11ch.at-blobs", "nested", "bb"), 300) |
| 29 | // A sibling instance's files must not leak into the count. |
| 30 | write(filepath.Join(dir, "other.11ch.at.db"), 9999) |
| 31 | |
| 32 | if got := storageBytesForDB(db); got != 6500 { |
| 33 | t.Errorf("storageBytesForDB = %d, want 6500 (db+wal+blobs)", got) |
| 34 | } |
| 35 | // Missing everything → 0, no error. |
| 36 | if got := storageBytesForDB(filepath.Join(dir, "absent.db")); got != 0 { |
| 37 | t.Errorf("absent instance = %d, want 0", got) |
| 38 | } |
| 39 | } |
| 40 | |
| 41 | func TestMeterOnceRecordsReadings(t *testing.T) { |
| 42 | db := openTestCloudDB(t) |
| 43 | idA := insertSpaceRow(t, db, "a.11ch.at", 1000) |
| 44 | idB := insertSpaceRow(t, db, "b.11ch.at", 1000) |
| 45 | // A pending create (NULL host) is never swept. |
| 46 | if _, err := db.sql.Exec(`INSERT INTO spaces(host, caller_id, requester, created_at) VALUES(NULL,'c','r',1)`); err != nil { |
| 47 | t.Fatal(err) |
| 48 | } |
| 49 | |
| 50 | bytesByHost := map[string]int64{"a.11ch.at": 12345, "b.11ch.at": 0} |
| 51 | mux := http.NewServeMux() |
| 52 | mux.HandleFunc("GET /instances/{host}/stats", func(w http.ResponseWriter, r *http.Request) { |
| 53 | h := r.PathValue("host") |
| 54 | if h == "b.11ch.at" { |
| 55 | w.WriteHeader(http.StatusInternalServerError) // stats down for b |
| 56 | return |
| 57 | } |
| 58 | writeJSON(w, http.StatusOK, map[string]any{"storage_bytes": bytesByHost[h], "setup_complete": true, "members": 3}) |
| 59 | }) |
| 60 | sock := fakeControl(t, mux.ServeHTTP) |
| 61 | s := &cloudServer{db: db, control: controlClient(sock), controlToken: "t"} |
| 62 | |
| 63 | if metered := s.meterOnce(); metered != 1 { |
| 64 | t.Fatalf("metered = %d, want 1 (a recorded, b skipped)", metered) |
| 65 | } |
| 66 | spaces, err := db.listSpaces("c", "r") |
| 67 | if err != nil { |
| 68 | t.Fatal(err) |
| 69 | } |
| 70 | byHost := map[string]Space{} |
| 71 | for _, sp := range spaces { |
| 72 | byHost[sp.Host] = sp |
| 73 | } |
| 74 | if got := byHost["a.11ch.at"]; got.StorageBytes != 12345 || got.MeteredAt == 0 || !got.SetupComplete || got.Members != 3 { |
| 75 | t.Errorf("a = %+v, want 12345 bytes, claimed-ness recorded, timestamp", got) |
| 76 | } |
| 77 | if got := byHost["b.11ch.at"]; got.StorageBytes != 0 || got.MeteredAt != 0 || got.SetupComplete { |
| 78 | t.Errorf("b = %+v, want untouched (stats error is a skip)", got) |
| 79 | } |
| 80 | _ = idA |
| 81 | _ = idB |
| 82 | |
| 83 | // A later sweep refreshes the reading, never accumulates. |
| 84 | bytesByHost["a.11ch.at"] = 200 |
| 85 | before := time.Now().Unix() |
| 86 | s.meterOnce() |
| 87 | spaces, _ = db.listSpaces("c", "r") |
| 88 | for _, sp := range spaces { |
| 89 | if sp.Host == "a.11ch.at" && (sp.StorageBytes != 200 || sp.MeteredAt < before) { |
| 90 | t.Errorf("refresh: %+v, want 200 bytes, fresh timestamp", sp) |
| 91 | } |
| 92 | } |
| 93 | } |
| 94 | |
| 95 | func TestParseByteSizeAndCaps(t *testing.T) { |
| 96 | for in, want := range map[string]int64{ |
| 97 | "25GB": 25e9, "2TB": 2e12, "500MB": 5e8, "1024": 1024, "10 GB": 10e9, "1kb": 1000, |
| 98 | } { |
| 99 | if got, err := parseByteSize(in); err != nil || got != want { |
| 100 | t.Errorf("parseByteSize(%q) = %d, %v; want %d", in, got, err, want) |
| 101 | } |
| 102 | } |
| 103 | for _, bad := range []string{"", "-5GB", "0", "GB", "ten GB"} { |
| 104 | if _, err := parseByteSize(bad); err == nil { |
| 105 | t.Errorf("parseByteSize(%q) accepted, want error", bad) |
| 106 | } |
| 107 | } |
| 108 | |
| 109 | caps, _, err := parseStorageCaps("free=10GB, family=25GB,business=2TB") |
| 110 | if err != nil || caps["free"] != 10e9 || caps["family"] != 25e9 || caps["business"] != 2e12 { |
| 111 | t.Fatalf("caps = %v, %v", caps, err) |
| 112 | } |
| 113 | if c, sc, err := parseStorageCaps(""); c != nil || sc != nil || err != nil { |
| 114 | t.Errorf("empty caps = %v, %v; want nil, nil (enforcement off)", c, err) |
| 115 | } |
| 116 | if _, _, err := parseStorageCaps("free=banana"); err == nil { |
| 117 | t.Error("garbage cap accepted") |
| 118 | } |
| 119 | } |
| 120 | |
| 121 | func TestSetInstanceQuota(t *testing.T) { |
| 122 | path := filepath.Join(t.TempDir(), "space.env") |
| 123 | if err := os.WriteFile(path, []byte("LCHAT_ADMIN_TOKEN=secret\nLCHAT_DB=/data/x.db\n"), 0o600); err != nil { |
| 124 | t.Fatal(err) |
| 125 | } |
| 126 | if err := setInstanceQuota(path, 25e9, "free"); err != nil { |
| 127 | t.Fatal(err) |
| 128 | } |
| 129 | env, err := readEnvFile(path) |
| 130 | if err != nil { |
| 131 | t.Fatal(err) |
| 132 | } |
| 133 | // The cap and entitlement are set and every other var survived. |
| 134 | if env["LCHAT_BLOB_QUOTA_BYTES"] != "25000000000" || env["LCHAT_ENTITLEMENT"] != "free" || |
| 135 | env["LCHAT_ADMIN_TOKEN"] != "secret" || env["LCHAT_DB"] != "/data/x.db" { |
| 136 | t.Fatalf("env after set = %v", env) |
| 137 | } |
| 138 | // Idempotent update; an empty entitlement leaves the flag alone. |
| 139 | if err := setInstanceQuota(path, 10e9, ""); err != nil { |
| 140 | t.Fatal(err) |
| 141 | } |
| 142 | env, _ = readEnvFile(path) |
| 143 | if env["LCHAT_BLOB_QUOTA_BYTES"] != "10000000000" || env["LCHAT_ENTITLEMENT"] != "free" { |
| 144 | t.Fatalf("env after update = %v", env) |
| 145 | } |
| 146 | // A billing flip rides the same write; removal at 0 spares the flag. |
| 147 | if err := setInstanceQuota(path, 0, "paid"); err != nil { |
| 148 | t.Fatal(err) |
| 149 | } |
| 150 | env, _ = readEnvFile(path) |
| 151 | if _, ok := env["LCHAT_BLOB_QUOTA_BYTES"]; ok { |
| 152 | t.Fatalf("cap not removed: %v", env) |
| 153 | } |
| 154 | if env["LCHAT_ENTITLEMENT"] != "paid" { |
| 155 | t.Fatalf("entitlement not flipped: %v", env) |
| 156 | } |
| 157 | // A missing instance surfaces not-exist (the handler's 404). |
| 158 | if err := setInstanceQuota(filepath.Join(t.TempDir(), "absent.env"), 1, ""); !os.IsNotExist(err) { |
| 159 | t.Errorf("absent env: err = %v, want not-exist", err) |
| 160 | } |
| 161 | } |
| 162 | |
| 163 | // quotaControl fakes controld's stats + quota endpoints, recording pushes. |
| 164 | func quotaControl(t *testing.T, bytesByHost map[string]int64, pushed map[string]int64) string { |
| 165 | t.Helper() |
| 166 | mux := http.NewServeMux() |
| 167 | mux.HandleFunc("GET /instances/{host}/stats", func(w http.ResponseWriter, r *http.Request) { |
| 168 | writeJSON(w, http.StatusOK, map[string]any{"storage_bytes": bytesByHost[r.PathValue("host")]}) |
| 169 | }) |
| 170 | mux.HandleFunc("POST /instances/{host}/quota", func(w http.ResponseWriter, r *http.Request) { |
| 171 | var req struct { |
| 172 | Bytes int64 `json:"bytes"` |
| 173 | } |
| 174 | json.NewDecoder(r.Body).Decode(&req) |
| 175 | pushed[r.PathValue("host")] = req.Bytes |
| 176 | writeJSON(w, http.StatusOK, map[string]any{"ok": true}) |
| 177 | }) |
| 178 | return fakeControl(t, mux.ServeHTTP) |
| 179 | } |
| 180 | |
| 181 | // TestQuotaPushOnUpgradeAndSweep pins the enforcement loop: a webhook upgrade |
| 182 | // pushes the tier's cap immediately; the meter sweep reconciles free spaces |
| 183 | // and downgrades. |
| 184 | func TestQuotaPushOnUpgradeAndSweep(t *testing.T) { |
| 185 | db := openTestCloudDB(t) |
| 186 | paidID := insertSpaceRow(t, db, "paid.11ch.at", 1000) |
| 187 | insertSpaceRow(t, db, "free.11ch.at", 1000) |
| 188 | |
| 189 | pushed := map[string]int64{} |
| 190 | sock := quotaControl(t, map[string]int64{"paid.11ch.at": 1, "free.11ch.at": 1}, pushed) |
| 191 | s := &cloudServer{ |
| 192 | db: db, control: controlClient(sock), controlToken: "t", |
| 193 | storageCaps: map[string]int64{"free": 10e9, "family": 25e9}, |
| 194 | } |
| 195 | |
| 196 | // The upgrade path (what the checkout webhook calls). |
| 197 | if err := db.markSpacePaid(paidID, "cus", "sub", "family", ""); err != nil { |
| 198 | t.Fatal(err) |
| 199 | } |
| 200 | s.pushQuotaByID(paidID) |
| 201 | if pushed["paid.11ch.at"] != 25e9 { |
| 202 | t.Fatalf("upgrade push = %d, want 25GB", pushed["paid.11ch.at"]) |
| 203 | } |
| 204 | |
| 205 | // The sweep reconciles the free space (never pushed) — and doesn't |
| 206 | // re-push the paid one (already at its cap). |
| 207 | delete(pushed, "paid.11ch.at") |
| 208 | s.meterOnce() |
| 209 | if pushed["free.11ch.at"] != 10e9 { |
| 210 | t.Fatalf("sweep push = %d, want 10GB", pushed["free.11ch.at"]) |
| 211 | } |
| 212 | if _, again := pushed["paid.11ch.at"]; again { |
| 213 | t.Error("sweep re-pushed an already-correct quota") |
| 214 | } |
| 215 | |
| 216 | // A lapse (subscription deleted → free) re-applies the free cap. |
| 217 | if err := db.setEntitlement(paidID, entFree, 0); err != nil { |
| 218 | t.Fatal(err) |
| 219 | } |
| 220 | s.pushQuotaByID(paidID) |
| 221 | if pushed["paid.11ch.at"] != 10e9 { |
| 222 | t.Fatalf("lapse push = %d, want back to 10GB", pushed["paid.11ch.at"]) |
| 223 | } |
| 224 | |
| 225 | // No caps configured → nothing pushes anywhere. |
| 226 | pushed2 := map[string]int64{} |
| 227 | sock2 := quotaControl(t, map[string]int64{"paid.11ch.at": 1, "free.11ch.at": 1}, pushed2) |
| 228 | s2 := &cloudServer{db: db, control: controlClient(sock2), controlToken: "t"} |
| 229 | s2.meterOnce() |
| 230 | s2.pushQuotaByID(paidID) |
| 231 | if len(pushed2) != 0 { |
| 232 | t.Errorf("unconfigured caps pushed: %v", pushed2) |
| 233 | } |
| 234 | } |
| 235 | |
| 236 | // The quota and meter pushes write DIFFERENT keys to the SAME env file from |
| 237 | // separate request goroutines; updateInstanceEnv serializes them so neither |
| 238 | // side's read-before-write snapshot can drop the other's keys. Hammered |
| 239 | // concurrently, every key family must survive. |
| 240 | func TestInstanceEnvWritesSerialized(t *testing.T) { |
| 241 | envPath := filepath.Join(t.TempDir(), "space.env") |
| 242 | if err := writeEnvFile(envPath, map[string]string{"LCHAT_ADMIN_TOKEN": "keepme"}); err != nil { |
| 243 | t.Fatal(err) |
| 244 | } |
| 245 | var wg sync.WaitGroup |
| 246 | for i := 0; i < 50; i++ { |
| 247 | wg.Add(2) |
| 248 | go func() { |
| 249 | defer wg.Done() |
| 250 | if err := setInstanceQuota(envPath, 12345, "free"); err != nil { |
| 251 | t.Error(err) |
| 252 | } |
| 253 | }() |
| 254 | go func() { |
| 255 | defer wg.Done() |
| 256 | if err := setInstanceMeter(envPath, "https://cloud.example/api/usage", "tok"); err != nil { |
| 257 | t.Error(err) |
| 258 | } |
| 259 | }() |
| 260 | } |
| 261 | wg.Wait() |
| 262 | env, err := readEnvFile(envPath) |
| 263 | if err != nil { |
| 264 | t.Fatal(err) |
| 265 | } |
| 266 | for _, key := range []string{"LCHAT_ADMIN_TOKEN", "LCHAT_BLOB_QUOTA_BYTES", "LCHAT_METER_URL", "LCHAT_METER_TOKEN"} { |
| 267 | if env[key] == "" { |
| 268 | t.Errorf("%s lost in concurrent pushes: %v", key, env) |
| 269 | } |
| 270 | } |
| 271 | } |
| 272 | |
amadan