Skip to content

Commit 15a91ae

Browse files
EpicEricEpicEric
committed
Improve coverage with more tests and removing noise from logs
1 parent 23d3a88 commit 15a91ae

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

50 files changed

+842
-259
lines changed

‎.github/workflows/validate.yml

+1-1
Original file line numberDiff line numberDiff line change
@@ -36,7 +36,7 @@ jobs:
3636
- name: Run Clippy
3737
run: |
3838
rustup component add clippy
39-
cargo clippy -- -D warnings
39+
cargo clippy --all-targets -- -D warnings
4040
build-msrv:
4141
name: Build on MSRV
4242
runs-on: ubuntu-latest

‎Cargo.lock

+21-2
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎Cargo.toml

+1
Original file line numberDiff line numberDiff line change
@@ -62,6 +62,7 @@ axum = { version = "0.8.1", features = ["macros", "ws"] }
6262
futures-util = "0.3.31"
6363
insta = { version = "1.42.1", features = ["yaml"] }
6464
mockall = "0.13.1"
65+
nix = "0.29.0"
6566
regex = "1.11.1"
6667
tokio-tungstenite = "0.26.2"
6768
tower = "0.5.2"

‎justfile

+3
Original file line numberDiff line numberDiff line change
@@ -12,3 +12,6 @@ book:
1212

1313
cli:
1414
to-html --no-prompt "cargo run --quiet -- --help" > cli.html
15+
16+
clippy:
17+
cargo clippy --all-targets --fix --allow-dirty && cargo fmt --all

‎src/addressing.rs

