collab: Add ability to revoke LLM service access tokens (#16143)

Marshall Bowers created 1 year ago

This PR adds the ability to revoke access tokens for the LLM service.

There is a new `revoked_access_tokens` table that contains the
identifiers (`jti`) of revoked access tokens.

To revoke an access token, insert a record into this table:

```sql
insert into revoked_access_tokens (jti) values ('1e887b9e-37f5-49e8-8feb-3274e5a86b67');
```

We now attach the `jti` as `authn.jti` to the tracing spans so that we
can associate an access token with a given request to the LLM service.

Release Notes:

- N/A

Change summary

crates/collab/migrations_llm/20240813002237_add_revoked_access_tokens_table.sql |  7 
crates/collab/src/llm.rs                                                        |  9 
crates/collab/src/llm/db/ids.rs                                                 |  1 
crates/collab/src/llm/db/queries.rs                                             |  1 
crates/collab/src/llm/db/queries/revoked_access_tokens.rs                       | 15 
crates/collab/src/llm/db/tables.rs                                              |  1 
crates/collab/src/llm/db/tables/revoked_access_token.rs                         | 19 
crates/collab/src/main.rs                                                       |  1 
8 files changed, 54 insertions(+)

Detailed changes

crates/collab/migrations_llm/20240813002237_add_revoked_access_tokens_table.sql

@@ -0,0 +1,7 @@
+create table revoked_access_tokens (
+    id serial primary key,
+    jti text not null,
+    revoked_at timestamp without time zone not null default now()
+);
+
+create unique index uix_revoked_access_tokens_on_jti on revoked_access_tokens (jti);

crates/collab/src/llm.rs

@@ -131,6 +131,15 @@ async fn validate_api_token<B>(mut req: Request<B>, next: Next<B>) -> impl IntoR
     let state = req.extensions().get::<Arc<LlmState>>().unwrap();
     match LlmTokenClaims::validate(&token, &state.config) {
         Ok(claims) => {
+            if state.db.is_access_token_revoked(&claims.jti).await? {
+                return Err(Error::http(
+                    StatusCode::UNAUTHORIZED,
+                    "unauthorized".to_string(),
+                ));
+            }
+
+            tracing::Span::current().record("authn.jti", &claims.jti);
+
             req.extensions_mut().insert(claims);
             Ok::<_, Error>(next.run(req).await.into_response())
         }

crates/collab/src/llm/db/ids.rs

@@ -7,3 +7,4 @@ id_type!(ModelId);
 id_type!(ProviderId);
 id_type!(UsageId);
 id_type!(UsageMeasureId);
+id_type!(RevokedAccessTokenId);

crates/collab/src/llm/db/queries.rs

@@ -1,4 +1,5 @@
 use super::*;
 
 pub mod providers;
+pub mod revoked_access_tokens;
 pub mod usages;

crates/collab/src/llm/db/queries/revoked_access_tokens.rs

@@ -0,0 +1,15 @@
+use super::*;
+
+impl LlmDatabase {
+    /// Returns whether the access token with the given `jti` has been revoked.
+    pub async fn is_access_token_revoked(&self, jti: &str) -> Result<bool> {
+        self.transaction(|tx| async move {
+            Ok(revoked_access_token::Entity::find()
+                .filter(revoked_access_token::Column::Jti.eq(jti))
+                .one(&*tx)
+                .await?
+                .is_some())
+        })
+        .await
+    }
+}

crates/collab/src/llm/db/tables.rs

@@ -1,5 +1,6 @@
 pub mod lifetime_usage;
 pub mod model;
 pub mod provider;
+pub mod revoked_access_token;
 pub mod usage;
 pub mod usage_measure;

crates/collab/src/llm/db/tables/revoked_access_token.rs

@@ -0,0 +1,19 @@
+use chrono::NaiveDateTime;
+use sea_orm::entity::prelude::*;
+
+use crate::llm::db::RevokedAccessTokenId;
+
+/// A revoked access token.
+#[derive(Clone, Debug, PartialEq, DeriveEntityModel)]
+#[sea_orm(table_name = "revoked_access_tokens")]
+pub struct Model {
+    #[sea_orm(primary_key)]
+    pub id: RevokedAccessTokenId,
+    pub jti: String,
+    pub revoked_at: NaiveDateTime,
+}
+
+#[derive(Copy, Clone, Debug, EnumIter, DeriveRelation)]
+pub enum Relation {}
+
+impl ActiveModelBehavior for ActiveModel {}

crates/collab/src/main.rs

@@ -150,6 +150,7 @@ async fn main() -> Result<()> {
                             "http_request",
                             method = ?request.method(),
                             matched_path,
+                            authn.jti = tracing::field::Empty
                         )
                     })
                     .on_response(