//! Invite-code enforcement on `com.atproto.server.createAccount`. //! //! # Why this file starts its own PDS //! //! Every other integration suite in this crate talks to whatever PDS the //! developer already has running on `:2583` and skips itself when there //! isn't one. That works because those tests only need *a* PDS. These //! need a PDS with `PDS_INVITE_REQUIRED=true`, and the ambient one is //! (correctly) started with the default `false` — otherwise every other //! suite, which creates throwaway accounts with no code, would fail. //! //! Asking the developer to restart their PDS with a different flag //! before this file passes would mean the flag's behaviour is only ever //! tested by hand. So each test here spawns its own `pds-server` on a //! free port with the flag set the way that test needs it, and kills it //! on the way out ([`Pds`]'s `Drop`). `env!("CARGO_BIN_EXE_pds-server")` //! is cargo's own path to the binary it just built for this test run, so //! the process under test is always the current code. //! //! The suite still fails open, in the same spirit as its neighbours: if //! the child never becomes healthy — no Postgres, no `.env`, no //! `DATABASE_URL_PDS` — the tests print why and return green rather than //! failing a workstation that simply isn't running the stack. use serde_json::{json, Value}; use std::process::{Child, Command}; use std::time::Duration; /// A `pds-server` child process bound to its own port, killed when the /// test that started it goes out of scope. /// /// The `Drop` impl is the reason this is a struct at all: a test that /// panics mid-way must not leave a server holding a port and a pool of /// Postgres connections for the rest of the run. struct Pds { child: Child, port: u16, http: reqwest::Client, } impl Drop for Pds { fn drop(&mut self) { let _ = self.child.kill(); let _ = self.child.wait(); } } impl Pds { fn url(&self, path: &str) -> String { format!("http://127.0.0.1:{}{}", self.port, path) } async fn create_account(&self, body: Value) -> (u16, Value) { let resp = self .http .post(self.url("/xrpc/com.atproto.server.createAccount")) .json(&body) .send() .await .expect("createAccount request"); let status = resp.status().as_u16(); let body: Value = resp.json().await.unwrap_or(Value::Null); (status, body) } async fn describe(&self) -> Value { self.http .get(self.url("/xrpc/com.atproto.server.describeServer")) .send() .await .expect("describeServer") .json() .await .expect("describeServer json") } } /// Ask the OS for a port nobody is using, then let go of it. /// /// There is a window between the drop and the child's `bind` in which /// something else could take the port; on a test machine that window is /// theoretical, and the alternative (a fixed port) would make two /// concurrently running tests in this file collide *reliably* instead of /// never. fn free_port() -> Option { let l = std::net::TcpListener::bind("127.0.0.1:0").ok()?; let p = l.local_addr().ok()?.port(); drop(l); Some(p) } /// Start a `pds-server` with `PDS_INVITE_REQUIRED` set to `required`. /// /// Returns `None` when the stack this needs isn't available, which the /// callers turn into a skip. The child inherits the ambient environment /// (so `DATABASE_URL_PDS` and friends come from `.env` exactly as they /// do for the real server — `dotenvy` does not override real variables, /// so our overrides below win). async fn start_pds(required: bool) -> Option { let port = free_port()?; let child = Command::new(env!("CARGO_BIN_EXE_pds-server")) .env("PDS_HOST", "127.0.0.1") .env("PDS_PORT", port.to_string()) .env("PDS_PUBLIC_URL", format!("http://127.0.0.1:{port}")) .env("PDS_INVITE_REQUIRED", if required { "true" } else { "false" }) // Point the PLC submit at a closed port. `create_account` // tolerates a failed submit by design (the DID is computed // locally), and a connection refused on loopback fails in // microseconds — whereas the real directory would add a network // round-trip to every account this file creates, and might // actually publish throwaway test DIDs. .env("PLC_DIRECTORY_URL", "http://127.0.0.1:1") .env("RUST_LOG", "warn") .stdout(std::process::Stdio::null()) .stderr(std::process::Stdio::null()) .spawn() .ok()?; let http = reqwest::Client::builder() .timeout(Duration::from_secs(30)) .build() .ok()?; let mut pds = Pds { child, port, http }; for _ in 0..80 { // If the child already exited (bad/missing env, no Postgres), // stop waiting — there is nothing to become healthy. if let Ok(Some(_)) = pds.child.try_wait() { return None; } if let Ok(r) = pds.http.get(pds.url("/healthz")).send().await { if r.status().is_success() { return Some(pds); } } tokio::time::sleep(Duration::from_millis(150)).await; } None } async fn db_pool() -> Option { let url = std::env::var("DATABASE_URL_PDS") .unwrap_or_else(|_| "postgres://pds:pds@127.0.0.1:5434/pds".to_string()); sqlx::postgres::PgPoolOptions::new() .max_connections(4) .acquire_timeout(Duration::from_secs(3)) .connect(&url) .await .ok() } /// Put a code straight into the table in whatever state the test needs. /// /// Tests seed through SQL rather than through `pds-server invite create` /// because the states that matter here — already spent, disabled — are /// not states the CLI can mint directly, and because a test that had to /// shell out to a second binary to arrange its fixture would be testing /// two things at once. async fn seed_code(db: &sqlx::PgPool, max_uses: i32, used: i32, disabled: bool) -> String { let code = format!("mt-test-{}", uuid::Uuid::new_v4().simple()); sqlx::query( "INSERT INTO invite_codes (code, max_uses, used_count, disabled) VALUES ($1, $2, $3, $4)", ) .bind(&code) .bind(max_uses) .bind(used) .bind(disabled) .execute(db) .await .expect("seed invite code"); code } async fn used_count(db: &sqlx::PgPool, code: &str) -> i32 { sqlx::query_scalar::<_, i32>("SELECT used_count FROM invite_codes WHERE code = $1") .bind(code) .fetch_one(db) .await .expect("read used_count") } async fn use_rows(db: &sqlx::PgPool, code: &str) -> Vec<(String, String)> { sqlx::query_as::<_, (String, String)>( "SELECT did, handle FROM invite_code_uses WHERE code = $1 ORDER BY used_at", ) .bind(code) .fetch_all(db) .await .expect("read invite_code_uses") } /// A unique throwaway handle. /// /// `createAccount` caps handles at 64 characters, and /// `_<32 hex>.maarcadetweet.local` overshoots that for anything /// but the shortest prefix — a limit that shows up as a confusing /// `InvalidHandle` in a test that is about invite codes. Half the UUID /// is 64 bits of uniqueness, which is plenty for a test fixture and /// leaves room for a readable prefix. fn handle(prefix: &str) -> String { let uniq = uuid::Uuid::new_v4().simple().to_string(); format!("{}_{}.maarcadetweet.local", prefix, &uniq[..16]) } /// Every invite rejection must look the same to a client: `400` with the /// module's usual `{error, message}` body and the name /// `InvalidInviteCode`. fn assert_invalid_invite(status: u16, body: &Value, what: &str) { assert_eq!(status, 400, "{what}: expected 400, body = {body}"); assert_eq!( body["error"], "InvalidInviteCode", "{what}: wrong error name, body = {body}" ); assert!( body["message"].is_string(), "{what}: error body must carry a message, body = {body}" ); // No DID may have been minted on a rejected request. assert!( body["did"].is_null(), "{what}: rejected request returned a did, body = {body}" ); } // -- enforcement ------------------------------------------------------------ /// The happy path, and the property that makes a single-use code /// single-use: after the account exists, the same code is dead. /// /// Also checks the audit trail, which is the reason /// `invite_code_uses` exists at all — "which account did this code /// create" has to be answerable after the fact. #[tokio::test] async fn valid_code_admits_one_account_then_is_spent() { let Some(db) = db_pool().await else { eprintln!("no pds database, skipping"); return; }; let Some(pds) = start_pds(true).await else { eprintln!("could not start a pds with PDS_INVITE_REQUIRED=true, skipping"); return; }; // The switch must be advertised, not just enforced — a client reads // this before it asks the user for anything. assert_eq!( pds.describe().await["invite_code_required"], json!(true), "describeServer must report the actual PDS_INVITE_REQUIRED value" ); let code = seed_code(&db, 1, 0, false).await; let h = handle("inv_ok"); let (status, body) = pds .create_account(json!({ "handle": h, "password": "hunter2hunter2", "invite_code": code, })) .await; assert_eq!(status, 200, "valid code must create an account: {body}"); let did = body["did"].as_str().expect("did").to_string(); assert_eq!(used_count(&db, &code).await, 1, "code must be counted as used"); let uses = use_rows(&db, &code).await; assert_eq!(uses.len(), 1); assert_eq!(uses[0].0, did, "audit row must name the account it created"); assert_eq!(uses[0].1, h, "audit row must snapshot the handle"); // Second attempt on the now-spent code. let (status2, body2) = pds .create_account(json!({ "handle": handle("inv_second"), "password": "hunter2hunter2", "invite_code": code, })) .await; assert_invalid_invite(status2, &body2, "spent code"); assert_eq!( used_count(&db, &code).await, 1, "a rejected attempt must not move the counter" ); } /// Every way a code can fail, and the missing-code case, all land on the /// same `400 InvalidInviteCode`. #[tokio::test] async fn unknown_disabled_spent_and_missing_codes_are_rejected() { let Some(db) = db_pool().await else { eprintln!("no pds database, skipping"); return; }; let Some(pds) = start_pds(true).await else { eprintln!("could not start a pds with PDS_INVITE_REQUIRED=true, skipping"); return; }; // Unknown. let (s, b) = pds .create_account(json!({ "handle": handle("inv_unknown"), "password": "hunter2hunter2", "invite_code": "mt-zzzzz-zzzzz", })) .await; assert_invalid_invite(s, &b, "unknown code"); // Disabled, with uses left — proves `disabled` is checked and not // just the counter. let disabled = seed_code(&db, 5, 0, true).await; let (s, b) = pds .create_account(json!({ "handle": handle("inv_disabled"), "password": "hunter2hunter2", "invite_code": disabled, })) .await; assert_invalid_invite(s, &b, "disabled code"); assert_eq!(used_count(&db, &disabled).await, 0); // Already at its limit. let spent = seed_code(&db, 2, 2, false).await; let (s, b) = pds .create_account(json!({ "handle": handle("inv_spent"), "password": "hunter2hunter2", "invite_code": spent, })) .await; assert_invalid_invite(s, &b, "exhausted code"); // No field at all. let (s, b) = pds .create_account(json!({ "handle": handle("inv_none"), "password": "hunter2hunter2", })) .await; assert_invalid_invite(s, &b, "missing code"); // Present but blank / whitespace — must be indistinguishable from // absent, not an attempt to look up the empty string. for blank in ["", " "] { let (s, b) = pds .create_account(json!({ "handle": handle("inv_blank"), "password": "hunter2hunter2", "invite_code": blank, })) .await; assert_invalid_invite(s, &b, "blank code"); } } /// The camelCase spelling from the atproto lexicon, and the multi-use /// case the schema exists for. #[tokio::test] async fn camel_case_spelling_works_and_multi_use_codes_stop_at_the_limit() { let Some(db) = db_pool().await else { eprintln!("no pds database, skipping"); return; }; let Some(pds) = start_pds(true).await else { eprintln!("could not start a pds with PDS_INVITE_REQUIRED=true, skipping"); return; }; // `inviteCode` is what an off-the-shelf atproto client sends. let code = seed_code(&db, 3, 0, false).await; let (s, b) = pds .create_account(json!({ "handle": handle("inv_camel"), "password": "hunter2hunter2", "inviteCode": code, })) .await; assert_eq!(s, 200, "inviteCode spelling must be accepted: {b}"); // Case and stray whitespace are normalised, so a code shouted or // pasted out of a chat window still works. let (s, b) = pds .create_account(json!({ "handle": handle("inv_case"), "password": "hunter2hunter2", "invite_code": format!(" {} ", code.to_uppercase()), })) .await; assert_eq!(s, 200, "normalised code must be accepted: {b}"); // Third and last use. let (s, _) = pds .create_account(json!({ "handle": handle("inv_third"), "password": "hunter2hunter2", "invite_code": code, })) .await; assert_eq!(s, 200); // Fourth is one too many. let (s, b) = pds .create_account(json!({ "handle": handle("inv_fourth"), "password": "hunter2hunter2", "invite_code": code, })) .await; assert_invalid_invite(s, &b, "one past max_uses"); assert_eq!(used_count(&db, &code).await, 3); assert_eq!(use_rows(&db, &code).await.len(), 3); } /// A failed account creation must not consume the code. /// /// The cheapest way to make the account creation fail *after* the /// redemption has already run is a handle that is already taken: the /// redeem happens first inside the transaction, the `users` insert then /// trips the unique index, and the whole transaction rolls back. If the /// redemption had been done outside the transaction (or committed /// separately) the user would have lost their code to someone else's /// handle. #[tokio::test] async fn a_failed_registration_does_not_burn_the_code() { let Some(db) = db_pool().await else { eprintln!("no pds database, skipping"); return; }; let Some(pds) = start_pds(true).await else { eprintln!("could not start a pds with PDS_INVITE_REQUIRED=true, skipping"); return; }; let taken = handle("inv_taken"); let first = seed_code(&db, 1, 0, false).await; let (s, b) = pds .create_account(json!({ "handle": taken, "password": "hunter2hunter2", "invite_code": first, })) .await; assert_eq!(s, 200, "{b}"); // Now a *different* code, used on a handle that cannot be created. let code = seed_code(&db, 1, 0, false).await; let (s, _b) = pds .create_account(json!({ "handle": taken, "password": "hunter2hunter2", "invite_code": code, })) .await; assert_eq!(s, 409, "duplicate handle is still a 409"); assert_eq!( used_count(&db, &code).await, 0, "the code must survive a registration that rolled back" ); assert!(use_rows(&db, &code).await.is_empty()); // And it still works afterwards. let (s, b) = pds .create_account(json!({ "handle": handle("inv_retry"), "password": "hunter2hunter2", "invite_code": code, })) .await; assert_eq!(s, 200, "unburned code must still be redeemable: {b}"); } // -- the race --------------------------------------------------------------- /// Two (here: eight) registrations arriving at the same instant on the /// last remaining use of a code. Exactly one may get in. /// /// This is the test the whole design is built around. A /// `SELECT`-then-`UPDATE` implementation passes every other test in this /// file and fails this one: all eight requests read `used_count = 0`, /// all eight decide they are allowed, and the server hands out eight /// accounts for a one-use code while the row afterwards claims a single /// redemption. The fix is that `invite::redeem` never reads before it /// writes — the `WHERE used_count < max_uses` is part of the `UPDATE`, /// so Postgres re-evaluates it against the committed row after the /// row lock is released and the losers match zero rows. /// /// Every request uses a distinct handle, so nothing but the invite code /// can be what serialises them. #[tokio::test] async fn concurrent_registrations_cannot_share_one_use() { let Some(db) = db_pool().await else { eprintln!("no pds database, skipping"); return; }; let Some(pds) = start_pds(true).await else { eprintln!("could not start a pds with PDS_INVITE_REQUIRED=true, skipping"); return; }; const N: usize = 8; let code = seed_code(&db, 1, 0, false).await; let mut tasks = Vec::with_capacity(N); for i in 0..N { let http = pds.http.clone(); let url = pds.url("/xrpc/com.atproto.server.createAccount"); let code = code.clone(); let h = handle(&format!("inv_race{i}")); tasks.push(tokio::spawn(async move { let resp = http .post(url) .json(&json!({ "handle": h, "password": "hunter2hunter2", "invite_code": code, })) .send() .await .expect("concurrent createAccount"); let status = resp.status().as_u16(); let body: Value = resp.json().await.unwrap_or(Value::Null); (status, body) })); } let mut ok = Vec::new(); let mut rejected = 0usize; for t in tasks { let (status, body) = t.await.unwrap(); match status { 200 => ok.push(body), 400 => { assert_eq!(body["error"], "InvalidInviteCode", "body = {body}"); rejected += 1; } other => panic!("unexpected status {other}: {body}"), } } assert_eq!( ok.len(), 1, "a one-use code let {} concurrent registrations through — the redeem is racy", ok.len() ); assert_eq!(rejected, N - 1); assert_eq!(used_count(&db, &code).await, 1); let uses = use_rows(&db, &code).await; assert_eq!(uses.len(), 1, "counter and audit rows disagree: {uses:?}"); assert_eq!(uses[0].0, ok[0]["did"].as_str().unwrap()); } /// The same race with room for more than one winner: a three-use code /// hit by eight simultaneous registrations must admit exactly three. /// /// Worth having next to the one-use case because an implementation can /// be "safe" by accident for a single use (e.g. by serialising every /// registration globally) and still lose count when several are /// genuinely allowed to proceed. #[tokio::test] async fn concurrent_registrations_respect_a_multi_use_limit() { let Some(db) = db_pool().await else { eprintln!("no pds database, skipping"); return; }; let Some(pds) = start_pds(true).await else { eprintln!("could not start a pds with PDS_INVITE_REQUIRED=true, skipping"); return; }; const N: usize = 8; const USES: i32 = 3; let code = seed_code(&db, USES, 0, false).await; let mut tasks = Vec::with_capacity(N); for i in 0..N { let http = pds.http.clone(); let url = pds.url("/xrpc/com.atproto.server.createAccount"); let code = code.clone(); let h = handle(&format!("inv_mrace{i}")); tasks.push(tokio::spawn(async move { let resp = http .post(url) .json(&json!({ "handle": h, "password": "hunter2hunter2", "invite_code": code, })) .send() .await .expect("concurrent createAccount"); let status = resp.status().as_u16(); let body: Value = resp.json().await.unwrap_or(Value::Null); (status, body) })); } let mut ok = 0usize; for t in tasks { let (status, body) = t.await.unwrap(); match status { 200 => ok += 1, 400 => assert_eq!(body["error"], "InvalidInviteCode", "body = {body}"), other => panic!("unexpected status {other}: {body}"), } } assert_eq!(ok, USES as usize, "a {USES}-use code admitted {ok} accounts"); assert_eq!(used_count(&db, &code).await, USES); assert_eq!(use_rows(&db, &code).await.len(), USES as usize); } // -- the switch off --------------------------------------------------------- /// With `PDS_INVITE_REQUIRED=false` — the default, and what every other /// test suite in this workspace relies on — nothing about `createAccount` /// changes. /// /// This is the regression test for the whole feature's blast radius: the /// switch is off by default precisely so that the existing suites keep /// creating accounts with no code, and if that ever stopped being true /// the failure would show up as dozens of unrelated tests breaking. It /// shows up here instead. #[tokio::test] async fn switch_off_leaves_create_account_untouched() { let Some(_db) = db_pool().await else { eprintln!("no pds database, skipping"); return; }; let Some(pds) = start_pds(false).await else { eprintln!("could not start a pds with PDS_INVITE_REQUIRED=false, skipping"); return; }; assert_eq!( pds.describe().await["invite_code_required"], json!(false), "describeServer must report the actual PDS_INVITE_REQUIRED value" ); // No code at all: the historical behaviour. let (s, b) = pds .create_account(json!({ "handle": handle("inv_off"), "password": "hunter2hunter2", })) .await; assert_eq!(s, 200, "no-code registration must still work: {b}"); assert!(b["did"].as_str().unwrap().starts_with("did:")); assert!(b["access_jwt"].is_string()); // A code that does not exist is simply ignored rather than becoming // a new way to fail — a client that was talking to an invite-only // PDS yesterday must not break when the operator opens the server up. let (s, b) = pds .create_account(json!({ "handle": handle("inv_off_bogus"), "password": "hunter2hunter2", "invite_code": "mt-does-notexist", })) .await; assert_eq!(s, 200, "an ignored code must not fail the request: {b}"); // The other validations are untouched. let (s, b) = pds .create_account(json!({ "handle": handle("inv_off_short"), "password": "short", })) .await; assert_eq!(s, 400); assert_eq!(b["error"], "InvalidPassword"); } // -- the CLI ---------------------------------------------------------------- /// `pds-server invite create` / `list` / `disable`, run as the operator /// would run them, against the real database. /// /// The point is not that the SQL works (the tests above cover that) but /// that the *binary* exposes it: that `invite` short-circuits before the /// server starts, that `create` prints bare codes one per line so they /// can be pasted, and that a code it minted is actually redeemable. #[tokio::test] async fn invite_cli_mints_listable_redeemable_codes() { let Some(db) = db_pool().await else { eprintln!("no pds database, skipping"); return; }; let out = Command::new(env!("CARGO_BIN_EXE_pds-server")) .args(["invite", "create", "--count", "3", "--uses", "2", "--note", "cli test"]) .env("RUST_LOG", "warn") .output() .expect("run invite create"); if !out.status.success() { eprintln!( "invite create failed (no env/db?), skipping: {}", String::from_utf8_lossy(&out.stderr) ); return; } let stdout = String::from_utf8_lossy(&out.stdout); let codes: Vec<&str> = stdout.lines().filter(|l| !l.trim().is_empty()).collect(); assert_eq!(codes.len(), 3, "one code per line, nothing else: {stdout:?}"); for c in &codes { // Bare, paste-ready: no labels, no quotes, no indentation. assert_eq!(*c, c.trim(), "code line has surrounding whitespace: {c:?}"); assert!(c.starts_with("mt-"), "unexpected code shape: {c}"); assert_eq!( sqlx::query_scalar::<_, i32>("SELECT max_uses FROM invite_codes WHERE code = $1") .bind(c) .fetch_one(&db) .await .expect("minted code must be in the table"), 2, "--uses must reach the row" ); } // All three distinct — a generator that returned a constant would // otherwise only show up as a primary-key error. let unique: std::collections::HashSet<&&str> = codes.iter().collect(); assert_eq!(unique.len(), 3); // `list` must show what `redeem` would accept. let listed = Command::new(env!("CARGO_BIN_EXE_pds-server")) .args(["invite", "list"]) .env("RUST_LOG", "warn") .output() .expect("run invite list"); assert!(listed.status.success()); let listed = String::from_utf8_lossy(&listed.stdout); for c in &codes { assert!(listed.contains(*c), "invite list omitted {c}"); } // `disable` takes a code out without deleting it. let disabled = Command::new(env!("CARGO_BIN_EXE_pds-server")) .args(["invite", "disable", codes[0]]) .env("RUST_LOG", "warn") .output() .expect("run invite disable"); assert!(disabled.status.success()); assert!( sqlx::query_scalar::<_, bool>("SELECT disabled FROM invite_codes WHERE code = $1") .bind(codes[0]) .fetch_one(&db) .await .unwrap() ); let listed = Command::new(env!("CARGO_BIN_EXE_pds-server")) .args(["invite", "list"]) .env("RUST_LOG", "warn") .output() .expect("run invite list"); let listed = String::from_utf8_lossy(&listed.stdout); assert!( !listed.contains(codes[0]), "a disabled code must not show in the default listing" ); // And a minted code really lets an account through. let Some(pds) = start_pds(true).await else { eprintln!("could not start a pds with PDS_INVITE_REQUIRED=true, skipping redeem check"); return; }; let (s, b) = pds .create_account(json!({ "handle": handle("inv_cli"), "password": "hunter2hunter2", "invite_code": codes[1], })) .await; assert_eq!(s, 200, "CLI-minted code must be redeemable: {b}"); } /// An unknown subcommand must not silently boot a server, and `help` /// must not need a database. #[tokio::test] async fn unknown_subcommand_fails_instead_of_starting_a_server() { let out = Command::new(env!("CARGO_BIN_EXE_pds-server")) .args(["invit"]) .env("RUST_LOG", "warn") .output() .expect("run bad subcommand"); assert!(!out.status.success(), "a typo'd subcommand must not exit 0"); let stderr = String::from_utf8_lossy(&out.stderr); assert!(stderr.contains("unknown command"), "stderr = {stderr}"); let out = Command::new(env!("CARGO_BIN_EXE_pds-server")) .args(["invite", "help"]) .env("RUST_LOG", "warn") .output() .expect("run invite help"); assert!(out.status.success()); assert!(String::from_utf8_lossy(&out.stdout).contains("pds-server invite")); }