sessions: JWT (#19)

* jwt token sessions

# Conflicts:
#	crates/embucket-lambda/src/main.rs

* Requested changes

* clippy

* fmt

* clippy

Author: DanCodedThis

Cargo.lock

@@ -19,6 +19,8 @@ snafu = { workspace = true }
 tracing = { workspace = true }
 uuid = { workspace = true }
 regex = { workspace = true }
+jsonwebtoken = { workspace = true }
+chrono = { workspace = true }
 
 [lints]
 workspace = true

crates/api-snowflake-rest-sessions/Cargo.toml

@@ -2,6 +2,7 @@ use axum::{Json, http, response::IntoResponse};
 use error_stack_trace;
 use http::header::InvalidHeaderValue;
 use http::{StatusCode, header::MaxSizeReached};
+use jsonwebtoken::errors::Error as JwtError;
 use serde::{Deserialize, Serialize};
 use snafu::Location;
 use snafu::prelude::*;
@@ -33,6 +34,14 @@ pub enum Error {
         #[snafu(source)]
         error: executor::Error,
     },
+
+    #[snafu(display("Bad authentication token. {error}"))]
+    BadAuthToken {
+        #[snafu(source)]
+        error: JwtError,
+        #[snafu(implicit)]
+        location: Location,
+    },
 }
 
 #[derive(Debug, Serialize, Deserialize)]

crates/api-snowflake-rest-sessions/src/error.rs

@@ -0,0 +1,62 @@
+use chrono::offset::Local;
+use jsonwebtoken::{DecodingKey, EncodingKey, Header, Validation, decode, encode};
+use serde::{Deserialize, Serialize};
+use time::Duration;
+use uuid::Uuid;
+
+#[derive(Serialize, Deserialize)]
+#[cfg_attr(test, derive(Debug))]
+pub struct Claims {
+    pub sub: String, // token issued to a particular user
+    pub iat: i64,    // Issued At
+    pub exp: i64,    // Expiration Time
+    pub session_id: String,
+}
+
+#[must_use]
+pub fn jwt_claims(username: &str, expiration: Duration) -> Claims {
+    let now = Local::now();
+    let iat = now.timestamp();
+    let exp = now.timestamp() + expiration.whole_seconds();
+
+    Claims {
+        sub: username.to_string(),
+        iat,
+        exp,
+        session_id: Uuid::new_v4().to_string(),
+    }
+}
+
+pub fn get_claims_validate_jwt_token(
+    token: &str,
+    jwt_secret: &str,
+) -> Result<Claims, jsonwebtoken::errors::Error> {
+    let mut validation = Validation::default();
+    validation.leeway = 5;
+    validation.set_required_spec_claims(&["exp"]);
+
+    let decoding_key = DecodingKey::from_secret(jwt_secret.as_bytes());
+
+    let decoded = decode::<Claims>(token, &decoding_key, &validation)?;
+
+    Ok(decoded.claims)
+}
+
+pub fn create_jwt<T>(claims: &T, jwt_secret: &str) -> Result<String, jsonwebtoken::errors::Error>
+where
+    T: Serialize,
+{
+    encode(
+        &Header::default(),
+        &claims,
+        &EncodingKey::from_secret(jwt_secret.as_bytes()),
+    )
+}
+
+#[must_use]
+pub fn ensure_jwt_secret_is_valid(jwt_secret: &str) -> Option<String> {
+    if jwt_secret.is_empty() {
+        return None;
+    }
+    Some(jwt_secret.to_string())
+}

crates/api-snowflake-rest-sessions/src/helpers.rs

@@ -1,4 +1,5 @@
 pub mod error;
+pub mod helpers;
 pub mod layer;
 pub mod session;
 

crates/api-snowflake-rest-sessions/src/lib.rs

@@ -1,4 +1,6 @@
 use crate::error as session_error;
+use crate::error::BadAuthTokenSnafu;
+use crate::helpers::get_claims_validate_jwt_token;
 use axum::extract::FromRequestParts;
 use executor::ExecutionAppState;
 use executor::service::ExecutionService;
@@ -32,12 +34,16 @@ impl SessionStore {
     }
 }
 
+pub trait JwtSecret {
+    fn jwt_secret(&self) -> &str;
+}
+
 #[derive(Debug, Clone)]
 pub struct DFSessionId(pub String);
 
 impl<S> FromRequestParts<S> for DFSessionId
 where