+15-23
Original file line numberDiff line numberDiff line change
@@ -481,11 +481,7 @@ mod address_delegator_tests {
481481
.fingerprint(HashAlg::Sha256);
482482
let mut mock = MockResolver::new();
483483
mock.expect_has_txt_record_for_fingerprint()
484-
.with(
485-
eq("_some_prefix"),
486-
eq("some.address"),
487-
eq(fingerprint.clone()),
488-
)
484+
.with(eq("_some_prefix"), eq("some.address"), eq(fingerprint))
489485
.return_const(true);
490486
mock.expect_has_cname_record_for_domain()
491487
.once()
@@ -525,11 +521,7 @@ mod address_delegator_tests {
525521
.fingerprint(HashAlg::Sha256);
526522
let mut mock = MockResolver::new();
527523
mock.expect_has_txt_record_for_fingerprint()
528-
.with(
529-
eq("_some_prefix"),
530-
eq("some.address"),
531-
eq(fingerprint.clone()),
532-
)
524+
.with(eq("_some_prefix"), eq("some.address"), eq(fingerprint))
533525
.return_const(true);
534526
let delegator = AddressDelegator::new(AddressDelegatorData {
535527
resolver: mock,
@@ -599,7 +591,7 @@ mod address_delegator_tests {
599591
.fingerprint(HashAlg::Sha256);
600592
let mut mock = MockResolver::new();
601593
mock.expect_has_txt_record_for_fingerprint()
602-
.with(eq("_some_prefix"), eq("something"), eq(fingerprint.clone()))
594+
.with(eq("_some_prefix"), eq("something"), eq(fingerprint))
603595
.return_const(false);
604596
let delegator = AddressDelegator::new(AddressDelegatorData {
605597
resolver: mock,
@@ -1067,95 +1059,95 @@ mod address_delegator_tests {
10671059
.get_http_address(
10681060
"a1",
10691061
&None,
1070-
&Some(f1.clone()),
1062+
&Some(f1),
10711063
&"127.0.0.1:12301".parse::<SocketAddr>().unwrap(),
10721064
)
10731065
.await;
10741066
let address2_f1_a1_u0 = delegator
10751067
.get_http_address(
10761068
"a1",
10771069
&None,
1078-
&Some(f1.clone()),
1070+
&Some(f1),
10791071
&"127.0.0.1:12302".parse::<SocketAddr>().unwrap(),
10801072
)
10811073
.await;
10821074
let address3_f2_a1_u0 = delegator
10831075
.get_http_address(
10841076
"a1",
10851077
&None,
1086-
&Some(f2.clone()),
1078+
&Some(f2),
10871079
&"127.0.0.1:12303".parse::<SocketAddr>().unwrap(),
10881080
)
10891081
.await;
10901082
let address4_f1_a2_u0 = delegator
10911083
.get_http_address(
10921084
"a2",
10931085
&None,
1094-
&Some(f1.clone()),
1086+
&Some(f1),
10951087
&"127.0.0.1:12304".parse::<SocketAddr>().unwrap(),
10961088
)
10971089
.await;
10981090
let address5_f1_a1_u1 = delegator
10991091
.get_http_address(
11001092
"a1",
11011093
&Some("u1".into()),
1102-
&Some(f1.clone()),
1094+
&Some(f1),
11031095
&"127.0.0.1:12305".parse::<SocketAddr>().unwrap(),
11041096
)
11051097
.await;
11061098
let address6_f1_a1_u1 = delegator
11071099
.get_http_address(
11081100
"a1",
11091101
&Some("u1".into()),
1110-
&Some(f1.clone()),
1102+
&Some(f1),
11111103
&"127.0.0.1:12306".parse::<SocketAddr>().unwrap(),
11121104
)
11131105
.await;
11141106
let address7_f2_a1_u1 = delegator
11151107
.get_http_address(
11161108
"a1",
11171109
&Some("u1".into()),
1118-
&Some(f2.clone()),
1110+
&Some(f2),
11191111
&"127.0.0.1:12307".parse::<SocketAddr>().unwrap(),
11201112
)
11211113
.await;
11221114
let address8_f1_a2_u1 = delegator
11231115
.get_http_address(
11241116
"a2",
11251117
&Some("u1".into()),
1126-
&Some(f1.clone()),
1118+
&Some(f1),
11271119
&"127.0.0.1:12308".parse::<SocketAddr>().unwrap(),
11281120
)
11291121
.await;
11301122
let address9_f1_a1_u2 = delegator
11311123
.get_http_address(
11321124
"a1",
11331125
&Some("u2".into()),
1134-
&Some(f1.clone()),
1126+
&Some(f1),
11351127
&"127.0.0.1:12309".parse::<SocketAddr>().unwrap(),
11361128
)
11371129
.await;
11381130
let address10_f1_a1_u2 = delegator
11391131
.get_http_address(
11401132
"a1",
11411133
&Some("u2".into()),
1142-
&Some(f1.clone()),
1134+
&Some(f1),
11431135
&"127.0.0.1:12310".parse::<SocketAddr>().unwrap(),
11441136
)
11451137
.await;
11461138
let address11_f2_a1_u2 = delegator
11471139
.get_http_address(
11481140
"a1",
11491141
&Some("u2".into()),
1150-
&Some(f2.clone()),
1142+
&Some(f2),
11511143
&"127.0.0.1:12311".parse::<SocketAddr>().unwrap(),
11521144
)
11531145
.await;
11541146
let address12_f1_a2_u2 = delegator
11551147
.get_http_address(
11561148
"a2",
11571149
&Some("u2".into()),
1158-
&Some(f1.clone()),
1150+
&Some(f1),
11591151
&"127.0.0.1:12312".parse::<SocketAddr>().unwrap(),
11601152
)
11611153
.await;

‎src/admin.rs

+1-2
Original file line numberDiff line numberDiff line change
@@ -203,8 +203,7 @@ impl AdminState {
203203
// If this is not a pseudo-terminal, show some information on how to allocate one
204204
let text = Text::from(vec![
205205
Line::from(
206-
"PTY not detected! Make sure to connect with \"ssh -t ... admin\" instead."
207-
.red(),
206+
r#"PTY not detected! Make sure to connect with "ssh -t" instead."#.red(),
208207
),
209208
Line::from("Press Ctrl-C to close this connection."),
210209
]);

‎src/lib.rs

+27-11
Original file line numberDiff line numberDiff line change
@@ -48,7 +48,7 @@ use tokio::{
4848
io::AsyncWriteExt,
4949
net::{TcpListener, TcpStream},
5050
pin,
51-
time::sleep,
51+
time::{sleep, timeout},
5252
};
5353
use tokio_rustls::LazyConfigAcceptor;
5454
use tokio_util::sync::CancellationToken;
@@ -186,6 +186,8 @@ pub async fn entrypoint(config: ApplicationConfig) -> anyhow::Result<()> {
186186
)
187187
.into());
188188
}
189+
let http_request_timeout = config.http_request_timeout.map(Into::into);
190+
let tcp_connection_timeout = config.tcp_connection_timeout.map(Into::into);
189191
// Initialize crypto and credentials
190192
rustls::crypto::aws_lc_rs::default_provider()
191193
.install_default()
@@ -309,7 +311,7 @@ pub async fn entrypoint(config: ApplicationConfig) -> anyhow::Result<()> {
309311
Arc::clone(&tcp_connections),
310312
Arc::clone(&telemetry),
311313
Arc::clone(&ip_filter),
312-
config.tcp_connection_timeout.map(Into::into),
314+
tcp_connection_timeout,
313315
config.disable_tcp_logs,
314316
));
315317
// Add TCP handler service as a listener for TCP port updates.
@@ -492,8 +494,8 @@ pub async fn entrypoint(config: ApplicationConfig) -> anyhow::Result<()> {
492494
},
493495
// Always use aliasing channels instead of tunneling channels.
494496
proxy_type: ProxyType::Aliasing,
495-
http_request_timeout: config.http_request_timeout.map(Into::into),
496-
websocket_timeout: config.tcp_connection_timeout.map(Into::into),
497+
http_request_timeout,
498+
websocket_timeout: tcp_connection_timeout,
497499
disable_http_logs: config.disable_http_logs,
498500
_phantom_data: PhantomData,
499501
});
@@ -530,7 +532,7 @@ pub async fn entrypoint(config: ApplicationConfig) -> anyhow::Result<()> {
530532
.unproxied_connection_timeout
531533
.map(Into::into)
532534
.unwrap_or(config.idle_connection_timeout.into()),
533-
tcp_connection_timeout: config.tcp_connection_timeout.map(Into::into),
535+
tcp_connection_timeout,
534536
});
535537

536538
// HTTP handler
@@ -562,8 +564,8 @@ pub async fn entrypoint(config: ApplicationConfig) -> anyhow::Result<()> {
562564
},
563565
// Always use tunneling channels.
564566
proxy_type: ProxyType::Tunneling,
565-
http_request_timeout: config.http_request_timeout.map(Into::into),
566-
websocket_timeout: config.tcp_connection_timeout.map(Into::into),
567+
http_request_timeout,
568+
websocket_timeout: tcp_connection_timeout,
567569
disable_http_logs: config.disable_http_logs,
568570
_phantom_data: PhantomData,
569571
});
@@ -593,7 +595,14 @@ pub async fn entrypoint(config: ApplicationConfig) -> anyhow::Result<()> {
593595
tokio::spawn(async move {
594596
let server = http1::Builder::new();
595597
let conn = server.serve_connection(io, service).with_upgrades();
596-
let _ = conn.await;
598+
match tcp_connection_timeout {
599+
Some(duration) => {
600+
let _ = timeout(duration, conn).await;
601+
}
602+
None => {
603+
let _ = conn.await;
604+
}
605+
}
597606
});
598607
}
599608
})
@@ -626,8 +635,8 @@ pub async fn entrypoint(config: ApplicationConfig) -> anyhow::Result<()> {
626635
},
627636
// Always use tunneling channels.
628637
proxy_type: ProxyType::Tunneling,
629-
http_request_timeout: config.http_request_timeout.map(Into::into),
630-
websocket_timeout: config.tcp_connection_timeout.map(Into::into),
638+
http_request_timeout,
639+
websocket_timeout: tcp_connection_timeout,
631640
disable_http_logs: config.disable_http_logs,
632641
_phantom_data: PhantomData,
633642
});
@@ -694,7 +703,14 @@ pub async fn entrypoint(config: ApplicationConfig) -> anyhow::Result<()> {
694703
TokioIo::new(stream),
695704
service,
696705
);
697-
let _ = conn.await;
706+
match tcp_connection_timeout {
707+
Some(duration) => {
708+
let _ = timeout(duration, conn).await;
709+
}
710+
None => {
711+
let _ = conn.await;
712+
}
713+
}
698714
}
699715
Err(err) => {
700716
warn!(

‎src/login.rs

+3-5
Original file line numberDiff line numberDiff line change
@@ -232,8 +232,7 @@ mod api_login_tests {
232232
let listener = TcpListener::bind("127.0.0.1:28011").await.unwrap();
233233
let jh = tokio::spawn(async move { axum::serve(listener, app).await.unwrap() });
234234

235-
let api_login =
236-
ApiLogin::new("http://localhost:28011/authentication".into(), mock).unwrap();
235+
let api_login = ApiLogin::new("http://localhost:28011/authentication", mock).unwrap();
237236
assert!(
238237
api_login
239238
.authenticate(&AuthenticationRequest {
@@ -270,8 +269,7 @@ mod api_login_tests {
270269
let listener = TcpListener::bind("127.0.0.1:28012").await.unwrap();
271270
let jh = tokio::spawn(async move { axum::serve(listener, app).await.unwrap() });
272271

273-
let api_login =
274-
ApiLogin::new("http://localhost:28012/authentication".into(), mock).unwrap();
272+
let api_login = ApiLogin::new("http://localhost:28012/authentication", mock).unwrap();
275273
assert!(
276274
!api_login
277275
.authenticate(&AuthenticationRequest {
@@ -357,7 +355,7 @@ mod api_login_tests {
357355
});
358356

359357
let api_login =
360-
ApiLogin::new("https://localhost:28013/secure_authentication".into(), mock).unwrap();
358+
ApiLogin::new("https://localhost:28013/secure_authentication", mock).unwrap();
361359
assert!(
362360
api_login
363361
.authenticate(&AuthenticationRequest {

‎src/ssh.rs

+17-2
Original file line numberDiff line numberDiff line change
@@ -1692,6 +1692,7 @@ impl Handler for ServerHandler {
16921692
proxy_handler(req, peer, fingerprint, Arc::clone(&proxy_data))
16931693
});
16941694
let io = TokioIo::new(channel.into_stream());
1695+
let tcp_connection_timeout = self.server.tcp_connection_timeout;
16951696
match self.auth_data {
16961697
// Serve HTTP for unauthed user, then add disconnection timeout if this is the last proxy connection
16971698
AuthenticatedData::None { ref proxy_count } => {
@@ -1705,7 +1706,14 @@ impl Handler for ServerHandler {
17051706
tokio::spawn(async move {
17061707
let server = auto::Builder::new(TokioExecutor::new());
17071708
let conn = server.serve_connection_with_upgrades(io, service);
1708-
let _ = conn.await;
1709+
match tcp_connection_timeout {
1710+
Some(duration) => {
1711+
let _ = timeout(duration, conn).await;
1712+
}
1713+
None => {
1714+
let _ = conn.await;
1715+
}
1716+
}
17091717
if proxy_count.fetch_sub(1, Ordering::AcqRel) == 1 {
17101718
*timeout_handle.lock().await =
17111719
Some(DroppableHandle(tokio::spawn(async move {
@@ -1720,7 +1728,14 @@ impl Handler for ServerHandler {
17201728
tokio::spawn(async move {
17211729
let server = auto::Builder::new(TokioExecutor::new());
17221730
let conn = server.serve_connection_with_upgrades(io, service);
1723-
let _ = conn.await;
1731+
match tcp_connection_timeout {
1732+
Some(duration) => {
1733+
let _ = timeout(duration, conn).await;
1734+
}
1735+
None => {
1736+
let _ = conn.await;
1737+
}
1738+
}
17241739
});
17251740
}
17261741
}

0 commit comments

Comments
 (0)