Add the set_storage_cap internal API

src/auth.rs

@@ -16,6 +16,7 @@ pub enum Right {
     Write,
     Delete,
     Usage,
+    Payment,
 }
 
 impl Right {
@@ -26,6 +27,7 @@ impl Right {
             Right::Write => "write",
             Right::Delete => "delete",
             Right::Usage => "usage",
+            Right::Payment => "payment",
         }
     }
 }

src/internal.rs

@@ -1,13 +1,16 @@
 use std::sync::Arc;
 
-use bfsp::internal::internal_file_server_message::Message;
+use bfsp::internal::internal_file_server_message::{Message, SetStorageCap};
 use bfsp::internal::GetStorageCapResp;
-/// The internal API
-use bfsp::Message as ProtoMessage;
 use bfsp::{
     chacha20poly1305::XChaCha20Poly1305,
-    internal::{decrypt_internal_message, EncryptedInternalFileServerMessage, GetUsageResp},
+    internal::{
+        decrypt_internal_message, EncryptedInternalFileServerMessage, GetUsageResp,
+        SetStorageCapResp,
+    },
 };
+/// The internal API
+use bfsp::{Message as ProtoMessage, PrependLen};
 use tokio::io::{AsyncReadExt, AsyncWriteExt};
 use tokio::net::TcpStream;
 use tracing::{event, Level};
@@ -44,7 +47,14 @@ async fn handle_internal_message<M: MetaDB>(
             }
             .encode_to_vec()
         }
+        Message::SetStorageCap(query) => {
+            let caps = query.storage_caps;
+            meta_db.set_storage_caps(caps).await.unwrap();
+
+            SetStorageCapResp { err: None }.encode_to_vec()
+        }
     }
+    .prepend_len()
 }
 
 #[tracing::instrument(skip(stream, internal_private_key))]

src/main.rs

@@ -570,6 +570,13 @@ async fn handle_upload_chunk<M: MetaDB + 'static, C: ChunkDB + 'static>(
     let storage_caps = meta_db.storage_caps(&[user_id]).await.unwrap();
     let storage_cap = *storage_caps.get(&user_id).unwrap();
 
+    event!(
+        Level::INFO,
+        storage_usage = storage_usage,
+        storage_cap = storage_cap,
+        chunk_size = chunk.len(),
+    );
+
     if storage_usage + chunk.len() as u64 > storage_cap {
         todo!("Deny uploads that exceed storage cap");
     }

src/meta_db.rs

@@ -67,6 +67,7 @@ pub trait MetaDB: Sized + Send + Sync + std::fmt::Debug {
         &self,
         user_ids: &[i64],
     ) -> impl Future<Output = Result<HashMap<i64, u64>>> + Send;
+    fn set_storage_caps(&self, caps: HashMap<i64, u64>) -> impl Future<Output = Result<()>> + Send;
 }
 
 #[derive(Debug)]
@@ -432,7 +433,7 @@ impl MetaDB for PostgresMetaDB {
         let mut caps: HashMap<i64, u64> = rows
             .into_iter()
             .map(|row| {
-                let storage_cap: i64 = row.get("sum");
+                let storage_cap: i64 = row.get("max_bytes");
                 let user_id: i64 = row.get("user_id");
 
                 (user_id.try_into().unwrap(), storage_cap.try_into().unwrap())
@@ -450,4 +451,23 @@ impl MetaDB for PostgresMetaDB {
 
         Ok(caps)
     }
+
+    #[tracing::instrument(err)]
+    async fn set_storage_caps(&self, caps: HashMap<i64, u64>) -> Result<()> {
+        let mut query = QueryBuilder::new("insert into storage_caps (user_id, max_bytes) values ");
+
+        let mut separated = query.separated(",");
+        for (user_id, cap) in caps {
+            println!("{} {}", user_id, cap);
+            separated.push(format!("({}, {})", user_id, cap));
+        }
+
+        query.push(" on conflict (user_id) do update set max_bytes = excluded.max_bytes");
+
+        let query = query.build();
+
+        query.execute(&self.pool).await?;
+
+        Ok(())
+    }
 }