drive-foundation-followups: レビュー nit 8 件の集約是正 (バリデーション・競合安全・キャッシュ・文書化)

README.md

@@ -1,2 +1,4 @@
 # Booskiff
 ShuttlePub 群のファイル保管サービス (Misskey Drive 的 UX)。core は Rust (Axum) + PostgreSQL + S3 互換ストレージ。
+
+運用仕様メモは [docs/operations.md](docs/operations.md) を参照。

compose.yml

@@ -17,7 +17,9 @@ services:
       retries: 30
 
   minio:
-    image: minio/minio:latest
+    # quay.io: MinIO no longer publishes to Docker Hub (pull access denied),
+    # which broke the CI e2e job.
+    image: quay.io/minio/minio:latest
     command: server /data --console-address ":9001"
     environment:
       MINIO_ROOT_USER: booskiff
@@ -36,7 +38,7 @@ services:
   # Create the dev bucket once MinIO is healthy; always exits 0 so the
   # compose project reports healthy even when the bucket already exists.
   mc-init:
-    image: minio/mc
+    image: quay.io/minio/mc:latest
     depends_on:
       minio:
         condition: service_healthy

core/src/admin/billing_cache_tests.rs

@@ -0,0 +1,162 @@
+use super::*;
+use crate::billing::cache::{BillingCache, OwnerEntry};
+use crate::config::Config;
+use axum::body::Body;
+use axum::http::Request;
+use http_body_util::BodyExt;
+use std::sync::Arc;
+use tower::ServiceExt;
+
+#[tokio::test]
+#[ignore = "requires compose postgres"]
+async fn pg_admin_billing_writes_invalidate_cache() {
+    let config = Config::default();
+    let pool = sqlx::PgPool::connect(&config.database_url).await.unwrap();
+    let mut guard = pool.begin().await.unwrap();
+    sqlx::query("SELECT pg_advisory_xact_lock($1)")
+        .bind(0x0062_6969_i64)
+        .execute(&mut *guard)
+        .await
+        .unwrap();
+    let token = uuid::Uuid::now_v7().to_string();
+    sqlx::query("INSERT INTO admin_tokens (name, token_hash) VALUES ($1, $2)")
+        .bind(&token)
+        .bind(hash_admin_token(&token))
+        .execute(&pool)
+        .await
+        .unwrap();
+    let owner = Owner::new("test-admin-cache", uuid::Uuid::now_v7().to_string());
+    let cache = Arc::new(BillingCache::new(config.billing_cache_ttl_secs));
+    let state = AppState {
+        s3: crate::storage::Storage::build(&config).await.unwrap(),
+        jwks_cache: crate::auth::jwks::JwksCache::new(Vec::new()),
+        pool: pool.clone(),
+        config,
+        rate_limiters: Arc::new(crate::state::RateLimiters::default()),
+        public_rate_limiter: Arc::new(crate::auth::rate_limit::PublicRateLimiter::new(300)),
+        billing_cache: Arc::clone(&cache),
+    };
+    let router = admin_router().with_state(state);
+    let plan_uri = format!("/owners/{}/{}/plan", owner.owner_type, owner.owner_id);
+    let owner_rule = serde_json::json!({"owner_type": owner.owner_type, "owner_id": owner.owner_id,
+        "key": "max_file_bytes", "value": 71, "enabled": true});
+    let mut rule_id = String::new();
+    for (method, uri, body, all) in [
+        (
+            "PUT",
+            plan_uri.clone(),
+            serde_json::json!({"plan":"premium"}),
+            false,
+        ),
+        ("DELETE", plan_uri, serde_json::Value::Null, false),
+        ("POST", "/billing/rules".into(), owner_rule, false),
+        ("DELETE", String::new(), serde_json::Value::Null, true),
+        (
+            "POST",
+            "/billing/rules".into(),
+            serde_json::json!({"key":"max_file_bytes", "value":72, "enabled":true}),
+            true,
+        ),
+        ("DELETE", String::new(), serde_json::Value::Null, true),
+    ] {
+        cache.store_global(vec![]);
+        let other = Owner::new("other", "owner");
+        for cached_owner in [&owner, &other] {
+            cache.store_owner(
+                cached_owner,
+                OwnerEntry {
+                    assignment: None,
+                    rules: vec![],
+                },
+            );
+        }
+        let uri = if uri.is_empty() {
+            format!("/billing/rules/{rule_id}")
+        } else {
+            uri
+        };
+        let request = Request::builder()
+            .method(method)
+            .uri(uri)
+            .header("x-admin-token", &token)
+            .header("content-type", "application/json")
+            .body(Body::from(body.to_string()))
+            .unwrap();
+        let response = router.clone().oneshot(request).await.unwrap();
+        assert!(
+            response.status().is_success(),
+            "{method}: {}",
+            response.status()
+        );
+        assert!(
+            cache.get_owner(&owner).is_none(),
+            "{method} must invalidate owner"
+        );
+        assert_eq!(cache.get_owner(&other).is_none(), all);
+        assert_eq!(cache.get_global().is_none(), all);
+        if method == "POST" {
+            let bytes = response.into_body().collect().await.unwrap().to_bytes();
+            let item: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
+            rule_id = item["id"].as_str().unwrap().to_owned();
+        }
+    }
+
+    sqlx::query(
+        "INSERT INTO billing_rules (owner_type, owner_id, key, value, enabled) \
+         VALUES (NULL, NULL, $1, $2, true) \
+         ON CONFLICT ((COALESCE(owner_type, '')), (COALESCE(owner_id, '')), key) \
+         DO UPDATE SET value = EXCLUDED.value, enabled = EXCLUDED.enabled",
+    )
+    .bind("max_file_bytes")
+    .bind(serde_json::json!(72))
+    .execute(&pool)
+    .await
+    .unwrap();
+    cache.store_global(vec![("max_file_bytes".to_owned(), serde_json::json!(72))]);
+    let request = Request::builder()
+        .method("POST")
+        .uri("/billing/rules")
+        .header("x-admin-token", &token)
+        .header("content-type", "application/json")
+        .body(Body::from(
+            serde_json::json!({
+                "owner_type": "",
+                "owner_id": owner.owner_id,
+                "key": "max_file_bytes",
+                "value": 999,
+                "enabled": true
+            })
+            .to_string(),
+        ))
+        .unwrap();
+    let response = router.clone().oneshot(request).await.unwrap();
+    assert_eq!(response.status(), StatusCode::BAD_REQUEST);
+    let bytes = response.into_body().collect().await.unwrap().to_bytes();
+    let body: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
+    assert_eq!(body["error"]["code"], "validation");
+    assert!(
+        body["error"]["message"]
+            .as_str()
+            .unwrap()
+            .contains("owner_type must not be blank")
+    );
+
+    let global_value: serde_json::Value = sqlx::query_scalar(
+        "SELECT value FROM billing_rules WHERE owner_type IS NULL AND owner_id IS NULL AND key = $1",
+    )
+    .bind("max_file_bytes")
+    .fetch_one(&pool)
+    .await
+    .unwrap();
+    assert_eq!(global_value, serde_json::json!(72));
+    assert_eq!(
+        cache.get_global(),
+        Some(vec![("max_file_bytes".to_owned(), serde_json::json!(72))])
+    );
+
+    sqlx::query("DELETE FROM admin_tokens WHERE name = $1")
+        .bind(&token)
+        .execute(&pool)
+        .await
+        .unwrap();
+}

