Skip to main content

max / makenotwork

6.5 KB · 187 lines History Blame Raw
1 //! SSH banner check — TCP connect and verify SSH protocol banner.
2
3 use std::time::{Duration, Instant};
4
5 use tokio::io::AsyncReadExt;
6 use tokio::net::TcpStream;
7
8 use crate::config::SshBannerConfig;
9 use crate::types::{HealthSnapshot, HealthStatus};
10
11 /// Connect to an SSH server and verify the banner starts with "SSH-".
12 #[tracing::instrument(skip_all)]
13 pub async fn check_ssh_banner(target_name: &str, config: &SshBannerConfig) -> HealthSnapshot {
14 let addr = format!("{}:{}", config.host, config.port);
15 let timeout = Duration::from_secs(config.timeout_secs);
16 let checked_at = chrono::Utc::now().to_rfc3339();
17 let start = Instant::now();
18
19 let stream = match tokio::time::timeout(timeout, TcpStream::connect(&addr)).await {
20 Ok(Ok(stream)) => stream,
21 Ok(Err(e)) => {
22 return HealthSnapshot {
23 id: None,
24 target: target_name.to_string(),
25 status: HealthStatus::Unreachable,
26 checked_at,
27 response_time_ms: start.elapsed().as_millis() as i64,
28 details: None,
29 error: Some(format!("TCP connect to {addr} failed: {e}")),
30 };
31 }
32 Err(_) => {
33 return HealthSnapshot {
34 id: None,
35 target: target_name.to_string(),
36 status: HealthStatus::Unreachable,
37 checked_at,
38 response_time_ms: start.elapsed().as_millis() as i64,
39 details: None,
40 error: Some(format!("TCP connect to {addr} timed out after {timeout:?}")),
41 };
42 }
43 };
44
45 let mut buf = [0u8; 256];
46 let mut stream = stream;
47 match tokio::time::timeout(Duration::from_secs(5), stream.read(&mut buf)).await {
48 Ok(Ok(n)) if n > 0 => {
49 let banner = String::from_utf8_lossy(&buf[..n]);
50 let banner = banner.trim().to_string();
51 let response_time_ms = start.elapsed().as_millis() as i64;
52
53 if banner.starts_with("SSH-") {
54 HealthSnapshot {
55 id: None,
56 target: target_name.to_string(),
57 status: HealthStatus::Operational,
58 checked_at,
59 response_time_ms,
60 details: None,
61 error: None,
62 }
63 } else {
64 HealthSnapshot {
65 id: None,
66 target: target_name.to_string(),
67 status: HealthStatus::Degraded,
68 checked_at,
69 response_time_ms,
70 details: None,
71 error: Some(format!("unexpected banner: {banner}")),
72 }
73 }
74 }
75 Ok(Ok(_)) => HealthSnapshot {
76 id: None,
77 target: target_name.to_string(),
78 status: HealthStatus::Degraded,
79 checked_at,
80 response_time_ms: start.elapsed().as_millis() as i64,
81 details: None,
82 error: Some("empty response".to_string()),
83 },
84 Ok(Err(e)) => HealthSnapshot {
85 id: None,
86 target: target_name.to_string(),
87 status: HealthStatus::Degraded,
88 checked_at,
89 response_time_ms: start.elapsed().as_millis() as i64,
90 details: None,
91 error: Some(format!("banner read failed: {e}")),
92 },
93 Err(_) => HealthSnapshot {
94 id: None,
95 target: target_name.to_string(),
96 status: HealthStatus::Degraded,
97 checked_at,
98 response_time_ms: start.elapsed().as_millis() as i64,
99 details: None,
100 error: Some("banner read timed out".to_string()),
101 },
102 }
103 }
104
105 #[cfg(test)]
106 mod tests {
107 use super::*;
108 use tokio::io::AsyncWriteExt;
109 use tokio::net::TcpListener;
110
111 #[tokio::test]
112 async fn ssh_banner_unreachable() {
113 let config = SshBannerConfig {
114 host: "127.0.0.1".to_string(),
115 port: 19999,
116 timeout_secs: 1,
117 };
118 let result = check_ssh_banner("test", &config).await;
119 assert_eq!(result.status, HealthStatus::Unreachable);
120 assert!(result.error.unwrap().contains("TCP connect"));
121 }
122
123 #[tokio::test]
124 async fn ssh_banner_valid_operational() {
125 let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
126 let port = listener.local_addr().unwrap().port();
127
128 tokio::spawn(async move {
129 let (mut stream, _) = listener.accept().await.unwrap();
130 stream.write_all(b"SSH-2.0-OpenSSH_9.0\r\n").await.unwrap();
131 });
132
133 let config = SshBannerConfig {
134 host: "127.0.0.1".to_string(),
135 port,
136 timeout_secs: 5,
137 };
138 let result = check_ssh_banner("test", &config).await;
139 assert_eq!(result.status, HealthStatus::Operational);
140 assert!(result.error.is_none());
141 }
142
143 #[tokio::test]
144 async fn ssh_banner_unexpected_degraded() {
145 let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
146 let port = listener.local_addr().unwrap().port();
147
148 tokio::spawn(async move {
149 let (mut stream, _) = listener.accept().await.unwrap();
150 stream.write_all(b"HTTP/1.1 200 OK\r\n").await.unwrap();
151 });
152
153 let config = SshBannerConfig {
154 host: "127.0.0.1".to_string(),
155 port,
156 timeout_secs: 5,
157 };
158 let result = check_ssh_banner("test", &config).await;
159 assert_eq!(result.status, HealthStatus::Degraded);
160 assert!(result.error.unwrap().contains("unexpected banner"));
161 }
162
163 #[tokio::test]
164 async fn ssh_banner_empty_response() {
165 let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
166 let port = listener.local_addr().unwrap().port();
167
168 tokio::spawn(async move {
169 let (stream, _) = listener.accept().await.unwrap();
170 drop(stream); // close immediately
171 });
172
173 let config = SshBannerConfig {
174 host: "127.0.0.1".to_string(),
175 port,
176 timeout_secs: 5,
177 };
178 let result = check_ssh_banner("test", &config).await;
179 assert_eq!(result.status, HealthStatus::Degraded);
180 assert!(result.error.unwrap().contains("empty response"));
181 }
182
183 // Note: the read-error branch (Ok(Err(e))) is difficult to trigger with
184 // a real TCP listener without platform-specific tricks. The 4 tested branches
185 // cover all reachable paths in normal operation.
186 }
187