// SPDX-License-Identifier: Apache-2.0 use super::test_utils::Key; use crate::common::test_utils::{generate_random_public_ip, get_ip_info}; use anyhow::Result; use detee_shared::app_proto::brain_app_daemon_client::BrainAppDaemonClient; use detee_shared::app_proto::{self, NewAppRes, RegisterAppNodeReq}; use detee_shared::common_proto::MappedPort; use futures::StreamExt; use tokio::sync::mpsc; use tokio_stream::wrappers::ReceiverStream; use tonic::transport::Channel; pub async fn mock_app_daemon( brain_channel: &Channel, daemon_error: Option, ) -> Result { let mut daemon_client = BrainAppDaemonClient::new(brain_channel.clone()); let daemon_key = Key::new(); register_app_node(&mut daemon_client, &daemon_key, &Key::new().pubkey).await?; let (tx, brain_msg_rx) = tokio::sync::mpsc::channel(1); tokio::spawn(daemon_listener(daemon_client.clone(), daemon_key.clone(), tx)); let (daemon_msg_tx, rx) = tokio::sync::mpsc::channel(1); tokio::spawn(daemon_msg_sender( daemon_client.clone(), daemon_key.clone(), daemon_msg_tx.clone(), rx, )); tokio::spawn(daemon_engine(daemon_msg_tx.clone(), brain_msg_rx, daemon_error)); tokio::time::sleep(tokio::time::Duration::from_secs(1)).await; Ok(daemon_key.pubkey) } pub async fn register_app_node( client: &mut BrainAppDaemonClient, key: &Key, operator_wallet: &str, ) -> Result> { log::info!("Registering app_node: {}", key.pubkey); let node_pubkey = key.pubkey.clone(); let ip = generate_random_public_ip().to_string(); let ip_info = get_ip_info(&ip).await?; let req = RegisterAppNodeReq { node_pubkey, operator_wallet: operator_wallet.to_string(), main_ip: ip_info.ip, city: ip_info.city, country: ip_info.country, region: ip_info.region, price: 1200, }; let mut grpc_stream = client.register_app_node(key.sign_request(req)?).await?.into_inner(); let mut deleted_app_reqs = Vec::new(); while let Some(stream_update) = grpc_stream.next().await { match stream_update { Ok(del_app_req) => { deleted_app_reqs.push(del_app_req); } Err(e) => { panic!("Received error instead of deleted_app_reqs: {e:?}"); } } } Ok(deleted_app_reqs) } pub async fn daemon_listener( mut client: BrainAppDaemonClient, key: Key, tx: mpsc::Sender, ) -> Result<()> { log::info!("listening app_daemon"); let mut grpc_stream = client.brain_messages(key.sign_stream_auth_app(vec![])?).await?.into_inner(); while let Some(Ok(stream_update)) = grpc_stream.next().await { log::info!("app deamon got notified: {:?}", &stream_update); let _ = tx.send(stream_update).await; } Ok(()) } pub async fn daemon_msg_sender( mut client: BrainAppDaemonClient, key: Key, tx: mpsc::Sender, rx: mpsc::Receiver, ) -> Result<()> { log::info!("sender app_daemon"); let rx_stream = ReceiverStream::new(rx); tx.send(app_proto::DaemonMessageApp { msg: Some(app_proto::daemon_message_app::Msg::Auth(key.sign_stream_auth_app(vec![])?)), }) .await?; client.daemon_messages(rx_stream).await?; Ok(()) } pub async fn daemon_engine( tx: mpsc::Sender, mut rx: mpsc::Receiver, new_app_err: Option, ) -> Result<()> { log::info!("daemon engine app_daemon"); while let Some(brain_msg) = rx.recv().await { match brain_msg.msg { Some(app_proto::brain_message_app::Msg::NewAppReq(new_app_req)) => { let exposed_ports = [vec![34500], new_app_req.resource.unwrap_or_default().ports].concat(); let mapped_ports = exposed_ports .into_iter() .map(|port| MappedPort { host_port: port, guest_port: port }) .collect::>(); let res_data = NewAppRes { app_id: new_app_req.app_id, mapped_ports, ip_address: "127.0.0.1".to_string(), error: new_app_err.clone().unwrap_or_default(), }; let res = app_proto::DaemonMessageApp { msg: Some(app_proto::daemon_message_app::Msg::NewAppRes(res_data)), }; tx.send(res).await?; } Some(app_proto::brain_message_app::Msg::DeleteAppReq(del_app_req)) => { println!("MOCK_APP_DAEMON::delete app request for {}", del_app_req.app_id); } None => todo!(), } } Ok(()) }