Merge pull request #1080 from fabriziosestito/feat/readiness-probe-http

feat: expose readiness probe over http

Author: jvanz

cli-docs.md

@@ -67,6 +67,9 @@ This document contains the help content for the `policy-server` command-line pro
   Default value: `2`
 * `--port <PORT>` — Listen on PORT
 
+  Default value: `3000`
+* `--readiness-probe-port <READINESS_PROBE_PORT>` — Expose readiness endpoint on READINESS_PROBE_PORT
+
   Default value: `3000`
 * `--sigstore-cache-dir <SIGSTORE_CACHE_DIR>` — Directory used to cache sigstore data
 

src/cli.rs

@@ -68,6 +68,13 @@ pub(crate) fn build_cli() -> Command {
             .env("KUBEWARDEN_PORT")
             .help("Listen on PORT"),
 
+        Arg::new("readiness-probe-port")
+            .long("readiness-probe-port")
+            .value_name("READINESS_PROBE_PORT")
+            .default_value("3000")
+            .env("KUBEWARDEN_READINESS_PROBE_PORT")
+            .help("Expose readiness endpoint on READINESS_PROBE_PORT"),
+
         Arg::new("workers")
             .long("workers")
             .value_name("WORKERS_NUMBER")

src/config.rs

@@ -28,6 +28,7 @@ lazy_static! {
 
 pub struct Config {
     pub addr: SocketAddr,
+    pub readiness_probe_addr: SocketAddr,
     pub sources: Option<Sources>,
     pub policies: HashMap<String, PolicyOrPolicyGroup>,
     pub policies_download_dir: PathBuf,
@@ -60,6 +61,7 @@ impl Config {
     pub fn from_args(matches: &ArgMatches) -> Result<Self> {
         // init some variables based on the cli parameters
         let addr = api_bind_address(matches)?;
+        let readiness_probe_addr = readiness_probe_bind_address(matches)?;
 
         let policies = policies(matches)?;
         let policies_download_dir = matches
@@ -142,6 +144,7 @@ impl Config {
 
         Ok(Self {
             addr,
+            readiness_probe_addr,
             sources,
             policies,
             policies_download_dir,
@@ -176,6 +179,16 @@ fn api_bind_address(matches: &clap::ArgMatches) -> Result<SocketAddr> {
     .map_err(|e| anyhow!("error parsing arguments: {}", e))
 }
 
+fn readiness_probe_bind_address(matches: &clap::ArgMatches) -> Result<SocketAddr> {
+    format!(
+        "{}:{}",
+        matches.get_one::<String>("address").unwrap(),
+        matches.get_one::<String>("readiness-probe-port").unwrap()
+    )
+    .parse()
+    .map_err(|e| anyhow!("error parsing arguments: {}", e))
+}
+
 fn build_tls_config(matches: &clap::ArgMatches) -> Result<TlsConfig> {
     let cert_file = matches.get_one::<String>("cert-file").unwrap().to_owned();
     let key_file = matches.get_one::<String>("key-file").unwrap().to_owned();

src/lib.rs

@@ -34,7 +34,7 @@ use profiling::activate_memory_profiling;
 use rayon::prelude::*;
 use std::{fs, net::SocketAddr, sync::Arc};
 use tokio::{
-    sync::{oneshot, Semaphore},
+    sync::{oneshot, Notify, Semaphore},
     time,
 };
 use tower_http::trace::{self, TraceLayer};
@@ -65,10 +65,12 @@ pub static malloc_conf: &[u8] = b"background_thread:true,tcache_max:4096,dirty_d
 
 pub struct PolicyServer {
     router: Router,
+    readiness_probe_router: Router,
     callback_handler: CallbackHandler,
     callback_handler_shutdown_channel_tx: oneshot::Sender<()>,
     addr: SocketAddr,
     tls_config: Option<RustlsConfig>,
+    readiness_probe_addr: SocketAddr,
 }
 
 impl PolicyServer {
@@ -211,10 +213,7 @@ impl PolicyServer {
                 TraceLayer::new_for_http()
                     .make_span_with(trace::DefaultMakeSpan::new().level(Level::INFO))
                     .on_response(trace::DefaultOnResponse::new().level(Level::INFO)),
-            )
-            // Adding the readiness probe handler after the tracing layer to avoid logging
-            // See: https://github.com/tokio-rs/axum/discussions/355
-            .route("/readiness", get(readiness_handler));
+            );
 
         if config.enable_pprof {
             activate_memory_profiling().await?;
@@ -225,40 +224,56 @@ impl PolicyServer {
             router = Router::new().merge(router).merge(pprof_router);
         }
 
+        let readiness_probe_router = Router::new().route("/readiness", get(readiness_handler));
+
         Ok(Self {
             router,
+            readiness_probe_router,
             callback_handler,
             callback_handler_shutdown_channel_tx,
             addr: config.addr,
             tls_config,
+            readiness_probe_addr: config.readiness_probe_addr,
         })
     }
 
     pub async fn run(self) -> Result<()> {
-        // Start the CallbackHandler
+        let notify = Notify::new();
+
         let mut callback_handler = self.callback_handler;
         let callback_handler = tokio::spawn(async move {
             info!(status = "init", "CallbackHandler task");
             callback_handler.loop_eval().await;
             info!(status = "exit", "CallbackHandler task");
         });
 
-        if let Some(tls_config) = self.tls_config {
-            axum_server::bind_rustls(self.addr, tls_config)
-                .serve(self.router.into_make_service())
-                .await?;
-        } else {
-            axum_server::bind(self.addr)
-                .serve(self.router.into_make_service())
-                .await?;
+        let api_server = async {
+            if let Some(tls_config) = self.tls_config {
+                let server_with_tls = axum_server::bind_rustls(self.addr, tls_config);
+                notify.notify_one();
+
+                server_with_tls.serve(self.router.into_make_service()).await
+            } else {
+                let server = axum_server::bind(self.addr);
+                notify.notify_one();
+
+                server.serve(self.router.into_make_service()).await
+            }
+        };
+
+        let readiness_probe_server = async {
+            notify.notified().await;
+
+            axum_server::bind(self.readiness_probe_addr)
+                .serve(self.readiness_probe_router.into_make_service())
+                .await
         };
 
-        // Stop the CallbackHandler
+        tokio::try_join!(api_server, readiness_probe_server)?;
+
         self.callback_handler_shutdown_channel_tx
             .send(())
             .expect("Cannot send shutdown signal to CallbackHandler");
-
-        // Wait for the CallbackHandler to exit
         callback_handler
             .await
             .expect("Cannot wait for CallbackHandler to exit");

tests/common/mod.rs

@@ -106,6 +106,7 @@ pub(crate) fn default_test_config() -> Config {
 
     Config {
         addr: SocketAddr::from(([127, 0, 0, 1], 3001)),
+        readiness_probe_addr: SocketAddr::from(([127, 0, 0, 1], 3002)),
         sources: None,
         policies,
         policies_download_dir: tempdir().unwrap().into_path(),

tests/integration_test.rs

@@ -654,23 +654,11 @@ mod certificate_reload_helpers {
         }
     }
 
-    pub async fn policy_server_is_ready(
-        address: &str,
-        client_tls_pem_bundle: Option<String>,
-    ) -> anyhow::Result<StatusCode> {
+    pub async fn policy_server_is_ready(address: &str) -> anyhow::Result<StatusCode> {
         // wait for the server to start
-        let mut client_builder = reqwest::Client::builder();
-
-        if let Some(tls_data) = client_tls_pem_bundle {
-            let identity = reqwest::Identity::from_pem(tls_data.as_bytes())?;
-            client_builder = client_builder.identity(identity)
-        };
-        let client = client_builder
-            .danger_accept_invalid_certs(true)
-            .build()
-            .unwrap();
+        let client = reqwest::Client::builder().build().unwrap();
 
-        let url = reqwest::Url::parse(&format!("https://{address}/readiness")).unwrap();
+        let url = reqwest::Url::parse(&format!("http://{address}/readiness")).unwrap();
         let response = client.get(url).send().await?;
         Ok(response.status())
     }
@@ -704,8 +692,9 @@ async fn test_detect_certificate_rotation() {
     });
     config.policies = HashMap::new();
 
-    let domain_ip = config.addr.ip().to_string();
-    let domain_port = config.addr.port().to_string();
+    let host = config.addr.ip().to_string();
+    let port = config.addr.port().to_string();
+    let readiness_probe_port = config.readiness_probe_addr.port().to_string();
 
     tokio::spawn(async move {
         let api_server = policy_server::PolicyServer::new_from_config(config)
@@ -719,29 +708,22 @@ async fn test_detect_certificate_rotation() {
         .with_max_delay(Duration::from_secs(30))
         .with_max_times(5);
 
-    let client_cert = tls_data_client.cert.clone();
-    let client_key = tls_data_client.key.clone();
     let status_code = (|| async {
-        policy_server_is_ready(
-            format!("{domain_ip}:{domain_port}").as_str(),
-            Some(format!("{client_cert}\n{client_key}")),
-        )
-        .await
+        policy_server_is_ready(format!("{host}:{readiness_probe_port}").as_str()).await
     })
     .retry(exponential_backoff)
     .await
     .unwrap();
     assert_eq!(status_code, reqwest::StatusCode::OK);
 
-    check_tls_san_name(&domain_ip, &domain_port, hostname1)
+    check_tls_san_name(&host, &port, hostname1)
         .await
         .expect("certificate served doesn't use the expected SAN name");
 
     // Generate a new certificate and key, and switch to them
 
     let hostname2 = "cert2.example.com";
     let tls_data2 = create_cert(hostname2);
-    let client_ca2 = create_cert(hostname2);
 
     // write only the cert file
     std::fs::write(&cert_file, tls_data2.cert).unwrap();
@@ -750,7 +732,7 @@ async fn test_detect_certificate_rotation() {
     tokio::time::sleep(std::time::Duration::from_secs(4)).await;
 
     // the old certificate should still be in use, since we didn't change also the key
-    check_tls_san_name(&domain_ip, &domain_port, hostname1)
+    check_tls_san_name(&host, &port, hostname1)
         .await
         .expect("certificate should not have been changed");
 
@@ -760,32 +742,9 @@ async fn test_detect_certificate_rotation() {
     // give inotify some time to ensure it detected the cert change,
     // also give axum some time to complete the certificate reload
     tokio::time::sleep(std::time::Duration::from_secs(4)).await;
-    check_tls_san_name(&domain_ip, &domain_port, hostname2)
+    check_tls_san_name(&host, &port, hostname2)
         .await
         .expect("certificate hasn't been reloaded");
-
-    // Let test if the server is reloading client certificate
-    std::fs::write(&client_ca, client_ca2.cert.clone()).unwrap();
-
-    // give inotify some time to ensure it detected the cert change
-    tokio::time::sleep(std::time::Duration::from_secs(4)).await;
-
-    assert!(policy_server_is_ready(
-        format!("{domain_ip}:{domain_port}").as_str(),
-        Some(format!("{client_cert}\n{client_key}")),
-    )
-    .await
-    .is_err());
-
-    let client_cert = client_ca2.cert.clone();
-    let client_key = client_ca2.key.clone();
-    let status_code = policy_server_is_ready(
-        format!("{domain_ip}:{domain_port}").as_str(),
-        Some(format!("{client_cert}\n{client_key}")),
-    )
-    .await
-    .unwrap();
-    assert_eq!(status_code, reqwest::StatusCode::OK);
 }
 
 #[tokio::test]