Merge pull request #11222 from Turbo87/trustpub-cleanup

Add `trustpub::DeleteExpiredJtis/Tokens` background jobs

Author: Turbo87

crates/crates_io_database/src/models/trustpub/mod.rs

@@ -1,3 +1,7 @@
 mod github_config;
+mod token;
+mod used_jti;
 
 pub use self::github_config::{GitHubConfig, NewGitHubConfig};
+pub use self::token::NewToken;
+pub use self::used_jti::NewUsedJti;

crates/crates_io_database/src/models/trustpub/token.rs

@@ -0,0 +1,22 @@
+use crate::schema::trustpub_tokens;
+use chrono::{DateTime, Utc};
+use diesel::prelude::*;
+use diesel_async::{AsyncPgConnection, RunQueryDsl};
+
+#[derive(Debug, Insertable)]
+#[diesel(table_name = trustpub_tokens, check_for_backend(diesel::pg::Pg))]
+pub struct NewToken<'a> {
+    pub expires_at: DateTime<Utc>,
+    pub hashed_token: &'a [u8],
+    pub crate_ids: &'a [i32],
+}
+
+impl NewToken<'_> {
+    pub async fn insert(&self, conn: &mut AsyncPgConnection) -> QueryResult<()> {
+        self.insert_into(trustpub_tokens::table)
+            .execute(conn)
+            .await?;
+
+        Ok(())
+    }
+}

crates/crates_io_database/src/models/trustpub/used_jti.rs

@@ -0,0 +1,24 @@
+use crate::schema::trustpub_used_jtis;
+use chrono::{DateTime, Utc};
+use diesel::prelude::*;
+use diesel_async::{AsyncPgConnection, RunQueryDsl};
+
+#[derive(Debug, Insertable)]
+#[diesel(table_name = trustpub_used_jtis, check_for_backend(diesel::pg::Pg))]
+pub struct NewUsedJti<'a> {
+    pub jti: &'a str,
+    pub expires_at: DateTime<Utc>,
+}
+
+impl<'a> NewUsedJti<'a> {
+    pub fn new(jti: &'a str, expires_at: DateTime<Utc>) -> Self {
+        Self { jti, expires_at }
+    }
+
+    pub async fn insert(&self, conn: &mut AsyncPgConnection) -> QueryResult<usize> {
+        diesel::insert_into(trustpub_used_jtis::table)
+            .values(self)
+            .execute(conn)
+            .await
+    }
+}

src/bin/crates-admin/enqueue_job.rs

@@ -49,6 +49,7 @@ pub enum Command {
         name: String,
     },
     SyncUpdatesFeed,
