nullnet_libappguard/
lib.rs

1mod proto;
2
3use proto::appguard::app_guard_client::AppGuardClient;
4pub use proto::*;
5use std::future::Future;
6use tokio::sync::mpsc;
7pub use tonic::Streaming;
8use tonic::codegen::tokio_stream::wrappers::ReceiverStream;
9use tonic::transport::{Channel, ClientTlsConfig};
10use tonic::{Request, Response, Status};
11use crate::appguard::{AppGuardHttpRequest, AppGuardHttpResponse, AppGuardResponse, AppGuardSmtpRequest, AppGuardSmtpResponse, AppGuardTcpConnection, AppGuardTcpInfo, AppGuardTcpResponse, Logs};
12use crate::appguard_commands::{ClientMessage, FirewallPolicy, ServerMessage};
13
14#[derive(Clone)]
15pub struct AppGuardGrpcInterface {
16    client: AppGuardClient<Channel>,
17}
18
19impl AppGuardGrpcInterface {
20    #[allow(clippy::missing_errors_doc)]
21    pub async fn new(host: &str, port: u16, tls: bool) -> Result<Self, String> {
22        let protocol = if tls { "https" } else { "http" };
23
24        let mut endpoint = Channel::from_shared(format!("{protocol}://{host}:{port}"))
25            .map_err(|e| e.to_string())?
26            .connect_timeout(std::time::Duration::from_secs(10));
27
28        if tls {
29            endpoint = endpoint
30                .tls_config(ClientTlsConfig::new().with_native_roots())
31                .map_err(|e| e.to_string())?;
32        }
33
34        let channel = endpoint.connect().await.map_err(|e| e.to_string())?;
35
36        Ok(Self {
37            client: AppGuardClient::new(channel),
38        })
39    }
40
41    #[allow(clippy::missing_errors_doc)]
42    pub async fn control_channel(
43        &self,
44        receiver: mpsc::Receiver<ClientMessage>,
45    ) -> Result<Streaming<ServerMessage>, String> {
46        let receiver = ReceiverStream::new(receiver);
47
48        Ok(self
49            .client
50            .clone()
51            .control_channel(Request::new(receiver))
52            .await
53            .map_err(|e| e.to_string())?
54            .into_inner())
55    }
56
57    #[allow(clippy::missing_errors_doc)]
58    pub async fn handle_logs(&mut self, message: Logs) -> Result<(), String> {
59        self.client
60            .handle_logs(Request::new(message))
61            .await
62            .map(|_| ())
63            .map_err(|e| e.to_string())
64    }
65
66    #[allow(clippy::missing_errors_doc)]
67    pub async fn handle_tcp_connection(
68        &mut self,
69        timeout: Option<u64>,
70        tcp_connection: AppGuardTcpConnection,
71    ) -> Result<AppGuardTcpResponse, Status> {
72        self.client
73            .handle_tcp_connection(Request::new(tcp_connection.clone()))
74            .wait_until_timeout(
75                timeout,
76                AppGuardTcpResponse {
77                    tcp_info: Some(AppGuardTcpInfo {
78                        connection: Some(tcp_connection),
79                        ..Default::default()
80                    }),
81                },
82            )
83            .await
84    }
85
86    #[allow(clippy::missing_errors_doc)]
87    pub async fn handle_http_request(
88        &mut self,
89        timeout: Option<u64>,
90        default_policy: FirewallPolicy,
91        http_request: AppGuardHttpRequest,
92    ) -> Result<AppGuardResponse, Status> {
93        self.client
94            .handle_http_request(Request::new(http_request))
95            .wait_until_timeout(
96                timeout,
97                AppGuardResponse {
98                    policy: default_policy as i32,
99                },
100            )
101            .await
102    }
103
104    #[allow(clippy::missing_errors_doc)]
105    pub async fn handle_http_response(
106        &mut self,
107        timeout: Option<u64>,
108        default_policy: FirewallPolicy,
109        http_response: AppGuardHttpResponse,
110    ) -> Result<AppGuardResponse, Status> {
111        self.client
112            .handle_http_response(Request::new(http_response))
113            .wait_until_timeout(
114                timeout,
115                AppGuardResponse {
116                    policy: default_policy as i32,
117                },
118            )
119            .await
120    }
121
122    #[allow(clippy::missing_errors_doc)]
123    pub async fn handle_smtp_request(
124        &mut self,
125        timeout: Option<u64>,
126        default_policy: FirewallPolicy,
127        smtp_request: AppGuardSmtpRequest,
128    ) -> Result<AppGuardResponse, Status> {
129        self.client
130            .handle_smtp_request(Request::new(smtp_request))
131            .wait_until_timeout(
132                timeout,
133                AppGuardResponse {
134                    policy: default_policy as i32,
135                },
136            )
137            .await
138    }
139
140    #[allow(clippy::missing_errors_doc)]
141    pub async fn handle_smtp_response(
142        &mut self,
143        timeout: Option<u64>,
144        default_policy: FirewallPolicy,
145        smtp_response: AppGuardSmtpResponse,
146    ) -> Result<AppGuardResponse, Status> {
147        self.client
148            .handle_smtp_response(Request::new(smtp_response))
149            .wait_until_timeout(
150                timeout,
151                AppGuardResponse {
152                    policy: default_policy as i32,
153                },
154            )
155            .await
156    }
157}
158
159trait WaitUntilTimeout<T> {
160    async fn wait_until_timeout(self, timeout: Option<u64>, default: T) -> Result<T, Status>;
161}
162
163impl<T, F: Future<Output = Result<Response<T>, Status>>> WaitUntilTimeout<T> for F {
164    async fn wait_until_timeout(self, timeout: Option<u64>, default: T) -> Result<T, Status> {
165        if let Some(t) = timeout {
166            if let Ok(res) = tokio::time::timeout(std::time::Duration::from_millis(t), self).await {
167                res.map(Response::into_inner)
168            } else {
169                // handler timed out, return default value
170                Ok(default)
171            }
172        } else {
173            self.await.map(Response::into_inner)
174        }
175    }
176}
nullnet_libappguard/lib.rs

nullnet_libappguard/
lib.rs
