nullnet_libappguard/
lib.rs

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