Skip to main content

max / makenotwork

21.0 KB · 587 lines History Blame Raw
1 //! License key management: CRUD, activation tracking, and revocation.
2
3 use sqlx::PgPool;
4
5 use super::models::{DbLicenseActivation, DbLicenseKey};
6 use super::validated_types::KeyCode;
7 use super::{ItemId, LicenseActivationId, LicenseKeyId, TransactionId, UserId};
8 use crate::error::Result;
9
10 /// Create a new license key for an item.
11 ///
12 /// Retries once on a 23505 unique-violation with a freshly-generated code.
13 /// A real collision out of the wordlist generator is vanishingly rare (the
14 /// six-word space gives ~6B coin-flip headroom), but the alternative is
15 /// surfacing a 500 to whatever flow is creating the key.
16 #[tracing::instrument(skip_all)]
17 pub async fn create_license_key(
18 pool: &PgPool,
19 item_id: ItemId,
20 owner_id: UserId,
21 transaction_id: Option<TransactionId>,
22 key_code: &KeyCode,
23 max_activations: Option<i32>,
24 ) -> Result<DbLicenseKey> {
25 let first = sqlx::query_as!(
26 DbLicenseKey,
27 r#"
28 INSERT INTO license_keys (item_id, owner_id, transaction_id, key_code, max_activations)
29 VALUES ($1, $2, $3, $4, $5)
30 RETURNING id AS "id: LicenseKeyId", item_id AS "item_id: ItemId",
31 owner_id AS "owner_id: UserId", transaction_id AS "transaction_id: TransactionId",
32 key_code AS "key_code: KeyCode", max_activations, activation_count,
33 revoked_at AS "revoked_at: chrono::DateTime<chrono::Utc>",
34 created_at AS "created_at: chrono::DateTime<chrono::Utc>"
35 "#,
36 item_id as ItemId,
37 owner_id as UserId,
38 transaction_id as Option<TransactionId>,
39 key_code as &KeyCode,
40 max_activations,
41 )
42 .fetch_one(pool)
43 .await;
44
45 match first {
46 Ok(key) => Ok(key),
47 Err(sqlx::Error::Database(e)) if e.code().as_deref() == Some("23505") => {
48 // Distinguish which unique index fired. A `transaction_id` collision
49 // (mig 151's partial unique index) means this purchase already minted
50 // its one key, a duplicate finalize from crash-recovery redelivery.
51 // Return the existing key as idempotent success; retrying with a fresh
52 // code would only collide on the same index again (Pay-M1). A `key_code`
53 // collision is a random clash, regenerate and retry once.
54 let constraint = e.constraint().map(str::to_string);
55 if let (Some("license_keys_transaction_id_key"), Some(tx_id)) =
56 (constraint.as_deref(), transaction_id)
57 {
58 return get_license_key_by_transaction_id(pool, tx_id)
59 .await?
60 .ok_or_else(|| {
61 crate::error::AppError::Internal(anyhow::anyhow!(
62 "license_keys transaction_id unique violation but no existing \
63 row found for {tx_id:?}"
64 ))
65 });
66 }
67 let retry_code = crate::helpers::generate_key_code();
68 tracing::warn!(item_id = %item_id, "license key key_code 23505 collision; retrying once");
69 let key = sqlx::query_as!(
70 DbLicenseKey,
71 r#"
72 INSERT INTO license_keys (item_id, owner_id, transaction_id, key_code, max_activations)
73 VALUES ($1, $2, $3, $4, $5)
74 RETURNING id AS "id: LicenseKeyId", item_id AS "item_id: ItemId",
75 owner_id AS "owner_id: UserId", transaction_id AS "transaction_id: TransactionId",
76 key_code AS "key_code: KeyCode", max_activations, activation_count,
77 revoked_at AS "revoked_at: chrono::DateTime<chrono::Utc>",
78 created_at AS "created_at: chrono::DateTime<chrono::Utc>"
79 "#,
80 item_id as ItemId,
81 owner_id as UserId,
82 transaction_id as Option<TransactionId>,
83 &retry_code as &KeyCode,
84 max_activations,
85 )
86 .fetch_one(pool)
87 .await?;
88 Ok(key)
89 }
90 Err(e) => Err(e.into()),
91 }
92 }
93
94 /// Look up a license key by its code.
95 #[tracing::instrument(skip_all)]
96 pub async fn get_license_key_by_code(
97 pool: &PgPool,
98 key_code: &KeyCode,
99 ) -> Result<Option<DbLicenseKey>> {
100 let key = sqlx::query_as!(
101 DbLicenseKey,
102 r#"
103 SELECT id AS "id: LicenseKeyId", item_id AS "item_id: ItemId",
104 owner_id AS "owner_id: UserId", transaction_id AS "transaction_id: TransactionId",
105 key_code AS "key_code: KeyCode", max_activations, activation_count,
106 revoked_at AS "revoked_at: chrono::DateTime<chrono::Utc>",
107 created_at AS "created_at: chrono::DateTime<chrono::Utc>"
108 FROM license_keys WHERE key_code = $1
109 "#,
110 key_code as &KeyCode,
111 )
112 .fetch_optional(pool)
113 .await?;
114
115 Ok(key)
116 }
117
118 /// Look up the auto-minted license key for a purchase transaction, if any.
119 ///
120 /// Used by the finalize pre-check so a crash-recovery redelivery does not mint
121 /// a second key (the `license_keys_transaction_id_key` partial unique index is
122 /// the structural backstop). At most one such key exists per transaction.
123 #[tracing::instrument(skip_all)]
124 pub async fn get_license_key_by_transaction_id(
125 pool: &PgPool,
126 transaction_id: TransactionId,
127 ) -> Result<Option<DbLicenseKey>> {
128 let key = sqlx::query_as!(
129 DbLicenseKey,
130 r#"
131 SELECT id AS "id: LicenseKeyId", item_id AS "item_id: ItemId",
132 owner_id AS "owner_id: UserId", transaction_id AS "transaction_id: TransactionId",
133 key_code AS "key_code: KeyCode", max_activations, activation_count,
134 revoked_at AS "revoked_at: chrono::DateTime<chrono::Utc>",
135 created_at AS "created_at: chrono::DateTime<chrono::Utc>"
136 FROM license_keys WHERE transaction_id = $1
137 "#,
138 transaction_id as TransactionId,
139 )
140 .fetch_optional(pool)
141 .await?;
142
143 Ok(key)
144 }
145
146 /// Get a license key by ID. UNSCOPED: returns any user's key, so the caller MUST
147 /// authorize against the returned `owner_id` / `item_id` before acting on it
148 /// (the `_unchecked` suffix makes that contract legible at the call site, Sec-M2).
149 #[tracing::instrument(skip_all)]
150 pub async fn get_license_key_by_id_unchecked(
151 pool: &PgPool,
152 id: LicenseKeyId,
153 ) -> Result<Option<DbLicenseKey>> {
154 let key = sqlx::query_as!(
155 DbLicenseKey,
156 r#"
157 SELECT id AS "id: LicenseKeyId", item_id AS "item_id: ItemId",
158 owner_id AS "owner_id: UserId", transaction_id AS "transaction_id: TransactionId",
159 key_code AS "key_code: KeyCode", max_activations, activation_count,
160 revoked_at AS "revoked_at: chrono::DateTime<chrono::Utc>",
161 created_at AS "created_at: chrono::DateTime<chrono::Utc>"
162 FROM license_keys WHERE id = $1
163 "#,
164 id as LicenseKeyId,
165 )
166 .fetch_optional(pool)
167 .await?;
168
169 Ok(key)
170 }
171
172 /// Count license keys for an item.
173 #[tracing::instrument(skip_all)]
174 pub async fn count_keys_by_item(pool: &PgPool, item_id: ItemId) -> Result<i64> {
175 let count = sqlx::query_scalar!(
176 r#"SELECT COUNT(*) AS "count!" FROM license_keys WHERE item_id = $1"#,
177 item_id as ItemId,
178 )
179 .fetch_one(pool)
180 .await?;
181
182 Ok(count)
183 }
184
185 /// List all license keys for an item, newest first.
186 ///
187 /// Hard-caps at 500 rows to bound memory and response size for the creator
188 /// dashboard list view. Items with more than 500 keys are uncommon;
189 /// future work could add cursor-based pagination if needed.
190 #[tracing::instrument(skip_all)]
191 pub async fn get_license_keys_by_item(pool: &PgPool, item_id: ItemId) -> Result<Vec<DbLicenseKey>> {
192 let keys = sqlx::query_as!(
193 DbLicenseKey,
194 r#"
195 SELECT id AS "id: LicenseKeyId", item_id AS "item_id: ItemId",
196 owner_id AS "owner_id: UserId", transaction_id AS "transaction_id: TransactionId",
197 key_code AS "key_code: KeyCode", max_activations, activation_count,
198 revoked_at AS "revoked_at: chrono::DateTime<chrono::Utc>",
199 created_at AS "created_at: chrono::DateTime<chrono::Utc>"
200 FROM license_keys WHERE item_id = $1 ORDER BY created_at DESC LIMIT 500
201 "#,
202 item_id as ItemId,
203 )
204 .fetch_all(pool)
205 .await?;
206
207 Ok(keys)
208 }
209
210 /// Batch-load license keys for multiple items, grouped by item_id.
211 #[tracing::instrument(skip_all)]
212 pub async fn get_license_keys_by_items(
213 pool: &PgPool,
214 item_ids: &[ItemId],
215 ) -> Result<std::collections::HashMap<ItemId, Vec<DbLicenseKey>>> {
216 let keys = sqlx::query_as!(
217 DbLicenseKey,
218 r#"
219 SELECT id AS "id: LicenseKeyId", item_id AS "item_id: ItemId",
220 owner_id AS "owner_id: UserId", transaction_id AS "transaction_id: TransactionId",
221 key_code AS "key_code: KeyCode", max_activations, activation_count,
222 revoked_at AS "revoked_at: chrono::DateTime<chrono::Utc>",
223 created_at AS "created_at: chrono::DateTime<chrono::Utc>"
224 FROM license_keys WHERE item_id = ANY($1) ORDER BY item_id, created_at DESC
225 "#,
226 item_ids as &[ItemId],
227 )
228 .fetch_all(pool)
229 .await?;
230
231 let mut map: std::collections::HashMap<ItemId, Vec<DbLicenseKey>> =
232 std::collections::HashMap::new();
233 for k in keys {
234 map.entry(k.item_id).or_default().push(k);
235 }
236 Ok(map)
237 }
238
239 /// Find an existing activation for a key + machine combo.
240 #[tracing::instrument(skip_all)]
241 pub async fn get_activation(
242 pool: &PgPool,
243 license_key_id: LicenseKeyId,
244 machine_id: &str,
245 ) -> Result<Option<DbLicenseActivation>> {
246 let activation = sqlx::query_as!(
247 DbLicenseActivation,
248 r#"
249 SELECT id AS "id: LicenseActivationId", license_key_id AS "license_key_id: LicenseKeyId",
250 machine_id, label,
251 activated_at AS "activated_at: chrono::DateTime<chrono::Utc>",
252 last_validated_at AS "last_validated_at: chrono::DateTime<chrono::Utc>", is_active
253 FROM license_activations WHERE license_key_id = $1 AND machine_id = $2
254 "#,
255 license_key_id as LicenseKeyId,
256 machine_id,
257 )
258 .fetch_optional(pool)
259 .await?;
260
261 Ok(activation)
262 }
263
264 /// Update the last_validated_at timestamp for an existing activation.
265 #[tracing::instrument(skip_all)]
266 pub async fn touch_activation(pool: &PgPool, activation_id: LicenseActivationId) -> Result<()> {
267 sqlx::query!(
268 "UPDATE license_activations SET last_validated_at = NOW() WHERE id = $1",
269 activation_id as LicenseActivationId,
270 )
271 .execute(pool)
272 .await?;
273
274 Ok(())
275 }
276
277 /// Read the denormalized active-activation count for a key.
278 ///
279 /// `try_create_activation` keeps `license_keys.activation_count` authoritative
280 /// (it recomputes it under the row lock), so this is the value to report back to
281 /// the client rather than a pre-lock read or a manual `+ 1` guess that drifts
282 /// under concurrent activations (ultra-fuzz Run #1 Payments MINOR).
283 #[tracing::instrument(skip_all)]
284 pub async fn get_activation_count(pool: &PgPool, license_key_id: LicenseKeyId) -> Result<i32> {
285 let count = sqlx::query_scalar!(
286 "SELECT activation_count FROM license_keys WHERE id = $1",
287 license_key_id as LicenseKeyId,
288 )
289 .fetch_one(pool)
290 .await?;
291 Ok(count)
292 }
293
294 /// Activate a license key on a machine, atomically enforcing max_activations.
295 ///
296 /// Uses a transaction with `FOR UPDATE` to serialize concurrent activations
297 /// for the same key. Re-activations (same machine_id) always succeed via
298 /// upsert. New activations are rejected if the active count would exceed
299 /// `max_activations`.
300 ///
301 /// Returns `None` if the activation limit has been reached.
302 ///
303 /// After the upsert, the denormalized `activation_count` on `license_keys`
304 /// is refreshed with a full COUNT rather than an increment; this avoids
305 /// drift if a crash leaves the count out of sync.
306 #[tracing::instrument(skip_all)]
307 pub async fn try_create_activation(
308 pool: &PgPool,
309 license_key_id: LicenseKeyId,
310 machine_id: &str,
311 label: Option<&str>,
312 ) -> Result<Option<DbLicenseActivation>> {
313 let mut tx = pool.begin().await?;
314
315 // Lock the license key row to serialize concurrent activations, re-check
316 // revocation, AND read `max_activations` from the locked row, not from a
317 // caller-supplied argument (Pay-M2). The caller's value was read before the
318 // lock; if an admin lowered the cap in between, enforcing the stale arg would
319 // let an extra machine activate. Reading the column here makes the limit the
320 // authoritative one. No eligible (non-revoked) row => no activation.
321 let locked: Option<Option<i32>> = sqlx::query_scalar!(
322 r#"SELECT max_activations FROM license_keys WHERE id = $1 AND revoked_at IS NULL FOR UPDATE"#,
323 license_key_id as LicenseKeyId,
324 )
325 .fetch_optional(&mut *tx)
326 .await?;
327 let Some(max_activations) = locked else {
328 tx.rollback().await?;
329 return Ok(None);
330 };
331
332 // Check if this machine already has an activation (re-activation is always OK)
333 let existing: Option<DbLicenseActivation> = sqlx::query_as!(
334 DbLicenseActivation,
335 r#"
336 SELECT id AS "id: LicenseActivationId", license_key_id AS "license_key_id: LicenseKeyId",
337 machine_id, label,
338 activated_at AS "activated_at: chrono::DateTime<chrono::Utc>",
339 last_validated_at AS "last_validated_at: chrono::DateTime<chrono::Utc>", is_active
340 FROM license_activations WHERE license_key_id = $1 AND machine_id = $2
341 "#,
342 license_key_id as LicenseKeyId,
343 machine_id,
344 )
345 .fetch_optional(&mut *tx)
346 .await?;
347
348 // For truly new activations, enforce the limit
349 if existing.is_none()
350 && let Some(max) = max_activations
351 {
352 let count = sqlx::query_scalar!(
353 r#"SELECT COUNT(*) AS "count!" FROM license_activations WHERE license_key_id = $1 AND is_active = true"#,
354 license_key_id as LicenseKeyId,
355 )
356 .fetch_one(&mut *tx)
357 .await?;
358
359 if count >= max as i64 {
360 tx.rollback().await?;
361 return Ok(None);
362 }
363 }
364
365 // Upsert: if same machine_id re-activates, reactivate it
366 let activation = sqlx::query_as!(
367 DbLicenseActivation,
368 r#"
369 INSERT INTO license_activations (license_key_id, machine_id, label)
370 VALUES ($1, $2, $3)
371 ON CONFLICT (license_key_id, machine_id)
372 DO UPDATE SET is_active = true, last_validated_at = NOW(),
373 label = COALESCE(EXCLUDED.label, license_activations.label)
374 RETURNING id AS "id: LicenseActivationId", license_key_id AS "license_key_id: LicenseKeyId",
375 machine_id, label,
376 activated_at AS "activated_at: chrono::DateTime<chrono::Utc>",
377 last_validated_at AS "last_validated_at: chrono::DateTime<chrono::Utc>", is_active
378 "#,
379 license_key_id as LicenseKeyId,
380 machine_id,
381 label,
382 )
383 .fetch_one(&mut *tx)
384 .await?;
385
386 // Recount active activations to keep denormalized count accurate
387 sqlx::query!(
388 r#"
389 UPDATE license_keys
390 SET activation_count = (
391 SELECT COUNT(*) FROM license_activations
392 WHERE license_key_id = $1 AND is_active = true
393 )
394 WHERE id = $1
395 "#,
396 license_key_id as LicenseKeyId,
397 )
398 .execute(&mut *tx)
399 .await?;
400
401 tx.commit().await?;
402 Ok(Some(activation))
403 }
404
405 /// Deactivate a machine and update the key's activation_count.
406 ///
407 /// Only recounts if a row was actually deactivated (`rows_affected > 0`),
408 /// avoiding a wasted query when the machine wasn't active. Uses the same
409 /// full-recount strategy as [`try_create_activation`] for consistency.
410 #[tracing::instrument(skip_all)]
411 pub async fn deactivate_machine(
412 pool: &PgPool,
413 license_key_id: LicenseKeyId,
414 machine_id: &str,
415 ) -> Result<bool> {
416 let mut tx = pool.begin().await?;
417
418 let result = sqlx::query!(
419 r#"
420 UPDATE license_activations
421 SET is_active = false
422 WHERE license_key_id = $1 AND machine_id = $2 AND is_active = true
423 "#,
424 license_key_id as LicenseKeyId,
425 machine_id,
426 )
427 .execute(&mut *tx)
428 .await?;
429
430 if result.rows_affected() > 0 {
431 // Recount active activations
432 sqlx::query!(
433 r#"
434 UPDATE license_keys
435 SET activation_count = (
436 SELECT COUNT(*) FROM license_activations
437 WHERE license_key_id = $1 AND is_active = true
438 )
439 WHERE id = $1
440 "#,
441 license_key_id as LicenseKeyId,
442 )
443 .execute(&mut *tx)
444 .await?;
445
446 tx.commit().await?;
447 Ok(true)
448 } else {
449 tx.commit().await?;
450 Ok(false)
451 }
452 }
453
454 /// Create a manually-generated key for an item, atomically enforcing a per-item
455 /// cap. Locks the item row `FOR UPDATE` so concurrent manual issuance for the
456 /// same item serializes, the prior count-then-insert let N concurrent generates
457 /// each read `count = cap-1` and all insert, exceeding the cap (fuzz 2026-07-06
458 /// C6-1). Returns `None` if the cap is already reached (caller maps to a 400).
459 #[tracing::instrument(skip_all)]
460 pub async fn create_manual_key_capped(
461 pool: &PgPool,
462 item_id: ItemId,
463 owner_id: UserId,
464 key_code: &KeyCode,
465 max_activations: Option<i32>,
466 cap: i64,
467 ) -> Result<Option<DbLicenseKey>> {
468 let mut tx = pool.begin().await?;
469 // Serialize concurrent manual issuance for this item so the count below is
470 // stable through the insert (mirrors lock_project_for_splits' cap pattern).
471 sqlx::query!(
472 r#"SELECT id FROM items WHERE id = $1 FOR UPDATE"#,
473 item_id as ItemId
474 )
475 .fetch_one(&mut *tx)
476 .await?;
477 let count = sqlx::query_scalar!(
478 r#"SELECT COUNT(*) AS "count!" FROM license_keys WHERE item_id = $1"#,
479 item_id as ItemId,
480 )
481 .fetch_one(&mut *tx)
482 .await?;
483 if count >= cap {
484 tx.rollback().await?;
485 return Ok(None);
486 }
487 let key = sqlx::query_as!(
488 DbLicenseKey,
489 r#"
490 INSERT INTO license_keys (item_id, owner_id, transaction_id, key_code, max_activations)
491 VALUES ($1, $2, NULL, $3, $4)
492 RETURNING id AS "id: LicenseKeyId", item_id AS "item_id: ItemId",
493 owner_id AS "owner_id: UserId", transaction_id AS "transaction_id: TransactionId",
494 key_code AS "key_code: KeyCode", max_activations, activation_count,
495 revoked_at AS "revoked_at: chrono::DateTime<chrono::Utc>",
496 created_at AS "created_at: chrono::DateTime<chrono::Utc>"
497 "#,
498 item_id as ItemId,
499 owner_id as UserId,
500 key_code as &KeyCode,
501 max_activations,
502 )
503 .fetch_one(&mut *tx)
504 .await?;
505 tx.commit().await?;
506 Ok(Some(key))
507 }
508
509 /// Revoke a license key and deactivate all its activations.
510 ///
511 /// Wrapped in a transaction so the key revocation and activation
512 /// deactivation are atomic; a crash between the two statements
513 /// cannot leave the key revoked with activations still active.
514 #[tracing::instrument(skip_all)]
515 pub async fn revoke_license_key(pool: &PgPool, key_id: LicenseKeyId) -> Result<()> {
516 let mut tx = pool.begin().await?;
517
518 sqlx::query!(
519 r#"
520 UPDATE license_keys
521 SET revoked_at = NOW()
522 WHERE id = $1
523 "#,
524 key_id as LicenseKeyId,
525 )
526 .execute(&mut *tx)
527 .await?;
528
529 sqlx::query!(
530 "UPDATE license_activations SET is_active = false WHERE license_key_id = $1",
531 key_id as LicenseKeyId,
532 )
533 .execute(&mut *tx)
534 .await?;
535
536 tx.commit().await?;
537 Ok(())
538 }
539
540 /// Revoke all license keys for a given transaction and deactivate all activations.
541 /// Called from the Stripe `charge.refunded` webhook handler.
542 ///
543 /// Two-step approach: bulk-revoke keys, then bulk-deactivate activations.
544 /// Separate queries because `license_activations` is keyed by `license_key_id`,
545 /// not `transaction_id`.
546 #[tracing::instrument(skip_all)]
547 pub async fn revoke_keys_by_transaction(
548 conn: &mut sqlx::PgConnection,
549 transaction_id: TransactionId,
550 ) -> Result<u64> {
551 // Get all key IDs for this transaction
552 let key_ids: Vec<LicenseKeyId> = sqlx::query_scalar!(
553 r#"SELECT id AS "id: LicenseKeyId" FROM license_keys WHERE transaction_id = $1 AND revoked_at IS NULL"#,
554 transaction_id as TransactionId,
555 )
556 .fetch_all(&mut *conn)
557 .await?;
558
559 if key_ids.is_empty() {
560 return Ok(0);
561 }
562
563 // Revoke the keys
564 let result = sqlx::query!(
565 r#"
566 UPDATE license_keys
567 SET revoked_at = NOW()
568 WHERE transaction_id = $1 AND revoked_at IS NULL
569 "#,
570 transaction_id as TransactionId,
571 )
572 .execute(&mut *conn)
573 .await?;
574
575 // Deactivate all activations for those keys in a single query
576 if !key_ids.is_empty() {
577 sqlx::query!(
578 "UPDATE license_activations SET is_active = false WHERE license_key_id = ANY($1)",
579 &key_ids as &[LicenseKeyId],
580 )
581 .execute(&mut *conn)
582 .await?;
583 }
584
585 Ok(result.rows_affected())
586 }
587