nullnet_libappguard/
lib.rs

1mod proto;
2
3use crate::proto::appguard::{
4    HeartbeatRequest
5};
6use proto::appguard::app_guard_client::AppGuardClient;
7pub use proto::appguard::{
8    AppGuardHttpRequest, AppGuardHttpResponse, AppGuardResponse, AppGuardSmtpRequest,
9    AppGuardSmtpResponse, AppGuardTcpConnection, AppGuardTcpInfo, AppGuardTcpResponse,
10    DeviceStatus, FirewallPolicy,HeartbeatResponse
11};
12use std::future::Future;
13use tonic::transport::{Channel, ClientTlsConfig};
14use tonic::{Request, Response, Status, Streaming};
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 heartbeat(
45        &mut self,
46        app_id: String,
47        app_secret: String,
48        device_version: String,
49        device_uuid: String,
50    ) -> Result<Streaming<HeartbeatResponse>, String> {
51        self.client
52            .heartbeat(Request::new(HeartbeatRequest {
53                app_id,
54                app_secret,
55                device_version,
56                device_uuid,
57            }))
58            .await
59            .map(tonic::Response::into_inner)
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}