From 5534a8b8770d663a46b20afddee914401ac19bad Mon Sep 17 00:00:00 2001
From: cpprian <cpprian456@gmail.com>
Date: Sun, 4 Aug 2024 18:20:51 +0200
Subject: [PATCH] impl health service for Auth

---
 Cargo.lock      | 14 ++++++++++++++
 Cargo.toml      |  1 +
 src/grpc/mod.rs |  7 +++++++
 3 files changed, 22 insertions(+)

diff --git a/Cargo.lock b/Cargo.lock
index 104e74e4ce..1b819bcbe7 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -1039,6 +1039,7 @@ dependencies = [
  "tokio-stream",
  "tonic",
  "tonic-build",
+ "tonic-health",
  "tower-http",
  "tracing",
  "tracing-subscriber",
@@ -4945,6 +4946,19 @@ dependencies = [
  "syn 2.0.68",
 ]
 
+[[package]]
+name = "tonic-health"
+version = "0.11.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "2cef6e24bc96871001a7e48e820ab240b3de2201e59b517cf52835df2f1d2350"
+dependencies = [
+ "async-stream",
+ "prost",
+ "tokio",
+ "tokio-stream",
+ "tonic",
+]
+
 [[package]]
 name = "tower"
 version = "0.4.13"
diff --git a/Cargo.toml b/Cargo.toml
index 13e39f3a11..02a9b22cc5 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -100,6 +100,7 @@ x25519-dalek = { version = "2.0", features = ["static_secrets"] }
 # openapi
 utoipa = { version = "4", features = ["axum_extras"] }
 utoipa-swagger-ui = { version = "7", features = ["axum"] }
+tonic-health = "0.11"
 
 [dev-dependencies]
 bytes = "1.6"
diff --git a/src/grpc/mod.rs b/src/grpc/mod.rs
index 784b40ce89..17e4e1b788 100644
--- a/src/grpc/mod.rs
+++ b/src/grpc/mod.rs
@@ -547,6 +547,12 @@ pub async fn run_grpc_server(
         GatewayServer::new(pool, gateway_state, wireguard_tx, mail_tx),
         JwtInterceptor::new(ClaimsType::Gateway),
     );
+
+    let (mut health_reporter, health_service) = tonic_health::server::health_reporter();
+    health_reporter
+        .set_serving::<AuthServiceServer<AuthServer>>()
+        .await;
+
     // Run gRPC server
     let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::UNSPECIFIED), server_config().grpc_port);
     debug!("Starting gRPC services");
@@ -559,6 +565,7 @@ pub async fn run_grpc_server(
     let router = builder
         .http2_keepalive_interval(Some(TEN_SECS))
         .tcp_keepalive(Some(TEN_SECS))
+        .add_service(health_service)
         .add_service(auth_service);
     #[cfg(feature = "wireguard")]
     let router = router.add_service(gateway_service);