use std::sync::Arc;
use std::sync::RwLock;
#[derive(Clone)]
pub(crate) struct AuthInterceptor {
token: Arc<RwLock<Option<String>>>,
}
impl AuthInterceptor {
pub(crate) fn new(token: Arc<RwLock<Option<String>>>) -> Self {
Self { token }
}
}
impl tonic::service::Interceptor for AuthInterceptor {
fn call(
&mut self,
mut request: tonic::Request<()>,
) -> std::result::Result<tonic::Request<()>, tonic::Status> {
let guard = self
.token
.read()
.map_err(|_| tonic::Status::internal("token lock poisoned"))?;
if let Some(ref token) = *guard {
request.metadata_mut().insert(
"authorization",
token
.parse()
.map_err(|_| tonic::Status::internal("invalid token"))?,
);
}
Ok(request)
}
}
#[cfg(test)]
mod tests {
use super::*;
use tonic::service::Interceptor;
#[test]
fn test_auth_interceptor_no_token() {
let token = Arc::new(RwLock::new(None));
let mut interceptor = AuthInterceptor::new(token);
let request = tonic::Request::new(());
let result = interceptor.call(request);
assert!(result.is_ok());
let req = result.unwrap();
assert!(req.metadata().get("authorization").is_none());
}
#[test]
fn test_auth_interceptor_with_token() {
let token = Arc::new(RwLock::new(Some("my-secret-token".to_string())));
let mut interceptor = AuthInterceptor::new(token);
let request = tonic::Request::new(());
let result = interceptor.call(request);
assert!(result.is_ok());
let req = result.unwrap();
let auth = req.metadata().get("authorization").unwrap();
assert_eq!(auth.to_str().unwrap(), "my-secret-token");
}
#[test]
fn test_auth_interceptor_clone() {
let token = Arc::new(RwLock::new(Some("token".to_string())));
let interceptor = AuthInterceptor::new(token);
let _cloned = interceptor.clone();
}
#[test]
fn test_auth_interceptor_token_update() {
let token = Arc::new(RwLock::new(Some("old-token".to_string())));
let mut interceptor = AuthInterceptor::new(token.clone());
let request = tonic::Request::new(());
let result = interceptor.call(request).unwrap();
assert_eq!(
result
.metadata()
.get("authorization")
.unwrap()
.to_str()
.unwrap(),
"old-token"
);
*token.write().unwrap() = Some("new-token".to_string());
let request = tonic::Request::new(());
let result = interceptor.call(request).unwrap();
assert_eq!(
result
.metadata()
.get("authorization")
.unwrap()
.to_str()
.unwrap(),
"new-token"
);
}
#[test]
fn test_auth_interceptor_token_removed() {
let token = Arc::new(RwLock::new(Some("token".to_string())));
let mut interceptor = AuthInterceptor::new(token.clone());
*token.write().unwrap() = None;
let request = tonic::Request::new(());
let result = interceptor.call(request).unwrap();
assert!(result.metadata().get("authorization").is_none());
}
}