• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1 // Copyright 2024 Google LLC
2 //
3 // Licensed under the Apache License, Version 2.0 (the "License");
4 // you may not use this file except in compliance with the License.
5 // You may obtain a copy of the License at
6 //
7 //     https://www.apache.org/licenses/LICENSE-2.0
8 //
9 // Unless required by applicable law or agreed to in writing, software
10 // distributed under the License is distributed on an "AS IS" BASIS,
11 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 // See the License for the specific language governing permissions and
13 // limitations under the License.
14 
15 use crate::util::{into_raw_descriptor, ProxyConfig};
16 use crate::{Connector, DnsManager, Error};
17 use bytes::Bytes;
18 use libslirp_rs::libslirp::{ProxyConnect, ProxyManager};
19 use log::{debug, warn};
20 use std::net::SocketAddr;
21 use std::sync::{mpsc, Arc};
22 use std::thread;
23 use tokio::runtime::Runtime;
24 
25 /// # Manager
26 ///
27 /// The `Manager` struct implements the `ProxyManager` trait from
28 /// `libslirp_rs`.  It is responsible for managing TCP connections
29 /// through an HTTP proxy using the `Connector` struct.
30 ///
31 /// The `Manager` uses a `tokio::runtime::Runtime` to spawn tasks for
32 /// establishing proxy connections.  It takes a proxy configuration
33 /// string as input, which is parsed into a `ProxyConfig` to create a
34 /// `Connector` instance.
35 ///
36 /// The `try_connect` method attempts to establish a connection to the
37 /// given `SocketAddr` through the proxy.  If successful, it calls the
38 /// `proxy_connect` function with the raw file descriptor of the
39 /// connected socket.
40 ///
41 /// # Example
42 ///
43 /// ```
44 /// use std::net::SocketAddr;
45 /// use libslirp_rs::libslirp::ProxyConnect;
46 ///
47 /// struct MyProxyConnect;
48 ///
49 /// impl ProxyConnect for MyProxyConnect {
50 ///     fn proxy_connect(&self, fd: i32, sockaddr: SocketAddr) {
51 ///         // Handle the connected socket
52 ///     }
53 /// }
54 ///
55 /// #[tokio::main]
56 /// async fn main() {
57 /// }
58 /// ```
59 pub struct Manager {
60     runtime: Arc<Runtime>,
61     connector: Connector,
62     dns_manager: Arc<DnsManager>,
63 }
64 
65 impl Manager {
66     /// Creates a new `LibSlirp` instance.
67     ///
68     /// This function initializes the libslirp library and spawns the necessary threads
69     /// for handling network traffic and polling.
new(proxy: &str, rx_proxy_bytes: mpsc::Receiver<Bytes>) -> Result<Self, Error>70     pub fn new(proxy: &str, rx_proxy_bytes: mpsc::Receiver<Bytes>) -> Result<Self, Error> {
71         let config = ProxyConfig::from_string(proxy)?;
72         let dns_manager = Arc::new(DnsManager::new());
73         let dns_manager_clone = dns_manager.clone();
74         let _ = thread::Builder::new().name("Dns Manager".to_string()).spawn(move || {
75             while let Ok(bytes) = rx_proxy_bytes.recv() {
76                 dns_manager_clone.add_from_ethernet_slice(&bytes);
77             }
78         });
79 
80         Ok(Self {
81             runtime: Arc::new(Runtime::new()?),
82             connector: Connector::new(config.addr, config.username, config.password),
83             dns_manager,
84         })
85     }
86 }
87 
88 impl ProxyManager for Manager {
89     /// Attempts to establish a TCP connection to the given `sockaddr` through the proxy.
90     ///
91     /// This function spawns a new task in the `tokio` runtime to handle the connection process.
92     /// If the connection is successful, it calls the `proxy_connect` function of the provided
93     /// `ProxyConnect` object with the raw file descriptor of the connected socket.
94     ///
95     /// # Arguments
96     ///
97     /// * `sockaddr` - The target socket address to connect to.
98     /// * `connect_id` - An identifier for the connection.
99     /// * `connect_func` - A `ProxyConnect` object that will be called with the connected socket.
100     ///
101     /// # Returns
102     ///
103     /// `true` if the connection attempt was initiated, `false` otherwise.
try_connect( &self, sockaddr: SocketAddr, connect_id: usize, connect_func: Box<dyn ProxyConnect + Send>, ) -> bool104     fn try_connect(
105         &self,
106         sockaddr: SocketAddr,
107         connect_id: usize,
108         connect_func: Box<dyn ProxyConnect + Send>,
109     ) -> bool {
110         debug!("Connecting to {sockaddr:?} with connect ID {connect_id}");
111         let connector = self.connector.clone();
112 
113         self.runtime.handle().spawn(async move {
114             let fd = match connector.connect(sockaddr).await {
115                 Ok(tcp_stream) => into_raw_descriptor(tcp_stream),
116                 Err(e) => {
117                     warn!("Failed to connect to proxy {}. {}", sockaddr, e);
118                     -1
119                 }
120             };
121             connect_func.proxy_connect(fd, sockaddr);
122         });
123 
124         true
125     }
126 
127     /// Removes a connection with the given `connect_id`.
128     ///
129     /// Currently, this function only logs a debug message.
130     ///
131     /// # Arguments
132     ///
133     /// * `connect_id` - The identifier of the connection to remove.
remove(&self, connect_id: usize)134     fn remove(&self, connect_id: usize) {
135         debug!("Remove connect ID {}", connect_id);
136     }
137 }
138