nullnet_libappguard/
lib.rs

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