From 7349bf16e1217a498fab151b5b638e08ee3424c5 Mon Sep 17 00:00:00 2001
From: ghe0 <gheorghe@gheo.tech>
Date: Tue, 3 Dec 2024 03:51:36 +0200
Subject: [PATCH] fix bug where containers are not connecting

---
 src/grpc/server.rs | 19 ++++++++++++++++---
 1 file changed, 16 insertions(+), 3 deletions(-)

diff --git a/src/grpc/server.rs b/src/grpc/server.rs
index 3625c8f..6aecc97 100644
--- a/src/grpc/server.rs
+++ b/src/grpc/server.rs
@@ -1,6 +1,7 @@
 use super::challenge::{update_server::UpdateServer, Empty, Keys, NodeUpdate};
 use crate::{datastore::State, grpc::challenge::update_server::Update};
 use detee_sgx::RaTlsConfig;
+use rustls::pki_types::CertificateDer;
 use std::{pin::Pin, sync::Arc};
 use tokio::sync::broadcast::Sender;
 use tokio_stream::{Stream, StreamExt};
@@ -50,7 +51,7 @@ impl MyServer {
         let tls_acceptor = TlsAcceptor::from(Arc::new(tls));
 
         loop {
-            let (conn, _addr) = match listener.accept().await {
+            let (conn, addr) = match listener.accept().await {
                 Ok(incoming) => incoming,
                 Err(e) => {
                     println!("Error accepting connection: {}", e);
@@ -85,7 +86,12 @@ impl MyServer {
                     conn.unwrap()
                 };
 
-                let svc = ServiceBuilder::new().service(svc);
+                let extension_layer =
+                    tower_http::add_extension::AddExtensionLayer::new(Arc::new(ConnInfo {
+                        addr,
+                        certificates,
+                    }));
+                let svc = ServiceBuilder::new().layer(extension_layer).service(svc);
 
                 if let Err(e) = http
                     .serve_connection(
@@ -103,6 +109,12 @@ impl MyServer {
     }
 }
 
+#[derive(Debug)]
+struct ConnInfo {
+    addr: std::net::SocketAddr,
+    certificates: Vec<CertificateDer<'static>>,
+}
+
 #[tonic::async_trait]
 impl Update for MyServer {
     type GetUpdatesStream = Pin<Box<dyn Stream<Item = Result<NodeUpdate, Status>> + Send>>;
@@ -115,8 +127,9 @@ impl Update for MyServer {
         &self,
         req: Request<Streaming<NodeUpdate>>,
     ) -> Result<Response<Self::GetUpdatesStream>, Status> {
+        let conn_info = req.extensions().get::<Arc<ConnInfo>>().unwrap();
         self.state.increase_mratls_conns();
-        let remote_ip = req.remote_addr().unwrap().ip().to_string();
+        let remote_ip = conn_info.addr.ip().to_string();
         let tx = self.tx.clone();
         let mut rx = self.tx.subscribe();
         let mut inbound = req.into_inner();