Skip to main content

atomic_lti_tool/middleware/
jwt_authentication_middleware.rs

1use actix_web::HttpMessage;
2use actix_web::{
3  dev::{forward_ready, Service, ServiceRequest, ServiceResponse, Transform},
4  Error,
5};
6use atomic_lti::jwt::decode;
7use atomic_lti::stores::key_store::KeyStore;
8use jsonwebtoken::decode_header;
9use serde::de::DeserializeOwned;
10use std::boxed::Box;
11use std::sync::Arc;
12use std::{
13  future::{ready, Future, Ready},
14  pin::Pin,
15};
16
17// Fields:
18// key_store: a key store that implements the KeyStore trait wrapped in an Arc
19// marker: PhantomData is a zero-sized type used to mark generic structs and indicate to the Rust compiler that
20// JwtAuthenticationConfig is parameterized over T, even though T is not used in any of the struct's fields.
21pub struct JwtAuthenticationConfig<T> {
22  pub key_store: Arc<dyn KeyStore>,
23  pub marker: std::marker::PhantomData<T>,
24}
25
26impl<T> Clone for JwtAuthenticationConfig<T> {
27  fn clone(&self) -> Self {
28    JwtAuthenticationConfig {
29      key_store: Arc::clone(&self.key_store),
30      marker: std::marker::PhantomData,
31    }
32  }
33}
34
35#[derive(Clone)]
36pub struct JwtAuthentication<T> {
37  pub config: JwtAuthenticationConfig<T>,
38}
39
40impl<T> JwtAuthentication<T> {
41  pub fn new(config: JwtAuthenticationConfig<T>) -> Self {
42    JwtAuthentication { config }
43  }
44}
45
46// `S` - type of the next service
47// `B` - type of response's body
48impl<S, B, T> Transform<S, ServiceRequest> for JwtAuthentication<T>
49where
50  S: Service<ServiceRequest, Response = ServiceResponse<B>, Error = Error> + Clone + 'static,
51  S::Future: 'static,
52  B: 'static,
53  //T: DeserializeOwned + 'static,
54  T: DeserializeOwned + for<'de> serde::Deserialize<'de> + 'static + Clone,
55{
56  type Response = ServiceResponse<B>;
57  type Error = Error;
58  type InitError = ();
59  type Transform = JwtAuthenticationMiddleware<S, T>;
60  type Future = Ready<Result<Self::Transform, Self::InitError>>;
61
62  fn new_transform(&self, service: S) -> Self::Future {
63    ready(Ok(JwtAuthenticationMiddleware {
64      service,
65      config: self.config.clone(),
66    }))
67  }
68}
69
70pub struct JwtAuthenticationMiddleware<S, T> {
71  /// The next service to call
72  service: S,
73  config: JwtAuthenticationConfig<T>,
74}
75
76// This future doesn't have the requirement of being `Send`.
77// See: futures_util::future::LocalBoxFuture
78type LocalBoxFuture<T> = Pin<Box<dyn Future<Output = T> + 'static>>;
79
80// `S`: type of the wrapped service
81// `B`: type of the body - try to be generic over the body where possible
82impl<S, B, T> Service<ServiceRequest> for JwtAuthenticationMiddleware<S, T>
83where
84  S: Service<ServiceRequest, Response = ServiceResponse<B>, Error = Error> + Clone + 'static,
85  S::Future: 'static,
86  B: 'static,
87  T: DeserializeOwned + for<'de> serde::Deserialize<'de> + 'static + Clone,
88{
89  type Response = ServiceResponse<B>;
90  type Error = Error;
91  type Future = LocalBoxFuture<Result<Self::Response, Self::Error>>;
92
93  // This service is ready when its next service is ready
94  forward_ready!(service);
95
96  fn call(&self, req: ServiceRequest) -> Self::Future {
97    let headers = req.headers();
98    let auth_header = headers.get("Authorization");
99
100    let encoded_jwt = match auth_header {
101      Some(value) => {
102        let value_str = value.to_str().unwrap_or("");
103        if value_str.starts_with("Bearer ") {
104          value_str.trim_start_matches("Bearer ").to_string()
105        } else {
106          return Box::pin(async move { Err(actix_web::error::ErrorUnauthorized("Unauthorized")) });
107        }
108      }
109      None => {
110        return Box::pin(async move { Err(actix_web::error::ErrorUnauthorized("Unauthorized")) });
111      }
112    };
113
114    let header = match decode_header(&encoded_jwt) {
115      Ok(header) => header,
116      Err(e) => {
117        return Box::pin(async move {
118          Err(actix_web::error::ErrorUnauthorized(format!(
119            "Unauthorized: {e}"
120          )))
121        });
122      }
123    };
124
125    let kid = match header.kid {
126      Some(kid) => kid,
127      None => {
128        return Box::pin(async move {
129          Err(actix_web::error::ErrorUnauthorized(
130            "Unauthorized: No kid in JWT header".to_string(),
131          ))
132        });
133      }
134    };
135
136    let key_store = self.config.key_store.clone();
137    let service = self.service.clone();
138
139    Box::pin(async move {
140      let req = req;
141      let key = match key_store.get_key(&kid).await {
142        Ok(key) => key,
143        Err(e) => {
144          return Err(actix_web::error::ErrorUnauthorized(format!(
145            "Unauthorized: {e}"
146          )));
147        }
148      };
149
150      let jwt = match decode::<T>(&encoded_jwt, key) {
151        Ok(data) => data.claims,
152        Err(e) => {
153          return Err(actix_web::error::ErrorUnauthorized(format!(
154            "Unauthorized: {}",
155            e
156          )));
157        }
158      };
159
160      // Store the decoded JWT in the request's extensions.
161      req.extensions_mut().insert(jwt);
162
163      // Call the next service
164      let res = service.call(req).await?;
165      Ok(res)
166    })
167  }
168}