nullnet_libappguard/
lib.rs

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

nullnet_libappguard/
lib.rs
