fix bug where containers are not connecting

This commit is contained in:
ghe0 2024-12-03 03:51:36 +02:00
parent 9e9b906598
commit 7349bf16e1
Signed by: ghe0
GPG Key ID: 451028EE56A0FBB4

@ -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();