+    TrustpubCleanup,
 }
 
 pub async fn run(command: Command) -> Result<()> {
@@ -161,6 +162,13 @@ pub async fn run(command: Command) -> Result<()> {
         Command::SyncUpdatesFeed => {
             jobs::rss::SyncUpdatesFeed.enqueue(&mut conn).await?;
         }
+        Command::TrustpubCleanup => {
+            let job = jobs::trustpub::DeleteExpiredTokens;
+            job.enqueue(&mut conn).await?;
+
+            let job = jobs::trustpub::DeleteExpiredJtis;
+            job.enqueue(&mut conn).await?;
+        }
     };
 
     Ok(())

src/worker/jobs/mod.rs

@@ -12,6 +12,7 @@ mod readmes;
 pub mod rss;
 mod send_publish_notifications;
 mod sync_admins;
+pub mod trustpub;
 mod typosquat;
 mod update_default_version;
 

src/worker/jobs/trustpub/delete_jtis.rs

@@ -0,0 +1,68 @@
+use crate::worker::Environment;
+use crates_io_database::schema::trustpub_used_jtis;
+use crates_io_worker::BackgroundJob;
+use diesel::prelude::*;
+use diesel_async::RunQueryDsl;
+use std::sync::Arc;
+
+/// A background job that deletes expired JSON Web Token IDs (JTIs)
+/// tokens from the database.
+#[derive(Deserialize, Serialize)]
+pub struct DeleteExpiredJtis;
+
+impl BackgroundJob for DeleteExpiredJtis {
+    const JOB_NAME: &'static str = "trustpub::delete_expired_jtis";
+
+    type Context = Arc<Environment>;
+
+    async fn run(&self, ctx: Self::Context) -> anyhow::Result<()> {
+        let mut conn = ctx.deadpool.get().await?;
+
+        diesel::delete(trustpub_used_jtis::table)
+            .filter(trustpub_used_jtis::expires_at.lt(diesel::dsl::now))
+            .execute(&mut conn)
+            .await?;
+
+        Ok(())
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::tests::util::TestApp;
+    use chrono::{TimeDelta, Utc};
+    use crates_io_database::models::trustpub::NewUsedJti;
+    use insta::assert_compact_debug_snapshot;
+
+    #[tokio::test(flavor = "multi_thread")]
+    async fn test_expiry() -> anyhow::Result<()> {
+        let (app, _client) = TestApp::full().empty().await;
+        let mut conn = app.db_conn().await;
+
+        let jti = NewUsedJti {
+            expires_at: Utc::now() + TimeDelta::minutes(30),
+            jti: "foo",
+        };
+        jti.insert(&mut conn).await?;
+
+        let jti = NewUsedJti {
+            expires_at: Utc::now() - TimeDelta::minutes(5),
+            jti: "bar",
+        };
+        jti.insert(&mut conn).await?;
+
+        DeleteExpiredJtis.enqueue(&mut conn).await?;
+        app.run_pending_background_jobs().await;
+
+        // Check that the expired token was deleted
+        let known_jtis: Vec<String> = trustpub_used_jtis::table
+            .select(trustpub_used_jtis::jti)
+            .load(&mut conn)
+            .await?;
+
+        assert_compact_debug_snapshot!(known_jtis, @r#"["foo"]"#);
+
+        Ok(())
+    }
+}

src/worker/jobs/trustpub/delete_tokens.rs

@@ -0,0 +1,70 @@
+use crate::worker::Environment;
+use crates_io_database::schema::trustpub_tokens;
+use crates_io_worker::BackgroundJob;
+use diesel::prelude::*;
+use diesel_async::RunQueryDsl;
+use std::sync::Arc;
+
+/// A background job that deletes expired temporary access
+/// tokens from the database.
+#[derive(Deserialize, Serialize)]
+pub struct DeleteExpiredTokens;
+
+impl BackgroundJob for DeleteExpiredTokens {
+    const JOB_NAME: &'static str = "trustpub::delete_expired_tokens";
+
+    type Context = Arc<Environment>;
+
+    async fn run(&self, ctx: Self::Context) -> anyhow::Result<()> {
+        let mut conn = ctx.deadpool.get().await?;
+
+        diesel::delete(trustpub_tokens::table)
+            .filter(trustpub_tokens::expires_at.lt(diesel::dsl::now))
+            .execute(&mut conn)
+            .await?;
+
+        Ok(())
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::tests::util::TestApp;
+    use chrono::{TimeDelta, Utc};
+    use crates_io_database::models::trustpub::NewToken;
+    use insta::assert_compact_debug_snapshot;
+
+    #[tokio::test(flavor = "multi_thread")]
+    async fn test_expiry() -> anyhow::Result<()> {
+        let (app, _client) = TestApp::full().empty().await;
+        let mut conn = app.db_conn().await;
+
+        let token = NewToken {
+            expires_at: Utc::now() + TimeDelta::minutes(30),
+            hashed_token: &[0xC0, 0xFF, 0xEE],
+            crate_ids: &[1],
+        };
+        token.insert(&mut conn).await?;
+
+        let token = NewToken {
+            expires_at: Utc::now() - TimeDelta::minutes(5),
+            hashed_token: &[0xBA, 0xAD, 0xF0, 0x0D],
+            crate_ids: &[2],
+        };
+        token.insert(&mut conn).await?;
+
+        DeleteExpiredTokens.enqueue(&mut conn).await?;
+        app.run_pending_background_jobs().await;
+
+        // Check that the expired token was deleted
+        let crate_ids: Vec<Vec<Option<i32>>> = trustpub_tokens::table
+            .select(trustpub_tokens::crate_ids)
+            .load(&mut conn)
+            .await?;
+
+        assert_compact_debug_snapshot!(crate_ids, @"[[Some(1)]]");
+
+        Ok(())
+    }
+}

src/worker/jobs/trustpub/mod.rs

@@ -0,0 +1,5 @@
+mod delete_jtis;
+mod delete_tokens;
+
+pub use delete_jtis::DeleteExpiredJtis;
+pub use delete_tokens::DeleteExpiredTokens;

src/worker/mod.rs

@@ -43,5 +43,7 @@ impl RunnerExt for Runner<Arc<Environment>> {
             .register_job_type::<jobs::rss::SyncCrateFeed>()
             .register_job_type::<jobs::rss::SyncCratesFeed>()
             .register_job_type::<jobs::rss::SyncUpdatesFeed>()
+            .register_job_type::<jobs::trustpub::DeleteExpiredJtis>()
+            .register_job_type::<jobs::trustpub::DeleteExpiredTokens>()
     }
 }