Periodically prune pending tokens

Cargo.lock

@@ -305,6 +305,7 @@ checksum = "7f55667319111d593ba876406af7c409c0ebb44dc4be6132a783ccf163ea14c1"
 dependencies = [
  "futures-channel",
  "futures-core",
+ "futures-executor",
  "futures-io",
  "futures-sink",
  "futures-task",
@@ -327,12 +328,35 @@ version = "0.3.13"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "15496a72fabf0e62bdc3df11a59a3787429221dd0710ba8ef163d6f7a9112c94"
 
+[[package]]
+name = "futures-executor"
+version = "0.3.13"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "891a4b7b96d84d5940084b2a37632dd65deeae662c114ceaa2c879629c9c0ad1"
+dependencies = [
+ "futures-core",
+ "futures-task",
+ "futures-util",
+]
+
 [[package]]
 name = "futures-io"
 version = "0.3.13"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "d71c2c65c57704c32f5241c1223167c2c3294fd34ac020c807ddbe6db287ba59"
 
+[[package]]
+name = "futures-macro"
+version = "0.3.13"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "ea405816a5139fb39af82c2beb921d52143f556038378d6db21183a5c37fbfb7"
+dependencies = [
+ "proc-macro-hack",
+ "proc-macro2",
+ "quote",
+ "syn",
+]
+
 [[package]]
 name = "futures-sink"
 version = "0.3.13"
@@ -351,11 +375,17 @@ version = "0.3.13"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "1812c7ab8aedf8d6f2701a43e1243acdbcc2b36ab26e2ad421eb99ac963d96d1"
 dependencies = [
+ "futures-channel",
  "futures-core",
+ "futures-io",
+ "futures-macro",
  "futures-sink",
  "futures-task",
+ "memchr",
  "pin-project-lite",
  "pin-utils",
+ "proc-macro-hack",
+ "proc-macro-nested",
  "slab",
 ]
 
@@ -897,6 +927,12 @@ version = "0.5.19"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "dbf0c48bc1d91375ae5c3cd81e3722dff1abcf81a30960240640d223f59fe0e5"
 
+[[package]]
+name = "proc-macro-nested"
+version = "0.1.7"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "bc881b2c22681370c6a780e47af9840ef841837bc98118431d4e1868bd0c1086"
+
 [[package]]
 name = "proc-macro2"
 version = "1.0.24"
@@ -1202,6 +1238,7 @@ dependencies = [
  "base64",
  "chrono",
  "curve25519-parser",
+ "futures",
  "hex",
  "hmac",
  "http",

Cargo.toml

@@ -9,6 +9,7 @@ aes-gcm = "0.8"
 base64 = "0.13"
 chrono = "0.4"
 curve25519-parser = "0.2"
+futures = "0.3"
 hex = "0.4"
 hmac = "0.10"
 http = "0.2"

src/handlers.rs

@@ -10,9 +10,6 @@ use super::models;
 use super::rpc;
 use super::storage;
 
-/// The period after which a pending token is expired.
-const TOKEN_EXPIRATION: i64 = 10 * 60;
-
 // TODO: Expire tokens after 10 minutes
 
 pub async fn get_challenge(hex_public_key: String, pool: &storage::DatabaseConnectionPool) -> Result<Response, Rejection> {
@@ -75,7 +72,7 @@ pub async fn claim_token(public_key: String, token: String, pool: &storage::Data
         let raw_query = format!("SELECT timestamp, token FROM {} WHERE public_key = (?1) AND timestamp > (?2)", storage::PENDING_TOKENS_TABLE);
         let mut query = tx.prepare(&raw_query).map_err(|_| Error::DatabaseFailedInternally)?;
         let now = chrono::Utc::now().timestamp();
-        let expiration = now - TOKEN_EXPIRATION;
+        let expiration = now - storage::PENDING_TOKEN_EXPIRATION;
         let rows = match query.query_map(params![ public_key, expiration ], |row| {
             Ok((row.get(0)?, row.get(1)?))
         }) {

src/main.rs

@@ -1,3 +1,5 @@
+use futures::join;
+use tokio;
 use warp::Filter;
 
 mod crypto;
@@ -19,11 +21,12 @@ async fn main() {
     let pool = r2d2::Pool::new(db_manager).unwrap();
     let conn = pool.get().unwrap();
     storage::create_tables_if_needed(&conn);
-    let routes = routes::root().or(routes::lsrpc(pool));
-    warp::serve(routes)
-        .tls()
-        .cert_path("tls_certificate.pem")
-        .key_path("tls_private_key.pem")
-        .run(([0, 0, 0, 0], 443))
-        .await;
+    let f_0 = storage::prune_pending_tokens_periodically(pool.clone());
+    let routes = routes::root().or(routes::lsrpc(pool.clone()));
+    let f_1 = warp::serve(routes)
+        // .tls()
+        // .cert_path("tls_certificate.pem")
+        // .key_path("tls_private_key.pem")
+        .run(([127, 0, 0, 1], 8080));
+    join!(f_0, f_1);
 }

src/routes.rs

@@ -13,13 +13,13 @@ pub fn root() -> impl Filter<Extract = impl warp::Reply, Error = Rejection> + Cl
 
 /// POST /loki/v3/lsrpc
 pub fn lsrpc(
-    db_pool: storage::DatabaseConnectionPool
+    pool: storage::DatabaseConnectionPool
 ) -> impl Filter<Extract = impl warp::Reply, Error = Rejection> + Clone {
     return warp::post()
         .and(warp::path("loki")).and(warp::path("v3")).and(warp::path("lsrpc"))
         .and(warp::body::content_length_limit(10 * 1024 * 1024)) // Match storage server
         .and(warp::body::bytes()) // Expect bytes
-        .and(warp::any().map(move || db_pool.clone()))
+        .and(warp::any().map(move || pool.clone()))
         .and_then(onion_requests::handle_onion_request)
         // It's possible for an error to occur before we have the symmetric key needed
         // to encrypt the response. In this scenario we still want to return a useful

src/storage.rs

@@ -5,6 +5,8 @@ use r2d2_sqlite::SqliteConnectionManager;
 pub type DatabaseConnection = r2d2::PooledConnection<SqliteConnectionManager>;
 pub type DatabaseConnectionPool = r2d2::Pool<SqliteConnectionManager>;
 
+pub const PENDING_TOKEN_EXPIRATION: i64 = 10 * 60;
+
 pub const MESSAGES_TABLE: &str = "messages";
 pub const DELETED_MESSAGES_TABLE: &str = "deleted_messages";
 pub const MODERATORS_TABLE: &str = "moderators";
@@ -57,3 +59,31 @@ pub fn create_tables_if_needed(conn: &DatabaseConnection) {
     )", TOKENS_TABLE);
     conn.execute(&tokens_table_cmd, params![]).expect("Couldn't create tokens table.");
 }
+
+pub async fn prune_pending_tokens_periodically(pool: DatabaseConnectionPool) {
+    let mut timer = tokio::time::interval(chrono::Duration::minutes(10).to_std().unwrap());
+    loop {
+        let pool = pool.clone();
+        timer.tick().await;
+        tokio::spawn(async { prune_pending_tokens(pool).await; });
+    }
+}
+
+async fn prune_pending_tokens(pool: DatabaseConnectionPool) {
+    // It's not catastrophic if we fail to prune the database
+    let mut conn = match pool.get() {
+        Ok(conn) => conn,
+        Err(e) => return println!("Couldn't prune pending tokens due to error: {}.", e)
+    };
+    let tx = match conn.transaction() {
+        Ok(tx) => tx,
+        Err(e) => return println!("Couldn't prune pending tokens due to error: {}.", e)
+    };
+    let stmt = format!("DELETE FROM {} WHERE timestamp < (?1)", PENDING_TOKENS_TABLE);
+    let now = chrono::Utc::now().timestamp();
+    let expiration = now - PENDING_TOKEN_EXPIRATION;
+    match tx.execute(&stmt, params![ expiration ]) {
+        Ok(_) => (),
+        Err(e) => return println!("Couldn't prune pending tokens due to error: {}.", e)
+    };
+}