Skip to content

Commit 33a4dd2

Browse files
zh-jqzh-jq
committed
vey-proxy: add socket type to EgressNotes
1 parent d4cabd6 commit 33a4dd2

24 files changed

Lines changed: 167 additions & 118 deletions

File tree

‎vey-proxy/src/escape/direct_fixed/tcp_connect/mod.rs‎

Lines changed: 11 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -23,7 +23,7 @@ use vey_types::net::{
2323
};
2424

2525
use super::DirectFixedEscaper;
26-
use crate::escape::EgressNotes;
26+
use crate::escape::{EgressNotes, EgressSocketType};
2727
use crate::log::escape::tcp_connect::EscapeLogForTcpConnect;
2828
use crate::module::tcp_connect::{
2929
TcpConnectRemoteWrapperStats, TcpConnectResult, TcpConnectTaskConf, UnderlyingTcpConnectError,
@@ -126,7 +126,7 @@ impl DirectFixedEscaper {
126126
task_notes: &ServerTaskNotes,
127127
) -> Result<TcpStream, UnderlyingTcpConnectError> {
128128
let peer = SocketAddr::new(peer_ip, task_conf.upstream.port());
129-
egress_notes.next = Some(peer);
129+
egress_notes.tcp.peer = Some(peer);
130130

131131
let (sock, bind) =
132132
self.prepare_connect_socket(peer_ip, egress_notes.bind, task_notes, &config)?;
@@ -145,7 +145,7 @@ impl DirectFixedEscaper {
145145
.local_addr()
146146
.map_err(UnderlyingTcpConnectError::SetupSocketFailed)?;
147147
self.stats.tcp.connect.add_established();
148-
egress_notes.local = Some(local_addr);
148+
egress_notes.tcp.local = Some(local_addr);
149149
egress_notes.final_addr.target_addr = Some(peer);
150150
egress_notes.final_addr.outgoing_addr = Some(local_addr);
151151
Ok(ups_stream)
@@ -224,7 +224,7 @@ impl DirectFixedEscaper {
224224
loop {
225225
if spawn_new_connection && let Some(ip) = ips.pop() {
226226
let peer = SocketAddr::new(ip, port);
227-
egress_notes.next = Some(peer);
227+
egress_notes.tcp.peer = Some(peer);
228228

229229
let (sock, bind) =
230230
self.prepare_connect_socket(ip, egress_notes.bind, task_notes, &config)?;
@@ -268,15 +268,15 @@ impl DirectFixedEscaper {
268268
Some(Ok(r)) => {
269269
running_connection -= 1;
270270
let peer_addr = r.1;
271-
egress_notes.next = Some(peer_addr);
271+
egress_notes.tcp.peer = Some(peer_addr);
272272
egress_notes.bind = r.2;
273273
match r.0 {
274274
Ok(ups_stream) => {
275275
let local_addr = ups_stream
276276
.local_addr()
277277
.map_err(UnderlyingTcpConnectError::SetupSocketFailed)?;
278278
self.stats.tcp.connect.add_established();
279-
egress_notes.local = Some(local_addr);
279+
egress_notes.tcp.local = Some(local_addr);
280280
egress_notes.final_addr.target_addr = Some(peer_addr);
281281
egress_notes.final_addr.outgoing_addr = Some(local_addr);
282282
return Ok(ups_stream);
@@ -354,6 +354,8 @@ impl DirectFixedEscaper {
354354
egress_notes: &mut EgressNotes,
355355
task_notes: &ServerTaskNotes,
356356
) -> Result<TcpStream, UnderlyingTcpConnectError> {
357+
egress_notes.socket_type = Some(EgressSocketType::Tcp);
358+
357359
let mut config = DirectTcpConnectConfig {
358360
connect: self.config.general.tcp_connect,
359361
keepalive: self.config.tcp_keepalive,
@@ -397,6 +399,7 @@ impl DirectFixedEscaper {
397399
old_egress_notes: &EgressNotes,
398400
task_notes: &ServerTaskNotes,
399401
) -> Result<TcpStream, UnderlyingTcpConnectError> {
402+
new_egress_notes.socket_type = Some(EgressSocketType::Tcp);
400403
new_egress_notes.bind = old_egress_notes.bind;
401404
#[cfg(target_os = "linux")]
402405
if let BindAddr::Foreign(addr) = new_egress_notes.bind {
@@ -422,7 +425,8 @@ impl DirectFixedEscaper {
422425
}
423426

424427
if task_conf.upstream.host_eq(old_upstream) {
425-
let control_addr = old_egress_notes.next.ok_or_else(|| {
428+
// This escaper only set tcp.peer for TCP connections
429+
let control_addr = old_egress_notes.tcp.peer.ok_or_else(|| {
426430
UnderlyingTcpConnectError::SetupSocketFailed(io::Error::new(
427431
io::ErrorKind::InvalidInput,
428432
"no peer address for referenced connection found",

‎vey-proxy/src/escape/direct_float/tcp_connect/mod.rs‎

Lines changed: 11 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -21,8 +21,8 @@ use vey_types::acl::AclAction;
2121
use vey_types::net::{ConnectError, Host, TcpKeepAliveConfig, UpstreamAddr};
2222

2323
use super::{DirectFloatBindIp, DirectFloatEscaper};
24-
use crate::escape::EgressNotes;
2524
use crate::escape::direct_fixed::tcp_connect::DirectTcpConnectConfig;
25+
use crate::escape::{EgressNotes, EgressSocketType};
2626
use crate::log::escape::tcp_connect::EscapeLogForTcpConnect;
2727
use crate::module::tcp_connect::{
2828
TcpConnectRemoteWrapperStats, TcpConnectResult, TcpConnectTaskConf, UnderlyingTcpConnectError,
@@ -110,7 +110,7 @@ impl DirectFloatEscaper {
110110
task_notes: &ServerTaskNotes,
111111
) -> Result<(TcpStream, DirectFloatBindIp), UnderlyingTcpConnectError> {
112112
let peer = SocketAddr::new(peer_ip, task_conf.upstream.port());
113-
egress_notes.next = Some(peer);
113+
egress_notes.tcp.peer = Some(peer);
114114

115115
let (sock, bind) =
116116
self.prepare_connect_socket(peer_ip, egress_notes.bind, task_notes, &config)?;
@@ -131,7 +131,7 @@ impl DirectFloatEscaper {
131131
.local_addr()
132132
.map_err(UnderlyingTcpConnectError::SetupSocketFailed)?;
133133
self.stats.tcp.connect.add_established();
134-
egress_notes.local = Some(local_addr);
134+
egress_notes.tcp.local = Some(local_addr);
135135
egress_notes.final_addr.target_addr = Some(peer);
136136
egress_notes.final_addr.outgoing_addr = Some(local_addr);
137137
Ok((ups_stream, bind))
@@ -209,7 +209,7 @@ impl DirectFloatEscaper {
209209
loop {
210210
if spawn_new_connection && let Some(ip) = ips.pop() {
211211
let peer = SocketAddr::new(ip, task_conf.upstream.port());
212-
egress_notes.next = Some(peer);
212+
egress_notes.tcp.peer = Some(peer);
213213
let (sock, bind) =
214214
self.prepare_connect_socket(ip, egress_notes.bind, task_notes, &config)?;
215215
running_connection += 1;
@@ -253,7 +253,7 @@ impl DirectFloatEscaper {
253253
running_connection -= 1;
254254
let peer_addr = r.1;
255255
let bind = r.2;
256-
egress_notes.next = Some(peer_addr);
256+
egress_notes.tcp.peer = Some(peer_addr);
257257
egress_notes.bind = BindAddr::Ip(bind.ip);
258258
egress_notes.expire = bind.expire_datetime;
259259
egress_notes.egress = Some(bind.egress_info.clone());
@@ -263,7 +263,7 @@ impl DirectFloatEscaper {
263263
.local_addr()
264264
.map_err(UnderlyingTcpConnectError::SetupSocketFailed)?;
265265
self.stats.tcp.connect.add_established();
266-
egress_notes.local = Some(local_addr);
266+
egress_notes.tcp.local = Some(local_addr);
267267
egress_notes.final_addr.target_addr = Some(peer_addr);
268268
egress_notes.final_addr.outgoing_addr = Some(local_addr);
269269
return Ok((ups_stream, bind));
@@ -341,6 +341,8 @@ impl DirectFloatEscaper {
341341
egress_notes: &mut EgressNotes,
342342
task_notes: &ServerTaskNotes,
343343
) -> Result<(TcpStream, DirectFloatBindIp), UnderlyingTcpConnectError> {
344+
egress_notes.socket_type = Some(EgressSocketType::Tcp);
345+
344346
let mut config = DirectTcpConnectConfig {
345347
connect: self.config.general.tcp_connect,
346348
keepalive: self.config.tcp_keepalive,
@@ -384,6 +386,7 @@ impl DirectFloatEscaper {
384386
old_egress_notes: &EgressNotes,
385387
task_notes: &ServerTaskNotes,
386388
) -> Result<(TcpStream, DirectFloatBindIp), UnderlyingTcpConnectError> {
389+
new_egress_notes.socket_type = Some(EgressSocketType::Tcp);
387390
new_egress_notes.bind = old_egress_notes.bind;
388391

389392
let mut config = DirectTcpConnectConfig {
@@ -404,7 +407,8 @@ impl DirectFloatEscaper {
404407
}
405408

406409
if task_conf.upstream.host_eq(old_upstream) {
407-
let control_addr = old_egress_notes.next.ok_or_else(|| {
410+
// This escaper only set tcp.peer for TCP connections
411+
let control_addr = old_egress_notes.tcp.peer.ok_or_else(|| {
408412
UnderlyingTcpConnectError::SetupSocketFailed(io::Error::new(
409413
io::ErrorKind::InvalidInput,
410414
"no peer address for referenced connection found",

‎vey-proxy/src/escape/divert_tcp/tcp_connect/mod.rs‎

Lines changed: 10 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,7 @@ use vey_socket::BindAddr;
1717
use vey_types::net::{ConnectError, Host};
1818

1919
use super::DivertTcpEscaper;
20-
use crate::escape::EgressNotes;
20+
use crate::escape::{EgressNotes, EgressSocketType};
2121
use crate::log::escape::tcp_connect::EscapeLogForTcpConnect;
2222
use crate::module::tcp_connect::{
2323
TcpConnectRemoteWrapperStats, TcpConnectResult, TcpConnectTaskConf, UnderlyingTcpConnectError,
@@ -84,8 +84,8 @@ impl DivertTcpEscaper {
8484
egress_notes: &mut EgressNotes,
8585
task_notes: &ServerTaskNotes,
8686
) -> Result<TcpStream, UnderlyingTcpConnectError> {
87+
egress_notes.tcp.peer = Some(peer);
8788
let (sock, bind) = self.prepare_connect_socket(peer.ip())?;
88-
egress_notes.next = Some(peer);
8989
egress_notes.bind = bind;
9090

9191
let instant_now = Instant::now();
@@ -106,7 +106,7 @@ impl DivertTcpEscaper {
106106
.local_addr()
107107
.map_err(UnderlyingTcpConnectError::SetupSocketFailed)?;
108108
self.stats.tcp.connect.add_established();
109-
egress_notes.local = Some(local_addr);
109+
egress_notes.tcp.local = Some(local_addr);
110110
// the chained outgoing addr is not detected at here
111111
Ok(ups_stream)
112112
}
@@ -182,8 +182,10 @@ impl DivertTcpEscaper {
182182

183183
loop {
184184
if spawn_new_connection && let Some(ip) = ips.pop() {
185-
let (sock, bind) = self.prepare_connect_socket(ip)?;
186185
let peer = SocketAddr::new(ip, peer_port);
186+
egress_notes.tcp.peer = Some(peer);
187+
188+
let (sock, bind) = self.prepare_connect_socket(ip)?;
187189
running_connection += 1;
188190
spawn_new_connection = false;
189191
egress_notes.tries += 1;
@@ -224,15 +226,15 @@ impl DivertTcpEscaper {
224226
Some(Ok(r)) => {
225227
running_connection -= 1;
226228
let peer_addr = r.1;
227-
egress_notes.next = Some(peer_addr);
229+
egress_notes.tcp.peer = Some(peer_addr);
228230
egress_notes.bind = r.2;
229231
match r.0 {
230232
Ok(ups_stream) => {
231233
let local_addr = ups_stream
232234
.local_addr()
233235
.map_err(UnderlyingTcpConnectError::SetupSocketFailed)?;
234236
self.stats.tcp.connect.add_established();
235-
egress_notes.local = Some(local_addr);
237+
egress_notes.tcp.local = Some(local_addr);
236238
// the chained outgoing addr is not detected at here
237239
return Ok(ups_stream);
238240
}
@@ -309,6 +311,8 @@ impl DivertTcpEscaper {
309311
egress_notes: &mut EgressNotes,
310312
task_notes: &ServerTaskNotes,
311313
) -> Result<TcpStream, UnderlyingTcpConnectError> {
314+
egress_notes.socket_type = Some(EgressSocketType::Tcp);
315+
312316
let peer_proxy = self.get_next_proxy(task_notes, task_conf.upstream.host());
313317

314318
match peer_proxy.host() {

‎vey-proxy/src/escape/egress_notes.rs‎

Lines changed: 30 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -13,44 +13,56 @@ use vey_types::metrics::NodeName;
1313
use vey_types::net::{EgressInfo, UpstreamAddr};
1414

1515
/// This contains the final chained info about the client request
16-
#[derive(Debug, Clone, Default)]
16+
#[derive(Debug, Clone, Default, Copy)]
1717
pub(crate) struct FinalAddressNotes {
1818
pub(crate) target_addr: Option<SocketAddr>,
1919
pub(crate) outgoing_addr: Option<SocketAddr>,
2020
}
2121

22-
impl FinalAddressNotes {
23-
fn reset(&mut self) {
24-
self.target_addr = None;
25-
self.outgoing_addr = None;
26-
}
22+
#[derive(Debug, Default, Clone, Copy)]
23+
pub(crate) struct ConnectNotes {
24+
pub(crate) peer: Option<SocketAddr>,
25+
pub(crate) local: Option<SocketAddr>,
26+
}
27+
28+
#[derive(Debug, Clone, Copy)]
29+
pub(crate) enum EgressSocketType {
30+
Tcp,
31+
Socks5,
2732
}
2833

2934
#[derive(Debug, Default, Clone)]
3035
pub(crate) struct EgressNotes {
3136
pub(crate) escaper: NodeName,
3237
pub(crate) bind: BindAddr,
33-
pub(crate) next: Option<SocketAddr>,
3438
pub(crate) tries: usize,
35-
pub(crate) local: Option<SocketAddr>,
3639
pub(crate) expire: Option<DateTime<Utc>>,
3740
pub(crate) egress: Option<EgressInfo>,
41+
pub(crate) socket_type: Option<EgressSocketType>,
42+
pub(crate) tcp: ConnectNotes,
3843
pub(crate) final_addr: FinalAddressNotes,
3944
pub(crate) duration: Duration,
4045
pub(crate) override_peer: Option<UpstreamAddr>,
4146
}
4247

4348
impl EgressNotes {
4449
pub(crate) fn reset(&mut self) {
45-
self.escaper.clear();
46-
self.bind = BindAddr::None;
47-
self.next = None;
48-
self.tries = 0;
49-
self.local = None;
50-
self.expire = None;
51-
self.egress = None;
52-
self.final_addr.reset();
53-
self.duration = Duration::ZERO;
54-
self.override_peer = None;
50+
*self = Default::default();
51+
}
52+
53+
pub(crate) fn tcp_connect_peer_addr(&self) -> Option<SocketAddr> {
54+
let socket_type = self.socket_type?;
55+
match socket_type {
56+
EgressSocketType::Tcp => self.tcp.peer,
57+
EgressSocketType::Socks5 => self.tcp.peer,
58+
}
59+
}
60+
61+
pub(crate) fn tcp_connect_local_addr(&self) -> Option<SocketAddr> {
62+
let socket_type = self.socket_type?;
63+
match socket_type {
64+
EgressSocketType::Tcp => self.tcp.local,
65+
EgressSocketType::Socks5 => self.tcp.local,
66+
}
5567
}
5668
}

‎vey-proxy/src/escape/mod.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -47,7 +47,7 @@ mod egress_path;
4747
pub(crate) use egress_path::EgressPathSelection;
4848

4949
mod egress_notes;
50-
pub(crate) use egress_notes::{EgressNotes, FinalAddressNotes};
50+
pub(crate) use egress_notes::{EgressNotes, EgressSocketType, FinalAddressNotes};
5151

5252
mod comply_audit;
5353
mod comply_context;

‎vey-proxy/src/escape/proxy_float/peer/socks5/socks5_connect/mod.rs‎

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,7 @@ use vey_socks::v5;
1818
use vey_types::net::{SocketBufferConfig, UpstreamAddr};
1919

2020
use super::{ProxyFloatEscaper, ProxyFloatSocks5Peer};
21-
use crate::escape::EgressNotes;
21+
use crate::escape::{EgressNotes, EgressSocketType};
2222
use crate::log::escape::tls_handshake::TlsApplication;
2323
use crate::module::tcp_connect::{
2424
TcpConnectError, TcpConnectRemoteWrapperStats, TcpConnectResult, TcpConnectTaskConf,
@@ -38,6 +38,7 @@ impl ProxyFloatSocks5Peer {
3838
let mut stream = escaper
3939
.tcp_new_connection(self, task_conf, egress_notes, task_notes)
4040
.await?;
41+
egress_notes.socket_type = Some(EgressSocketType::Socks5);
4142
let outgoing_addr = v5::client::socks5_connect_to(
4243
&mut stream,
4344
&self.shared_config.auth_info,
@@ -83,8 +84,8 @@ impl ProxyFloatSocks5Peer {
8384
let mut ctl_stream = escaper
8485
.tcp_new_connection(self, &tcp_task_conf, egress_notes, task_notes)
8586
.await?;
86-
let local_tcp_addr = egress_notes.local.unwrap();
87-
let peer_tcp_addr = egress_notes.next.unwrap();
87+
let local_tcp_addr = egress_notes.tcp.local.unwrap();
88+
let peer_tcp_addr = egress_notes.tcp.peer.unwrap();
8889

8990
// bind early and send listen_addr if configured ?
9091
let send_udp_ip = match local_tcp_addr.ip() {

‎vey-proxy/src/escape/proxy_float/peer/socks5s/socks5_connect/mod.rs‎

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,7 @@ use vey_socks::v5;
2020
use vey_types::net::{SocketBufferConfig, UpstreamAddr};
2121

2222
use super::{ProxyFloatEscaper, ProxyFloatSocks5sPeer};
23-
use crate::escape::EgressNotes;
23+
use crate::escape::{EgressNotes, EgressSocketType};
2424
use crate::log::escape::tls_handshake::TlsApplication;
2525
use crate::module::tcp_connect::{
2626
TcpConnectError, TcpConnectResult, TcpConnectTaskConf, TlsConnectTaskConf,
@@ -39,6 +39,7 @@ impl ProxyFloatSocks5sPeer {
3939
let mut stream = escaper
4040
.tls_handshake_with_peer(task_conf, egress_notes, task_notes, &self.tls_name, self)
4141
.await?;
42+
egress_notes.socket_type = Some(EgressSocketType::Socks5);
4243
let outgoing_addr = v5::client::socks5_connect_to(
4344
&mut stream,
4445
&self.shared_config.auth_info,
@@ -97,8 +98,8 @@ impl ProxyFloatSocks5sPeer {
9798
self,
9899
)
99100
.await?;
100-
let local_tcp_addr = egress_notes.local.unwrap();
101-
let peer_tcp_addr = egress_notes.next.unwrap();
101+
let local_tcp_addr = egress_notes.tcp.local.unwrap();
102+
let peer_tcp_addr = egress_notes.tcp.peer.unwrap();
102103

103104
// bind early and send listen_addr if configured ?
104105
let send_udp_ip = match local_tcp_addr.ip() {

0 commit comments

Comments
 (0)