Update auth token validation to return account level and handle subsc…

…ription check

binaries/geph5-broker/src/auth.rs

@@ -1,9 +1,10 @@
 use std::ops::Deref as _;
 
 use argon2::{password_hash::Encoding, Argon2, PasswordHash, PasswordVerifier};
-use geph5_broker_protocol::AuthError;
+use geph5_broker_protocol::{AccountLevel, AuthError};
 
 use rand::Rng as _;
+use tracing::Level;
 
 use crate::{database::POSTGRES, log_error};
 
@@ -51,11 +52,20 @@ pub async fn new_auth_token(user_id: i32) -> anyhow::Result<String> {
     }
 }
 
-pub async fn valid_auth_token(token: &str) -> anyhow::Result<bool> {
-    let row: (i64,) = sqlx::query_as("SELECT COUNT(*) FROM auth_tokens WHERE token = $1")
+pub async fn valid_auth_token(token: &str) -> anyhow::Result<Option<AccountLevel>> {
+    let user_id: Option<(i32, i64)> = sqlx::query_as("SELECT user_id, (select count(*) from subscriptions where id = user_id) FROM auth_tokens WHERE token = $1")
         .bind(token)
-        .fetch_one(POSTGRES.deref())
+        .fetch_optional(POSTGRES.deref())
         .await?;
 
-    Ok(row.0 > 0)
+    if let Some((user_id, is_plus)) = user_id {
+        tracing::debug!(user_id, is_plus, "valid auth token");
+        if is_plus == 0 {
+            Ok(Some(AccountLevel::Free))
+        } else {
+            Ok(Some(AccountLevel::Plus))
+        }
+    } else {
+        Ok(None)
+    }
 }

binaries/geph5-broker/src/rpc_impl.rs

@@ -65,18 +65,23 @@ impl BrokerProtocol for BrokerImpl {
         epoch: u16,
         blind_token: BlindedClientToken,
     ) -> Result<BlindedSignature, AuthError> {
-        match valid_auth_token(&auth_token).await {
+        let user_level = match valid_auth_token(&auth_token).await {
             Ok(auth) => {
-                if !auth {
+                if let Some(level) = auth {
+                    level
+                } else {
                     return Err(AuthError::Forbidden);
                 }
             }
             Err(err) => {
                 tracing::warn!(err = debug(err), "database failed");
                 return Err(AuthError::RateLimited);
             }
-        }
+        };
         let start = Instant::now();
+        if user_level != level {
+            return Err(AuthError::WrongLevel);
+        }
         let signed = match level {
             AccountLevel::Free => &FREE_MIZARU_SK,
             AccountLevel::Plus => &PLUS_MIZARU_SK,

libraries/geph5-broker-protocol/src/lib.rs

@@ -47,7 +47,7 @@ pub trait BrokerProtocol {
     async fn incr_stat(&self, stat: String, value: i32);
 }
 
-#[derive(Clone, Copy, Debug, Serialize, Deserialize)]
+#[derive(Clone, Copy, Debug, Serialize, Deserialize, PartialEq, Eq)]
 pub enum AccountLevel {
     Free,
     Plus,