core/src/admin/handlers.rs

@@ -131,7 +131,19 @@ fn owner_scope(
 ) -> Result<Option<Owner>, AppError> {
     match (owner_type, owner_id) {
         (None, None) => Ok(None),
-        (Some(owner_type), Some(owner_id)) => Ok(Some(Owner::new(owner_type, owner_id))),
+        (Some(owner_type), Some(owner_id)) => {
+            if owner_type.trim().is_empty() {
+                return Err(AppError::Validation(
+                    "owner_type must not be blank".to_owned(),
+                ));
+            }
+            if owner_id.trim().is_empty() {
+                return Err(AppError::Validation(
+                    "owner_id must not be blank".to_owned(),
+                ));
+            }
+            Ok(Some(Owner::new(owner_type, owner_id)))
+        }
         (Some(_), None) | (None, Some(_)) => Err(AppError::Validation(
             "owner_type and owner_id must be given together; omit both for a global rule"
                 .to_owned(),
@@ -293,6 +305,10 @@ async fn create_billing_rule(
         request.enabled,
     )
     .await?;
+    match scope {
+        Some(owner) => state.billing_cache.invalidate_owner(&owner),
+        None => state.billing_cache.invalidate_all(),
+    }
     Ok(Json(billing_rule_item(rule)?))
 }
 
@@ -314,6 +330,8 @@ async fn delete_billing_rule(
     Path(id): Path<Uuid>,
 ) -> Result<StatusCode, AppError> {
     if rules::delete_rule(&state.pool, id).await? {
+        // Deletion returns no scope; conservatively invalidate every layer.
+        state.billing_cache.invalidate_all();
         Ok(StatusCode::NO_CONTENT)
     } else {
         Err(AppError::NotFound(format!("billing rule {id} not found")))
@@ -343,7 +361,9 @@ async fn set_plan(
     Json(request): Json<SetPlanRequest>,
 ) -> Result<StatusCode, AppError> {
     let plan = parse_plan(&request.plan)?;
-    assignments::set_plan(&state.pool, &Owner::new(owner_type, owner_id), plan).await?;
+    let owner = Owner::new(owner_type, owner_id);
+    assignments::set_plan(&state.pool, &owner, plan).await?;
+    state.billing_cache.invalidate_owner(&owner);
     Ok(StatusCode::NO_CONTENT)
 }
 
@@ -395,6 +415,7 @@ async fn delete_plan(
 ) -> Result<StatusCode, AppError> {
     let owner = Owner::new(owner_type, owner_id);
     if assignments::delete_plan(&state.pool, &owner).await? {
+        state.billing_cache.invalidate_owner(&owner);
         Ok(StatusCode::NO_CONTENT)
     } else {
         Err(AppError::NotFound(format!(
@@ -450,12 +471,18 @@ pub fn admin_router() -> Router<AppState> {
         )
         .route("/billing/rules/{id}", delete(delete_billing_rule))
         .route(
+            // Axum path parameters require non-empty segments, so these plan
+            // routes cannot receive an empty owner component.
             "/owners/{owner_type}/{owner_id}/plan",
             put(set_plan).get(get_plan).delete(delete_plan),
         )
         .route("/owners/{owner_type}/{owner_id}/usage", get(get_usage))
 }
 
+#[cfg(test)]
+#[path = "billing_cache_tests.rs"]
+mod billing_cache_tests;
+
 #[cfg(test)]
 mod tests {
     use std::collections::HashSet;
@@ -603,8 +630,14 @@ mod tests {
             pool: pool.clone(),
             s3: Storage::build(&config).await.unwrap(),
             jwks_cache: JwksCache::new(Vec::new()),
+            billing_cache: std::sync::Arc::new(crate::billing::cache::BillingCache::new(
+                config.billing_cache_ttl_secs,
+            )),
             config,
             rate_limiters: std::sync::Arc::new(RateLimiters::default()),
+            public_rate_limiter: std::sync::Arc::new(
+                crate::auth::rate_limit::PublicRateLimiter::new(300),
+            ),
         };
         let router = axum::Router::new()
             .nest("/v1/admin", admin_router())

core/src/auth/admin_auth.rs

@@ -140,8 +140,12 @@ mod tests {
                 .expect("lazy pool never connects"),
             s3: Storage::build(&config).await.unwrap(),
             jwks_cache: JwksCache::new(Vec::new()),
+            billing_cache: Arc::new(crate::billing::cache::BillingCache::new(
+                config.billing_cache_ttl_secs,
+            )),
             config,
             rate_limiters: Arc::new(RateLimiters::default()),
+            public_rate_limiter: Arc::new(crate::auth::rate_limit::PublicRateLimiter::new(300)),
         }
     }
 
@@ -204,8 +208,12 @@ mod tests {
             s3: Storage::build(&config).await.unwrap(),
             pool,
             jwks_cache: JwksCache::new(Vec::new()),
+            billing_cache: Arc::new(crate::billing::cache::BillingCache::new(
+                config.billing_cache_ttl_secs,
+            )),
             config,
             rate_limiters: Arc::new(RateLimiters::default()),
+            public_rate_limiter: Arc::new(crate::auth::rate_limit::PublicRateLimiter::new(300)),
         };
         let router = admin_router(state.clone());
 

core/src/auth/extractor.rs

@@ -5,7 +5,7 @@ use axum::http::request::Parts;
 use axum::response::{IntoResponse, Response};
 
 use crate::auth::jwt::{owner_from_claims, verify_token};
-use crate::billing::resolve::effective_limits;
+use crate::billing::resolve::effective_limits_cached;
 use crate::error::AppError;
 use crate::model::{Limits, Owner};
 use crate::state::AppState;
@@ -30,9 +30,10 @@ impl FromRequestParts<AppState> for AccountContext {
             .map_err(|err| err.into_response())?;
         let owner = owner_from_claims(&claims, &state.config.jwt_owner_type_claim)
             .map_err(|err| err.into_response())?;
-        let limits = effective_limits(&state.pool, &state.config, &owner)
-            .await
-            .map_err(|err| err.into_response())?;
+        let limits =
+            effective_limits_cached(&state.billing_cache, &state.pool, &state.config, &owner)
+                .await
+                .map_err(|err| err.into_response())?;
         if !state
             .rate_limiters
             .check(&owner.key(), limits.rate_limit_rpm)
@@ -122,8 +123,8 @@ mod tests {
     }
 
     /// Compose postgres (compose.yml), required by the DB-backed tests:
-    /// the billing `effective_limits` resolution queries `billing_rules`
-    /// on every request, so the 200/429 paths need a live database.
+    /// billing cache misses query `billing_rules`, so the 200/429 paths
+    /// need a live database.
     const PUBLISH_DB_URL: &str = "postgres://booskiff:booskiff@127.0.0.1:5432/booskiff";
 
     /// Router state for the DB-free paths (header parsing and token
@@ -144,8 +145,12 @@ mod tests {
                 .expect("lazy pool never connects"),
             s3: Storage::build(&config).await.unwrap(),
             jwks_cache: JwksCache::new(config.jwt_trusted_issuers.clone()),
+            billing_cache: Arc::new(crate::billing::cache::BillingCache::new(
+                config.billing_cache_ttl_secs,
+            )),
             config,
             rate_limiters: Arc::new(RateLimiters::default()),
+            public_rate_limiter: Arc::new(crate::auth::rate_limit::PublicRateLimiter::new(300)),
         }
     }
 
@@ -168,8 +173,12 @@ mod tests {
             pool: sqlx::PgPool::connect(PUBLISH_DB_URL).await.unwrap(),
             s3: Storage::build(&config).await.unwrap(),
             jwks_cache: JwksCache::new(config.jwt_trusted_issuers.clone()),
+            billing_cache: Arc::new(crate::billing::cache::BillingCache::new(
+                config.billing_cache_ttl_secs,
+            )),
             config,
             rate_limiters: Arc::new(RateLimiters::default()),
+            public_rate_limiter: Arc::new(crate::auth::rate_limit::PublicRateLimiter::new(300)),
         }
     }