Skip to main content

pwr_server/
listener.rs

1//! TLS listener and connection accept loop for pwr-server.
2//!
3//! Uses synchronous I/O wrapped in tokio::task::spawn_blocking to
4//! keep the handler code simple (Read + Write traits) while still
5//! running in the tokio async runtime.
6
7use std::fs;
8use std::io::BufReader;
9use std::net::{TcpListener, SocketAddr};
10use std::sync::Arc;
11
12use rustls::ServerConfig as TlsServerConfig;
13use rustls::pki_types::CertificateDer;
14
15use crate::auth::RateLimiter;
16use crate::config::ServerConfig;
17use crate::handler::{self, HandlerContext};
18use crate::storage::ProjectStorage;
19
20/// Build a rustls TLS server configuration from the cert and key files.
21fn build_tls_config(config: &ServerConfig) -> Result<TlsServerConfig, String> {
22    let cert_file = fs::File::open(&config.tls_cert_path)
23        .map_err(|e| format!("Cannot open TLS cert {}: {}", config.tls_cert_path.display(), e))?;
24    let mut cert_reader = BufReader::new(cert_file);
25    let certs: Vec<CertificateDer> = rustls_pemfile::certs(&mut cert_reader)
26        .collect::<Result<Vec<_>, _>>()
27        .map_err(|e| format!("Cannot parse TLS cert: {}", e))?;
28
29    if certs.is_empty() {
30        return Err("No certificates found in cert file".into());
31    }
32
33    let key_file = fs::File::open(&config.tls_key_path)
34        .map_err(|e| format!("Cannot open TLS key {}: {}", config.tls_key_path.display(), e))?;
35    let mut key_reader = BufReader::new(key_file);
36    let key = rustls_pemfile::private_key(&mut key_reader)
37        .map_err(|e| format!("Cannot parse TLS key: {}", e))?
38        .ok_or_else(|| "No private key found in key file".to_string())?;
39
40    let tls_config = TlsServerConfig::builder()
41        .with_no_client_auth()
42        .with_single_cert(certs, key)
43        .map_err(|e| format!("Cannot build TLS config: {}", e))?;
44
45    Ok(tls_config)
46}
47
48/// Run the server main loop. Blocks until shutdown.
49///
50/// Spawns one OS thread per connection (via spawn_blocking) to keep
51/// handler code synchronous while the tokio runtime drives I/O.
52pub fn run(config: ServerConfig) -> Result<(), String> {
53    let psk = pwr_core::crypto::psk_from_hex(&config.auth_token)
54        .map_err(|e| format!("Invalid auth token: {}", e))?;
55
56    let storage = ProjectStorage::new(config.clone())?;
57    let storage = Arc::new(std::sync::RwLock::new(storage));
58    let rate_limiter = Arc::new(std::sync::Mutex::new(RateLimiter::new()));
59
60    let tls_config = Arc::new(build_tls_config(&config)?);
61
62    let bind_addr = config.bind_addr();
63    let listener = bind_dual_stack(&bind_addr)?;
64
65    log::info!("pwr-server listening on {} (TLS 1.3)", bind_addr);
66    log::info!("Storage: {}", config.storage_base_path.display());
67
68    // Accept connections in a loop
69    for stream_result in listener.incoming() {
70        let stream = stream_result.map_err(|e| format!("Accept error: {}", e))?;
71        let peer_addr = stream
72            .peer_addr()
73            .unwrap_or_else(|_| "unknown".parse().unwrap());
74
75        let psk = psk;
76        let storage = storage.clone();
77        let rate_limiter = rate_limiter.clone();
78        let tls_config = tls_config.clone();
79
80        // Handle each connection on a dedicated thread
81        std::thread::spawn(move || {
82            log::debug!("Connection from {}", peer_addr);
83
84            // Perform TLS handshake
85            let conn = match rustls::ServerConnection::new(tls_config.clone()) {
86                Ok(c) => c,
87                Err(e) => {
88                    log::error!("Cannot create TLS connection: {}", e);
89                    return;
90                }
91            };
92            let mut tls_stream = rustls::StreamOwned::new(conn, stream);
93
94            let ctx = HandlerContext {
95                storage: storage.clone(),
96                rate_limiter: rate_limiter.clone(),
97                psk,
98                peer_addr,
99                connected_at: std::time::Instant::now(),
100            };
101
102            if let Err(e) = handler::handle_connection(&mut tls_stream, ctx) {
103                log::error!("Handler error for {}: {}", peer_addr, e);
104            }
105        });
106    }
107
108    Ok(())
109}
110
111// ---------------------------------------------------------------------------
112// Dual-stack binding
113// ---------------------------------------------------------------------------
114
115/// Bind to `addr`, preferring IPv6 dual-stack with IPv4 fallback.
116///
117/// On Linux, binding to `[::]:port` with `IPV6_V6ONLY=0` (the default)
118/// accepts both IPv4 and IPv6 connections. If the configured address
119/// fails to bind, we try `0.0.0.0:port` as a fallback for systems
120/// without IPv6.
121///
122/// Also explicitly sets `SO_REUSEADDR` so rapid restarts don't fail
123/// with "address already in use".
124fn bind_dual_stack(bind_addr: &str) -> Result<TcpListener, String> {
125    // Try the configured address first
126    if let Ok(addr) = bind_addr.parse::<SocketAddr>() {
127        match bind_with_reuse(addr) {
128            Ok(listener) => return Ok(listener),
129            Err(e) => log::warn!("Cannot bind to {}: {} (trying fallback)", addr, e),
130        }
131    } else {
132        // Hostname bind — let the OS resolve it
133        match TcpListener::bind(bind_addr) {
134            Ok(listener) => return Ok(listener),
135            Err(e) => log::warn!("Cannot bind to {}: {} (trying fallback)", bind_addr, e),
136        }
137    }
138
139    // Fallback: if the configured address was [::], try 0.0.0.0
140    // (system may not have IPv6 available)
141    if bind_addr.starts_with('[') || bind_addr == "[::]" {
142        let port = bind_addr
143            .rsplit(':')
144            .next()
145            .and_then(|p| p.trim_end_matches(']').parse::<u16>().ok())
146            .unwrap_or(9742);
147        let fallback = SocketAddr::from(([0, 0, 0, 0], port));
148        log::info!("Falling back to IPv4: {}", fallback);
149        return bind_with_reuse(fallback)
150            .map_err(|e| format!("Cannot bind to {} or fallback {}: {}", bind_addr, fallback, e));
151    }
152
153    Err(format!("Cannot bind to {}", bind_addr))
154}
155
156/// Bind a TcpListener with SO_REUSEADDR set.
157fn bind_with_reuse(addr: SocketAddr) -> std::io::Result<TcpListener> {
158    use std::os::unix::io::FromRawFd;
159
160    let domain = if addr.is_ipv4() {
161        libc::AF_INET
162    } else {
163        libc::AF_INET6
164    };
165
166    let sock = unsafe {
167        let fd = libc::socket(domain, libc::SOCK_STREAM, 0);
168        if fd < 0 {
169            return Err(std::io::Error::last_os_error());
170        }
171
172        // SO_REUSEADDR: allow rebinding to the same port quickly after restart
173        let opt: libc::c_int = 1;
174        if libc::setsockopt(
175            fd,
176            libc::SOL_SOCKET,
177            libc::SO_REUSEADDR,
178            &opt as *const _ as *const _,
179            std::mem::size_of::<libc::c_int>() as u32,
180        ) < 0
181        {
182            libc::close(fd);
183            return Err(std::io::Error::last_os_error());
184        }
185
186        // For IPv6: clear IPV6_V6ONLY so the socket accepts IPv4 connections too
187        if domain == libc::AF_INET6 {
188            let opt: libc::c_int = 0;
189            if libc::setsockopt(
190                fd,
191                libc::IPPROTO_IPV6,
192                libc::IPV6_V6ONLY,
193                &opt as *const _ as *const _,
194                std::mem::size_of::<libc::c_int>() as u32,
195            ) < 0
196            {
197                // Non-fatal: IPv4-mapped addresses might not work but IPv6 still will
198                log::warn!("Cannot clear IPV6_V6ONLY: {}", std::io::Error::last_os_error());
199            }
200        }
201
202        fd
203    };
204
205    // Bind the socket
206    let sockaddr = socket_addr_to_raw(&addr);
207    let sockaddr_len = if addr.is_ipv4() {
208        std::mem::size_of::<libc::sockaddr_in>()
209    } else {
210        std::mem::size_of::<libc::sockaddr_in6>()
211    };
212
213    let ret = unsafe {
214        libc::bind(
215            sock,
216            &sockaddr as *const _ as *const libc::sockaddr,
217            sockaddr_len as u32,
218        )
219    };
220
221    if ret < 0 {
222        let err = std::io::Error::last_os_error();
223        unsafe { libc::close(sock) };
224        return Err(err);
225    }
226
227    // Listen
228    let ret = unsafe { libc::listen(sock, 128) };
229    if ret < 0 {
230        let err = std::io::Error::last_os_error();
231        unsafe { libc::close(sock) };
232        return Err(err);
233    }
234
235    // Wrap in TcpListener. SAFETY: we created this socket ourselves and
236    // it's in a valid listening state.
237    Ok(unsafe { TcpListener::from_raw_fd(sock) })
238}
239
240/// Convert a SocketAddr to a libc sockaddr_storage.
241fn socket_addr_to_raw(addr: &SocketAddr) -> libc::sockaddr_storage {
242    let mut storage: libc::sockaddr_storage = unsafe { std::mem::zeroed() };
243    match addr {
244        SocketAddr::V4(v4) => {
245            let raw: &mut libc::sockaddr_in =
246                unsafe { &mut *(&mut storage as *mut _ as *mut libc::sockaddr_in) };
247            raw.sin_family = libc::AF_INET as u16;
248            raw.sin_port = v4.port().to_be();
249            raw.sin_addr = libc::in_addr {
250                s_addr: u32::from_ne_bytes(v4.ip().octets()),
251            };
252        }
253        SocketAddr::V6(v6) => {
254            let raw: &mut libc::sockaddr_in6 =
255                unsafe { &mut *(&mut storage as *mut _ as *mut libc::sockaddr_in6) };
256            raw.sin6_family = libc::AF_INET6 as u16;
257            raw.sin6_port = v6.port().to_be();
258            raw.sin6_addr = libc::in6_addr {
259                s6_addr: v6.ip().octets(),
260            };
261            raw.sin6_flowinfo = v6.flowinfo();
262            raw.sin6_scope_id = v6.scope_id();
263        }
264    }
265    storage
266}