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,
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_tcp_connection(
66        &mut self,
67        timeout: Option<u64>,
68        tcp_connection: AppGuardTcpConnection,
69    ) -> Result<AppGuardTcpResponse, Status> {
70        self.client
71            .handle_tcp_connection(Request::new(tcp_connection.clone()))
72            .wait_until_timeout(
73                timeout,
74                AppGuardTcpResponse {
75                    tcp_info: Some(AppGuardTcpInfo {
76                        connection: Some(tcp_connection),
77                        ..Default::default()
78                    }),
79                },
80            )
81            .await
82    }
83
84    #[allow(clippy::missing_errors_doc)]
85    pub async fn handle_http_request(
86        &mut self,
87        timeout: Option<u64>,
88        default_policy: FirewallPolicy,
89        http_request: AppGuardHttpRequest,
90    ) -> Result<AppGuardResponse, Status> {
91        self.client
92            .handle_http_request(Request::new(http_request))
93            .wait_until_timeout(
94                timeout,
95                AppGuardResponse {
96                    policy: default_policy as i32,
97                },
98            )
99            .await
100    }
101
102    #[allow(clippy::missing_errors_doc)]
103    pub async fn handle_http_response(
104        &mut self,
105        timeout: Option<u64>,
106        default_policy: FirewallPolicy,
107        http_response: AppGuardHttpResponse,
108    ) -> Result<AppGuardResponse, Status> {
109        self.client
110            .handle_http_response(Request::new(http_response))
111            .wait_until_timeout(
112                timeout,
113                AppGuardResponse {
114                    policy: default_policy as i32,
115                },
116            )
117            .await
118    }
119
120    #[allow(clippy::missing_errors_doc)]
121    pub async fn handle_smtp_request(
122        &mut self,
123        timeout: Option<u64>,
124        default_policy: FirewallPolicy,
125        smtp_request: AppGuardSmtpRequest,
126    ) -> Result<AppGuardResponse, Status> {
127        self.client
128            .handle_smtp_request(Request::new(smtp_request))
129            .wait_until_timeout(
130                timeout,
131                AppGuardResponse {
132                    policy: default_policy as i32,
133                },
134            )
135            .await
136    }
137
138    #[allow(clippy::missing_errors_doc)]
139    pub async fn handle_smtp_response(
140        &mut self,
141        timeout: Option<u64>,
142        default_policy: FirewallPolicy,
143        smtp_response: AppGuardSmtpResponse,
144    ) -> Result<AppGuardResponse, Status> {
145        self.client
146            .handle_smtp_response(Request::new(smtp_response))
147            .wait_until_timeout(
148                timeout,
149                AppGuardResponse {
150                    policy: default_policy as i32,
151                },
152            )
153            .await
154    }
155}
156
157trait WaitUntilTimeout<T> {
158    async fn wait_until_timeout(self, timeout: Option<u64>, default: T) -> Result<T, Status>;
159}
160
161impl<T, F: Future<Output = Result<Response<T>, Status>>> WaitUntilTimeout<T> for F {
162    async fn wait_until_timeout(self, timeout: Option<u64>, default: T) -> Result<T, Status> {
163        if let Some(t) = timeout {
164            if let Ok(res) = tokio::time::timeout(std::time::Duration::from_millis(t), self).await {
165                res.map(Response::into_inner)
166            } else {
167                // handler timed out, return default value
168                Ok(default)
169            }
170        } else {
171            self.await.map(Response::into_inner)
172        }
173    }
174}