A simple mock of L2 logic without proofs or signatures

Deposit is a big unanswered question

kairos-server/Cargo.toml

@@ -12,6 +12,7 @@ name = "kairos-server"
 
 [dependencies]
 axum = { version = "0.7", features = ["tracing"]}
+thiserror = "1.0"
 anyhow = "1.0"
 serde = { version = "1.0", features = ["derive"] }
 serde_json = "1.0"

kairos-server/src/errors.rs

@@ -8,8 +8,6 @@ use axum::{
     response::{IntoResponse, Response},
 };
 
-
-
 #[derive(Debug)]
 pub struct AppErr {
     error: anyhow::Error,

kairos-server/src/lib.rs

@@ -1,36 +1,16 @@
 pub mod errors;
 pub mod routes;
-
-use std::{
-    collections::{HashMap, HashSet},
-    sync::Arc,
-};
+pub mod state;
 
 use axum::{routing::post, Router};
-use routes::{transfer::Transfer, *};
-use tokio::sync::RwLock;
+use routes::*;
+use state::LockedBatchState;
 
 pub use errors::AppErr;
 
 type PublicKey = String;
 
-pub struct BatchState {
-    pub balances: HashMap<PublicKey, u64>,
-    pub batch_epoch: u64,
-    /// The set of transfers that will be batched in the next epoch.
-    pub batched_transfers: HashSet<Transfer>,
-}
-impl BatchState {
-    pub fn new() -> Arc<RwLock<Self>> {
-        Arc::new(RwLock::new(Self {
-            balances: HashMap::new(),
-            batch_epoch: 0,
-            batched_transfers: HashSet::new(),
-        }))
-    }
-}
-
-pub fn app_router(state: Arc<RwLock<BatchState>>) -> Router {
+pub fn app_router(state: LockedBatchState) -> Router {
     Router::new()
         .route("/api/v1/mock/deposit", post(deposit))
         .route("/api/v1/mock/withdraw", post(withdraw))

kairos-server/src/main.rs

@@ -1,5 +1,7 @@
 use std::net::SocketAddr;
 
+use kairos_server::state::BatchState;
+
 #[tokio::main]
 async fn main() {
     tracing_subscriber::fmt::init();
@@ -11,7 +13,7 @@ async fn main() {
         })
     });
 
-    let app = kairos_server::app_router(kairos_server::BatchState::new());
+    let app = kairos_server::app_router(BatchState::new());
 
     let axum_addr = SocketAddr::from(([127, 0, 0, 1], axum_port));
     tracing::info!("starting http server");

kairos-server/src/routes/deposit.rs

@@ -1,20 +1,41 @@
-use std::sync::Arc;
+use std::ops::Deref;
 
-use axum::{extract::State, Json};
+use anyhow::anyhow;
+use axum::{extract::State, http::StatusCode, Json};
 use serde::{Deserialize, Serialize};
-use tokio::sync::RwLock;
 
-use crate::{AppErr, BatchState, PublicKey};
+use crate::{state::LockedBatchState, AppErr, PublicKey};
 
 #[derive(Serialize, Deserialize)]
-pub struct DepositRequest {
+pub struct Deposit {
     pub public_key: PublicKey,
     pub amount: u64,
 }
 
 pub async fn deposit(
-    State(pool): State<Arc<RwLock<BatchState>>>,
-    Json(proof_request): Json<DepositRequest>,
+    state: State<LockedBatchState>,
+    Json(Deposit { public_key, amount }): Json<Deposit>,
 ) -> Result<(), AppErr> {
-    todo!("deposit")
+    tracing::info!("TODO: verifying deposit");
+
+    tracing::info!("TODO: adding deposit to batch");
+
+    let mut state = state.deref().write().await;
+    let account = state.balances.entry(public_key.clone());
+
+    let prior_balance = account.or_insert(0);
+    let updated_balance = prior_balance.checked_add(amount).ok_or_else(|| {
+        AppErr::set_status(
+            anyhow!("deposit would overflow account"),
+            StatusCode::CONFLICT,
+        )
+    })?;
+
+    tracing::info!(
+        "Updated account public_key={} balance={}",
+        public_key,
+        updated_balance
+    );
+
+    Ok(())
 }

kairos-server/src/routes/mod.rs

@@ -1,7 +1,7 @@
 pub mod deposit;
-pub mod withdraw;
 pub mod transfer;
+pub mod withdraw;
 
 pub use deposit::deposit;
-pub use withdraw::withdraw;
 pub use transfer::transfer;
+pub use withdraw::withdraw;

kairos-server/src/routes/transfer.rs

@@ -1,10 +1,8 @@
-use std::sync::Arc;
-
-use axum::{extract::State, Json};
+use anyhow::anyhow;
+use axum::{extract::State, http::StatusCode, Json};
 use serde::{Deserialize, Serialize};
-use tokio::sync::RwLock;
 
-use crate::{AppErr, BatchState, PublicKey};
+use crate::{state::LockedBatchState, AppErr, PublicKey};
 
 #[derive(Serialize, Deserialize)]
 pub struct Transfer {
@@ -20,8 +18,91 @@ pub struct TransferRequest {
 }
 
 pub async fn transfer(
-    State(pool): State<Arc<RwLock<BatchState>>>,
-    Json(proof_request): Json<TransferRequest>,
+    State(state): State<LockedBatchState>,
+    Json(TransferRequest {
+        transfer,
+        signature: _,
+    }): Json<TransferRequest>,
 ) -> Result<(), AppErr> {
-    todo!()
+    if transfer.amount == 0 {
+        return Err(AppErr::set_status(
+            anyhow!("transfer amount must be greater than 0"),
+            StatusCode::BAD_REQUEST,
+        ));
+    }
+
+    tracing::info!("TODO: verifying transfer signature");
+
+    // We pre-check this read-only to error early without acquiring the write lock.
+    // This prevents a DoS attack exploiting the write lock.
+    tracing::info!("verifying transfer sender has sufficient funds");
+    check_sender_funds(&state, &transfer).await?;
+
+    let mut state = state.write().await;
+    let from_balance = state.balances.get_mut(&transfer.from).ok_or_else(|| {
+        AppErr::set_status(
+            anyhow!(
+                "Sender no longer has an account.
+                The sender just removed all their funds."
+            ),
+            StatusCode::CONFLICT,
+        )
+    })?;
+
+    *from_balance = from_balance.checked_sub(transfer.amount).ok_or_else(|| {
+        AppErr::set_status(
+            anyhow!(
+                "Sender no longer has sufficient funds, balance={}, transfer_amount={}.
+                The sender just moved their funds in a concurrent request",
+                from_balance,
+                transfer.amount
+            ),
+            StatusCode::CONFLICT,
+        )
+    })?;
+
+    let to_balance = state
+        .balances
+        .entry(transfer.to.clone())
+        .or_insert_with(|| {
+            tracing::info!("creating new account for receiver");
+            0
+        });
+
+    *to_balance = to_balance.checked_add(transfer.amount).ok_or_else(|| {
+        AppErr::set_status(anyhow!("Receiver balance overflow"), StatusCode::CONFLICT)
+    })?;
+
+    Ok(())
+}
+
+async fn check_sender_funds(state: &LockedBatchState, transfer: &Transfer) -> Result<(), AppErr> {
+    let state = state.read().await;
+    let from_balance = state.balances.get(&transfer.from).ok_or_else(|| {
+        AppErr::set_status(
+            anyhow!("Sender does not have an account"),
+            StatusCode::BAD_REQUEST,
+        )
+    })?;
+
+    from_balance.checked_sub(transfer.amount).ok_or_else(|| {
+        AppErr::set_status(
+            anyhow!(
+                "Sender does not have sufficient funds, balance={}, transfer_amount={}",
+                from_balance,
+                transfer.amount
+            ),
+            StatusCode::FORBIDDEN,
+        )
+    })?;
+
+    let to_balance = state.balances.get(&transfer.to).unwrap_or(&0);
+    if to_balance.checked_add(transfer.amount).is_none() {
+        return Err(AppErr::set_status(
+            anyhow!("Receiver balance overflow"),
+            StatusCode::CONFLICT,
+        ));
+    }
+
+    Ok(())
 }

kairos-server/src/routes/withdraw.rs

@@ -1,20 +1,87 @@
-use std::sync::Arc;
-
-use axum::{extract::State, Json};
+use anyhow::anyhow;
+use axum::{extract::State, http::StatusCode, Json};
 use serde::{Deserialize, Serialize};
-use tokio::sync::RwLock;
 
-use crate::{AppErr, BatchState, PublicKey};
+use crate::{state::LockedBatchState, AppErr, PublicKey};
 
 #[derive(Serialize, Deserialize)]
-pub struct WithdrawRequest {
+pub struct Withdrawal {
     pub public_key: PublicKey,
+    pub signature: String,
     pub amount: u64,
 }
 
 pub async fn withdraw(
-    State(pool): State<Arc<RwLock<BatchState>>>,
-    Json(proof_request): Json<WithdrawRequest>,
+    State(state): State<LockedBatchState>,
+    Json(withdrawal): Json<Withdrawal>,
+) -> Result<(), AppErr> {
+    tracing::info!("TODO: verifying withdrawal signature");
+
+    tracing::info!("verifying withdrawal sender has sufficient funds");
+    check_sender_funds(&state, &withdrawal).await?;
+
+    tracing::info!("TODO: adding withdrawal to batch");
+
+    let mut state = state.write().await;
+    let from_balance = state
+        .balances
+        .get_mut(&withdrawal.public_key)
+        .ok_or_else(|| {
+            AppErr::set_status(
+                anyhow!(
+                    "Sender no longer has an account.
+                The sender just removed all their funds."
+                ),
+                StatusCode::CONFLICT,
+            )
+        })?;
+
+    let updated_balance = from_balance.checked_sub(withdrawal.amount).ok_or_else(|| {
+        AppErr::set_status(
+            anyhow!(
+                "Sender no longer has sufficient funds, balance={}, withdrawal_amount={}.
+                The sender just moved their funds in a concurrent request",
+                from_balance,
+                withdrawal.amount
+            ),
+            StatusCode::CONFLICT,
+        )
+    })?;
+
+    *from_balance = updated_balance;
+
+    if updated_balance == 0 {
+        state.balances.remove(&withdrawal.public_key);
+    }
+
+    tracing::info!(
+        "Updated account public_key={} balance={}",
+        withdrawal.public_key,
+        updated_balance
+    );
+
+    Ok(())
+}
+
+async fn check_sender_funds(
+    state: &LockedBatchState,
+    withdrawal: &Withdrawal,
 ) -> Result<(), AppErr> {
-    todo!()
+    let state = state.read().await;
+    let from_balance = state.balances.get(&withdrawal.public_key).ok_or_else(|| {
+        AppErr::set_status(anyhow!("Withdrawer has no account."), StatusCode::CONFLICT)
+    })?;
+
+    if *from_balance < withdrawal.amount {
+        return Err(AppErr::set_status(
+            anyhow!(
+                "Withdrawer has insufficient funds, balance={}, withdrawal_amount={}.",
+                from_balance,
+                withdrawal.amount
+            ),
+            StatusCode::FORBIDDEN,
+        ));
+    }
+
+    Ok(())
 }