-    S: Send + Sync + ExecutionAppState,
+    S: Send + Sync + ExecutionAppState + JwtSecret,
 {
     type Rejection = session_error::Error;
 
@@ -47,7 +53,11 @@ where
         let execution_svc = state.get_execution_svc();
 
         let (session_id, located_at) = if let Some(token) = extract_token_from_auth(&req.headers) {
-            (token, "auth header")
+            let jwt_secret = state.jwt_secret();
+            let jwt_claims =
+                get_claims_validate_jwt_token(&token, jwt_secret).context(BadAuthTokenSnafu)?;
+
+            (jwt_claims.session_id, "auth header")
         } else {
             //This is guaranteed by the `propagate_session_cookie`, so we can unwrap
             let Self(token) = req.extensions.get::<Self>().unwrap();
@@ -97,7 +107,9 @@ pub fn extract_token_from_auth(headers: &HeaderMap) -> Option<String> {
     headers.get("authorization").and_then(|value| {
         value.to_str().ok().and_then(|auth| {
             #[allow(clippy::unwrap_used)]
-            let re = Regex::new(r#"Snowflake Token="([a-f0-9\-]+)""#).unwrap();
+            let re = Regex::new(
+                r#"Snowflake Token="([A-Za-z0-9_\-]+\.[A-Za-z0-9_\-]+\.[A-Za-z0-9_\-]+|[0-9a-fA-F]{8}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{12})""#
+            ).unwrap();
             re.captures(auth)
                 .and_then(|caps| caps.get(1).map(|m| m.as_str().to_string()))
         })

crates/api-snowflake-rest-sessions/src/session.rs

@@ -51,6 +51,7 @@ time = { workspace = true }
 uuid = { workspace = true }
 tokio = { workspace = true }
 cfg-if = { workspace = true }
+jsonwebtoken = { workspace = true }
 
 [dev-dependencies]
 insta = { workspace = true }

crates/api-snowflake-rest/Cargo.toml

@@ -141,4 +141,15 @@ impl From<ColumnInfoModel> for ColumnInfo {
 pub struct Auth {
     pub demo_user: String,
     pub demo_password: String,
+    pub jwt_secret: String,
+}
+
+impl Auth {
+    #[must_use]
+    pub fn new(jwt_secret: String) -> Self {
+        Self {
+            jwt_secret,
+            ..Self::default()
+        }
+    }
 }

crates/api-snowflake-rest/src/models.rs

@@ -10,6 +10,7 @@ use executor::QueryRecordId;
 use executor::error::OperationOn;
 use executor::error_code::ErrorCode;
 use executor::snowflake_error::Entity;
+use jsonwebtoken::errors::Error as JwtError;
 use snafu::Location;
 use snafu::prelude::*;
 
@@ -70,6 +71,28 @@ pub enum Error {
 
     #[snafu(transparent)]
     Execution { source: executor::Error },
+
+    #[snafu(display("JWT secret is not set"))]
+    NoJwtSecret {
+        #[snafu(implicit)]
+        location: Location,
+    },
+
+    #[snafu(display("Failed to create JWT: {error}"))]
+    CreateJwt {
+        #[snafu(source)]
+        error: JwtError,
+        #[snafu(implicit)]
+        location: Location,
+    },
+
+    #[snafu(display("Bad authentication token. {error}"))]
+    BadAuthToken {
+        #[snafu(source)]
+        error: JwtError,
+        #[snafu(implicit)]
+        location: Location,
+    },
 }
 
 impl IntoResponse for Error {
@@ -190,7 +213,10 @@ impl Error {
             }
             Self::MissingAuthToken { .. }
             | Self::InvalidAuthData { .. }
-            | Self::InvalidAuthToken { .. } => (
+            | Self::InvalidAuthToken { .. }
+            | Self::NoJwtSecret { .. }
+            | Self::CreateJwt { .. }
+            | Self::BadAuthToken { .. } => (
                 http::StatusCode::UNAUTHORIZED,
                 SqlState::Success,
                 ErrorCode::Other,

crates/api-snowflake-rest/src/server/error.rs

@@ -14,6 +14,7 @@ use executor::utils::{
 };
 use serde_json::value::RawValue;
 use snafu::ResultExt;
+use tracing;
 use uuid::Uuid;
 
 // https://arrow.apache.org/docs/format/Columnar.html#buffer-alignment-and-padding