diff --git a/crates/crates_io_database/src/models/trustpub/mod.rs b/crates/crates_io_database/src/models/trustpub/mod.rs index 634f500f46..6a2ad6357b 100644 --- a/crates/crates_io_database/src/models/trustpub/mod.rs +++ b/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; diff --git a/crates/crates_io_database/src/models/trustpub/token.rs b/crates/crates_io_database/src/models/trustpub/token.rs new file mode 100644 index 0000000000..80e6fcf5c8 --- /dev/null +++ b/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, + 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(()) + } +} diff --git a/crates/crates_io_database/src/models/trustpub/used_jti.rs b/crates/crates_io_database/src/models/trustpub/used_jti.rs new file mode 100644 index 0000000000..eced690bf0 --- /dev/null +++ b/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, +} + +impl<'a> NewUsedJti<'a> { + pub fn new(jti: &'a str, expires_at: DateTime) -> Self { + Self { jti, expires_at } + } + + pub async fn insert(&self, conn: &mut AsyncPgConnection) -> QueryResult { + diesel::insert_into(trustpub_used_jtis::table) + .values(self) + .execute(conn) + .await + } +} diff --git a/src/bin/crates-admin/enqueue_job.rs b/src/bin/crates-admin/enqueue_job.rs index 192fd56227..1f85346ef2 100644 --- a/src/bin/crates-admin/enqueue_job.rs +++ b/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(()) diff --git a/src/worker/jobs/mod.rs b/src/worker/jobs/mod.rs index f7cbd6434b..af8a008a65 100644 --- a/src/worker/jobs/mod.rs +++ b/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; diff --git a/src/worker/jobs/trustpub/delete_jtis.rs b/src/worker/jobs/trustpub/delete_jtis.rs new file mode 100644 index 0000000000..6e058fa7b7 --- /dev/null +++ b/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; + + 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 = trustpub_used_jtis::table + .select(trustpub_used_jtis::jti) + .load(&mut conn) + .await?; + + assert_compact_debug_snapshot!(known_jtis, @r#"["foo"]"#); + + Ok(()) + } +} diff --git a/src/worker/jobs/trustpub/delete_tokens.rs b/src/worker/jobs/trustpub/delete_tokens.rs new file mode 100644 index 0000000000..4b528a93e3 --- /dev/null +++ b/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; + + 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>> = trustpub_tokens::table + .select(trustpub_tokens::crate_ids) + .load(&mut conn) + .await?; + + assert_compact_debug_snapshot!(crate_ids, @"[[Some(1)]]"); + + Ok(()) + } +} diff --git a/src/worker/jobs/trustpub/mod.rs b/src/worker/jobs/trustpub/mod.rs new file mode 100644 index 0000000000..04a130afa5 --- /dev/null +++ b/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; diff --git a/src/worker/mod.rs b/src/worker/mod.rs index 0838b8bff4..f22bd29d0a 100644 --- a/src/worker/mod.rs +++ b/src/worker/mod.rs @@ -43,5 +43,7 @@ impl RunnerExt for Runner> { .register_job_type::() .register_job_type::() .register_job_type::() + .register_job_type::() + .register_job_type::() } }