Dynamo cache

Author: MaeIsBad

Cargo.toml

@@ -15,7 +15,7 @@ default = ["tracing_opentelemetry"]
 
 # Feature set that should be used
 # This exists to avoid compatibility issues with otel version conflicts
-_docs = ["auth0", "grpc", "gzip", "redis-tls", "tracing_opentelemetry"]
+_docs = ["auth0", "cache-dynamodb", "grpc", "gzip", "redis-tls", "tracing_opentelemetry"]
 
 auth0 = [
   "rand",
@@ -30,6 +30,8 @@ grpc = [ "_any_otel_version", "tonic"]
 gzip = ["reqwest/gzip"]
 
 redis-tls = [ "redis", "redis/tls", "redis/tokio-native-tls-comp"]
+cache-dynamodb = [ "aws-sdk-dynamodb" ]
+
 tracing_opentelemetry = ["tracing_opentelemetry_0_27"]
 
 tracing_opentelemetry_0_21 = [
@@ -107,6 +109,7 @@ tonic = { version = "0.12", default-features = false, optional = true }
 tracing = { version = "0.1", optional = true }
 uuid = { version = ">=0.7.0, <2.0.0", features = ["serde", "v4"] }
 chacha20poly1305 = { version = "0.10.1", features = ["std"], optional = true }
+aws-sdk-dynamodb = { version = "1.63", optional = true } 
 
 reqwest-middleware = { version = "0.4.0", features = ["json", "multipart"] }
 http = "1.0.0"

src/auth0/cache/dynamodb.rs

@@ -1 +1,148 @@
+use std::{error::Error, time::Duration};
 
+pub use aws_sdk_dynamodb;
+use aws_sdk_dynamodb::{
+    client::Waiters,
+    operation::describe_table::DescribeTableError,
+    types::{AttributeDefinition, AttributeValue, TimeToLiveSpecification},
+};
+
+use crate::auth0::Token;
+
+use super::{Cache, CacheError};
+
+#[derive(thiserror::Error, Debug)]
+pub enum DynamoDBCacheError {
+    #[error("AWS error when interacting with dynamo cache: {0}")]
+    Aws(Box<dyn Error>),
+    #[error("Data in database is wrong. Key: {0}")]
+    SchemaError(String),
+}
+
+impl From<DynamoDBCacheError> for super::CacheError {
+    fn from(val: DynamoDBCacheError) -> Self {
+        CacheError(Box::new(val))
+    }
+}
+
+#[derive(Debug)]
+pub struct DynamoDBCache {
+    table_name: String,
+    client: aws_sdk_dynamodb::Client,
+}
+
+impl DynamoDBCache {
+    /// Construct a DynamoDBCache instance which uses a given table name and client
+    ///
+    /// Note: this method doesn't currectly check whether a table with the given name exists during creation.
+    pub fn new(client: aws_sdk_dynamodb::Client, table_name: String) -> Self {
+        Self { client, table_name }
+    }
+
+    pub async fn create_table_if_not_exists(&self) -> Result<(), DynamoDBCacheError> {
+        match self
+            .client
+            .describe_table()
+            .table_name(&self.table_name)
+            .send()
+            .await
+            .map_err(|e| e.into_service_error())
+        {
+            Ok(_) => return Ok(()),
+            Err(DescribeTableError::ResourceNotFoundException(_)) => (),
+            Err(e) => return Err(DynamoDBCacheError::Aws(Box::new(e))),
+        };
+
+        self.client
+            .create_table()
+            .table_name(self.table_name.clone())
+            .attribute_definitions(
+                AttributeDefinition::builder()
+                    .attribute_name("key".to_string())
+                    .attribute_type(aws_sdk_dynamodb::types::ScalarAttributeType::S)
+                    .build()
+                    // Unwraps here are fine, will be caught by tests
+                    .unwrap(),
+            )
+            .attribute_definitions(
+                AttributeDefinition::builder()
+                    .attribute_name("expiration".to_string())
+                    .attribute_type(aws_sdk_dynamodb::types::ScalarAttributeType::N)
+                    .build()
+                    .unwrap(),
+            )
+            .send()
+            .await
+            .map_err(|e| DynamoDBCacheError::Aws(Box::new(e)))?;
+
+        self.client
+            .wait_until_table_exists()
+            .table_name(&self.table_name)
+            .wait(Duration::from_secs(5))
+            .await
+            .map_err(|e| DynamoDBCacheError::Aws(Box::new(e)))?;
+
+        self.client
+            .update_time_to_live()
+            .table_name(self.table_name.clone())
+            .time_to_live_specification(
+                TimeToLiveSpecification::builder()
+                    .enabled(true)
+                    .attribute_name("expiration")
+                    .build()
+                    .unwrap(),
+            )
+            .send()
+            .await
+            .map_err(|e| DynamoDBCacheError::Aws(Box::new(e)))?;
+
+        Ok(())
+    }
+}
+
+#[async_trait::async_trait]
+impl Cache for DynamoDBCache {
+    async fn get_token(&self, client_id: &str, aud: &str) -> Result<Option<Token>, CacheError> {
+        let key = super::token_key(client_id, aud);
+        let Some(attrs) = self
+            .client
+            .get_item()
+            .table_name(&self.table_name)
+            .key("key", AttributeValue::S(key.clone()))
+            .send()
+            .await
+            .map_err(|e| DynamoDBCacheError::Aws(Box::new(e)))?
+            .item
+        else {
+            return Ok(None);
+        };
+
+        let token = attrs
+            .get("token")
+            .and_then(|t| t.as_s().ok())
+            .ok_or(DynamoDBCacheError::SchemaError(key.clone()))?;
+
+        let token: Token = serde_json::from_str(token).unwrap();
+
+        Ok(Some(token))
+    }
+
+    async fn put_token(&self, client_id: &str, aud: &str, token: &Token) -> Result<(), CacheError> {
+        let key = super::token_key(client_id, aud);
+        let encoded = serde_json::to_string(token).unwrap();
+        self.client
+            .put_item()
+            .table_name(&self.table_name)
+            .item("key", AttributeValue::S(key))
+            .item("token", AttributeValue::S(encoded))
+            .item(
+                "expiration",
+                AttributeValue::N(token.expire_date().timestamp().to_string()),
+            )
+            .send()
+            .await
+            .map_err(|e| DynamoDBCacheError::Aws(Box::new(e)))?;
+
+        Ok(())
+    }
+}

src/auth0/cache/mod.rs

@@ -1,13 +1,14 @@
+pub use dynamodb::DynamoDBCache;
 pub use inmemory::InMemoryCache;
 pub use redis_impl::RedisCache;
 use std::error::Error;
 
 use crate::auth0::Token;
 
 mod crypto;
-mod dynamodb;
-mod inmemory;
-mod redis_impl;
+pub mod dynamodb;
+pub mod inmemory;
+pub mod redis_impl;
 
 const TOKEN_PREFIX: &str = "auth0rs_tokens";
 // The version of the token for backwards incompatible changes