Add auth middleware, info endpoint to notary server (#368)

* Init auth middleware.

* Add auth module and info endpoint.

* Modify comment.

* Make failure to load auth list fallible.

* Uses hashmap for whitelist, remove expired session, and other fixes.

* Turn off auth.

* Fix argument type.

notary-server/Cargo.toml

@@ -10,6 +10,8 @@ axum = { version = "0.6.18", features = ["ws"] }
 axum-core = "0.3.4"
 axum-macros = "0.3.8"
 base64 = "0.21.0"
+chrono = "0.4.31"
+csv = "1.3.0"
 eyre = "0.6.8"
 futures = "0.3"
 futures-util = "0.3.28"

notary-server/config/config.yaml

@@ -12,6 +12,11 @@ tls-signature:
 
 notary-signature:
   private-key-pem-path: "./fixture/notary/notary.key"
+  public-key-pem-path: "./fixture/notary/notary.pub"
 
 tracing:
   default-level: DEBUG
+
+authorization:
+  enabled: false
+  whitelist-csv-path: "./fixture/auth/whitelist.csv"

notary-server/fixture/auth/whitelist.csv

@@ -0,0 +1,3 @@
+"Name","ApiKey","CreatedAt"
+"Jonas Nielsen","test_api_key_0","2023-09-18T07:38:53Z"
+"Eren Jaeger","test_api_key_1","2023-10-18T07:38:53Z"

notary-server/src/config.rs

@@ -13,6 +13,17 @@ pub struct NotaryServerProperties {
     pub notary_signature: NotarySignatureProperties,
     /// Setting for logging/tracing
     pub tracing: TracingProperties,
+    /// Setting for authorization
+    pub authorization: AuthorizationProperties,
+}
+
+#[derive(Clone, Debug, Deserialize)]
+#[serde(rename_all = "kebab-case")]
+pub struct AuthorizationProperties {
+    /// Switch to turn on or off auth middleware
+    pub enabled: bool,
+    /// File path of the whitelist API key csv
+    pub whitelist_csv_path: Option<String>,
 }
 
 #[derive(Clone, Debug, Deserialize)]
@@ -42,6 +53,7 @@ pub struct TLSSignatureProperties {
 #[serde(rename_all = "kebab-case")]
 pub struct NotarySignatureProperties {
     pub private_key_pem_path: String,
+    pub public_key_pem_path: String,
 }
 
 #[derive(Clone, Debug, Deserialize)]

notary-server/src/domain.rs

@@ -1,2 +1,15 @@
+pub mod auth;
 pub mod cli;
 pub mod notary;
+
+use serde::{Deserialize, Serialize};
+
+/// Response object of the /info API
+#[derive(Debug, Clone, Serialize, Deserialize)]
+#[serde(rename_all = "camelCase")]
+pub struct InfoResponse {
+    /// Current version of notary-server
+    pub version: String,
+    /// Public key of the notary signing key
+    pub public_key: String,
+}

notary-server/src/domain/auth.rs

@@ -0,0 +1,22 @@
+use serde::Deserialize;
+use std::collections::HashMap;
+
+/// Structure of each whitelisted record of the API key whitelist for authorization purpose
+#[derive(Clone, Debug, Deserialize)]
+#[serde(rename_all = "PascalCase")]
+pub struct AuthorizationWhitelistRecord {
+    pub name: String,
+    pub api_key: String,
+    pub created_at: String,
+}
+
+/// Convert whitelist data structure from vector to hashmap using api_key as the key to speed up lookup
+pub fn authorization_whitelist_vec_into_hashmap(
+    authorization_whitelist: Vec<AuthorizationWhitelistRecord>,
+) -> HashMap<String, AuthorizationWhitelistRecord> {
+    let mut hashmap = HashMap::new();
+    authorization_whitelist.iter().for_each(|record| {
+        hashmap.insert(record.api_key.clone(), record.to_owned());
+    });
+    hashmap
+}

notary-server/src/domain/notary.rs

@@ -1,10 +1,11 @@
 use std::{collections::HashMap, sync::Arc};
 
+use chrono::{DateTime, Utc};
 use p256::ecdsa::SigningKey;
 use serde::{Deserialize, Serialize};
 use tokio::sync::Mutex;
 
-use crate::config::NotarizationProperties;
+use crate::{config::NotarizationProperties, domain::auth::AuthorizationWhitelistRecord};
 
 /// Response object of the /session API
 #[derive(Debug, Clone, Serialize, Deserialize)]
@@ -40,24 +41,35 @@ pub enum ClientType {
     Websocket,
 }
 
+/// Session configuration data to be stored in temporary storage
+#[derive(Clone, Debug)]
+pub struct SessionData {
+    pub max_transcript_size: Option<usize>,
+    pub created_at: DateTime<Utc>,
+}
+
 /// Global data that needs to be shared with the axum handlers
 #[derive(Clone, Debug)]
 pub struct NotaryGlobals {
     pub notary_signing_key: SigningKey,
     pub notarization_config: NotarizationProperties,
     /// A temporary storage to store configuration data, mainly used for WebSocket client
-    pub store: Arc<Mutex<HashMap<String, Option<usize>>>>,
+    pub store: Arc<Mutex<HashMap<String, SessionData>>>,
+    /// Whitelist of API keys for authorization purpose
+    pub authorization_whitelist: Option<Arc<HashMap<String, AuthorizationWhitelistRecord>>>,
 }
 
 impl NotaryGlobals {
     pub fn new(
         notary_signing_key: SigningKey,
         notarization_config: NotarizationProperties,
+        authorization_whitelist: Option<Arc<HashMap<String, AuthorizationWhitelistRecord>>>,
     ) -> Self {
         Self {
             notary_signing_key,
             notarization_config,
             store: Default::default(),
+            authorization_whitelist,
         }
     }
 }

notary-server/src/error.rs

@@ -17,6 +17,8 @@ pub enum NotaryServerError {
     Notarization(Box<dyn Error + Send + 'static>),
     #[error("Invalid request from prover: {0}")]
     BadProverRequest(String),
+    #[error("Unauthorized request from prover: {0}")]
+    UnauthorizedProverRequest(String),
 }
 
 impl From<VerifierError> for NotaryServerError {
@@ -38,6 +40,11 @@ impl IntoResponse for NotaryServerError {
             bad_request_error @ NotaryServerError::BadProverRequest(_) => {
                 (StatusCode::BAD_REQUEST, bad_request_error.to_string()).into_response()
             }
+            unauthorized_request_error @ NotaryServerError::UnauthorizedProverRequest(_) => (
+                StatusCode::UNAUTHORIZED,
+                unauthorized_request_error.to_string(),
+            )
+                .into_response(),
             _ => (
                 StatusCode::INTERNAL_SERVER_ERROR,
                 "Something wrong happened.",

notary-server/src/lib.rs

@@ -1,14 +1,15 @@
 mod config;
 mod domain;
 mod error;
+mod middleware;
 mod server;
 mod server_tracing;
 mod service;
 mod util;
 
 pub use config::{
-    NotarizationProperties, NotaryServerProperties, NotarySignatureProperties, ServerProperties,
-    TLSSignatureProperties, TracingProperties,
+    AuthorizationProperties, NotarizationProperties, NotaryServerProperties,
+    NotarySignatureProperties, ServerProperties, TLSSignatureProperties, TracingProperties,
 };
 pub use domain::{
     cli::CliFields,

notary-server/src/middleware.rs

@@ -0,0 +1,104 @@
+use async_trait::async_trait;
+use axum::http::{header, request::Parts};
+use axum_core::extract::{FromRef, FromRequestParts};
+use std::collections::HashMap;
+use tracing::{error, trace};
+
+use crate::{
+    domain::{auth::AuthorizationWhitelistRecord, notary::NotaryGlobals},
+    NotaryServerError,
+};
+
+/// Auth middleware to prevent DOS
+pub struct AuthorizationMiddleware;
+
+#[async_trait]
+impl<S> FromRequestParts<S> for AuthorizationMiddleware
+where
+    NotaryGlobals: FromRef<S>,
+    S: Send + Sync,
+{
+    type Rejection = NotaryServerError;
+
+    async fn from_request_parts(parts: &mut Parts, state: &S) -> Result<Self, Self::Rejection> {
+        let notary_globals = NotaryGlobals::from_ref(state);
+        let Some(whitelist) = notary_globals.authorization_whitelist else {
+            trace!("Skipping authorization as whitelist is not set.");
+            return Ok(Self);
+        };
+        let auth_header = parts
+            .headers
+            .get(header::AUTHORIZATION)
+            .and_then(|value| std::str::from_utf8(value.as_bytes()).ok());
+
+        match auth_header {
+            Some(auth_header) => {
+                if api_key_is_valid(auth_header, &whitelist) {
+                    trace!("Request authorized.");
+                    Ok(Self)
+                } else {
+                    let err_msg = "Invalid API key.".to_string();
+                    error!(err_msg);
+                    Err(NotaryServerError::UnauthorizedProverRequest(err_msg))
+                }
+            }
+            None => {
+                let err_msg = "Missing API key.".to_string();
+                error!(err_msg);
+                Err(NotaryServerError::UnauthorizedProverRequest(err_msg))
+            }
+        }
+    }
+}
+
+/// Helper function to check if an API key is in whitelist
+fn api_key_is_valid(
+    api_key: &str,
+    whitelist: &HashMap<String, AuthorizationWhitelistRecord>,
+) -> bool {
+    whitelist.get(api_key).is_some()
+}
+
+#[cfg(test)]
+mod test {
+    use super::{api_key_is_valid, HashMap};
+    use crate::domain::auth::{
+        authorization_whitelist_vec_into_hashmap, AuthorizationWhitelistRecord,
+    };
+    use std::sync::Arc;
+
+    fn get_whitelist_fixture() -> HashMap<String, AuthorizationWhitelistRecord> {
+        authorization_whitelist_vec_into_hashmap(vec![
+            AuthorizationWhitelistRecord {
+                name: "test-name-0".to_string(),
+                api_key: "test-api-key-0".to_string(),
+                created_at: "2023-10-18T07:38:53Z".to_string(),
+            },
+            AuthorizationWhitelistRecord {
+                name: "test-name-1".to_string(),
+                api_key: "test-api-key-1".to_string(),
+                created_at: "2023-10-11T07:38:53Z".to_string(),
+            },
+            AuthorizationWhitelistRecord {
+                name: "test-name-2".to_string(),
+                api_key: "test-api-key-2".to_string(),
+                created_at: "2022-10-11T07:38:53Z".to_string(),
+            },
+        ])
+    }
+
+    #[test]
+    fn test_api_key_is_present() {
+        let whitelist = get_whitelist_fixture();
+        assert!(api_key_is_valid("test-api-key-0", &Arc::new(whitelist)));
+    }
+
+    #[test]
+    fn test_api_key_is_absent() {
+        let whitelist = get_whitelist_fixture();
+        assert_eq!(
+            api_key_is_valid("test-api-keY-0", &Arc::new(whitelist)),
+            false
+        );
+    }
+}

notary-server/src/server.rs

@@ -1,8 +1,9 @@
 use axum::{
     http::{Request, StatusCode},
+    middleware::from_extractor_with_state,
     response::IntoResponse,
     routing::{get, post},
-    Router,
+    Json, Router,
 };
 use eyre::{ensure, eyre, Result};
 use futures_util::future::poll_fn;
@@ -27,9 +28,15 @@ use tracing::{debug, error, info};
 
 use crate::{
     config::{NotaryServerProperties, NotarySignatureProperties, TLSSignatureProperties},
-    domain::notary::NotaryGlobals,
+    domain::{
+        auth::{authorization_whitelist_vec_into_hashmap, AuthorizationWhitelistRecord},
+        notary::NotaryGlobals,
+        InfoResponse,
+    },
     error::NotaryServerError,
+    middleware::AuthorizationMiddleware,
     service::{initialize, upgrade_protocol},
+    util::parse_csv_file,
 };
 
 /// Start a TLS-secured TCP server to accept notarization request for both TCP and WebSocket clients
@@ -39,6 +46,25 @@ pub async fn run_server(config: &NotaryServerProperties) -> Result<(), NotarySer
     let (tls_private_key, tls_certificates) = load_tls_key_and_cert(&config.tls_signature).await?;
     // Load the private key for notarized transcript signing from fixture folder — can be swapped out when we use proper ephemeral signing key
     let notary_signing_key = load_notary_signing_key(&config.notary_signature).await?;
+    // Load the authorization whitelist csv if it is turned on
+    let authorization_whitelist = if !config.authorization.enabled {
+        debug!("Skipping authorization as it is turned off.");
+        None
+    } else {
+        // Get the path of whitelist csv from config
+        let whitelist_csv_path = config
+            .authorization
+            .whitelist_csv_path
+            .as_ref()
+            .ok_or(eyre!(
+                "Failed to load authorization whitelist as its csv path is absent in config"
+            ))?;
+        // Load the csv
+        let whitelist_csv = parse_csv_file::<AuthorizationWhitelistRecord>(whitelist_csv_path)
+            .map_err(|err| eyre!("Failed to parse authorization whitelist csv: {:?}", err))?;
+        // Convert the whitelist record into hashmap for faster lookup
+        Some(authorization_whitelist_vec_into_hashmap(whitelist_csv))
+    };
 
     // Build a TCP listener with TLS enabled
     let mut server_config = ServerConfig::builder()
@@ -71,13 +97,45 @@ pub async fn run_server(config: &NotaryServerProperties) -> Result<(), NotarySer
     );
 
     let protocol = Arc::new(Http::new());
-    let notary_globals = NotaryGlobals::new(notary_signing_key, config.notarization.clone());
+    let notary_globals = NotaryGlobals::new(
+        notary_signing_key,
+        config.notarization.clone(),
+        authorization_whitelist.map(Arc::new),
+        // Use Arc to prevent cloning the whitelist for every request
+    );
+
+    // Parameters needed for the info endpoint
+    let public_key = std::fs::read_to_string(&config.notary_signature.public_key_pem_path)
+        .map_err(|err| eyre!("Failed to load notary public signing key for notarization: {err}"))?;
+    let version = env!("CARGO_PKG_VERSION").to_string();
     let router = Router::new()
         .route(
             "/healthcheck",
             get(|| async move { (StatusCode::OK, "Ok").into_response() }),
         )
+        .route(
+            "/info",
+            get(|| async move {
+                (
+                    StatusCode::OK,
+                    Json(InfoResponse {
+                        version,
+                        public_key,
+                    }),
+                )
+                    .into_response()
+            }),
+        )
         .route("/session", post(initialize))
+        // Not applying auth middleware to /notarize endpoint for now as we can rely on our
+        // short-lived session id generated from /session endpoint, as it is not possible
+        // to use header for API key for websocket /notarize endpoint due to browser restriction
+        // ref: https://stackoverflow.com/a/4361358; And putting it in url query param
+        // seems to be more insecured: https://stackoverflow.com/questions/5517281/place-api-key-in-headers-or-url
+        .route_layer(from_extractor_with_state::<
+            AuthorizationMiddleware,
+            NotaryGlobals,
+        >(notary_globals.clone()))
         .route("/notarize", get(upgrade_protocol))
         .with_state(notary_globals);
     let mut app = router.into_make_service();
@@ -187,6 +245,7 @@ mod test {
     async fn test_load_notary_signing_key() {
         let config = NotarySignatureProperties {
             private_key_pem_path: "./fixture/notary/notary.key".to_string(),
+            public_key_pem_path: "./fixture/notary/notary.pub".to_string(),
         };
         let result: Result<SigningKey> = load_notary_signing_key(&config).await;
         assert!(result.is_ok(), "Could not load notary private key");