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