1use std::fs::File;
19use std::io::{self, BufReader, Seek, Write};
20use std::net::SocketAddr;
21use std::path::Path;
22use std::sync::Arc;
23
24use rustls::client::danger::{ServerCertVerified, ServerCertVerifier};
25use rustls::pki_types::pem::PemObject;
26use rustls::pki_types::{CertificateDer, PrivateKeyDer, PrivatePkcs8KeyDer, ServerName, UnixTime};
27use rustls::{ClientConfig, DigitallySignedStruct, RootCertStore, ServerConfig};
28use tokio::net::{TcpListener, TcpStream};
29use tokio_rustls::client::TlsStream as ClientTlsStream;
30use tokio_rustls::server::TlsStream as ServerTlsStream;
31use tokio_rustls::{TlsAcceptor, TlsConnector};
32use tokio_util::codec::Framed;
33use tracing::{debug, error, info, instrument, warn};
34
35use crate::core::codec::PacketCodec;
36use crate::core::packet::Packet;
37use crate::error::{ProtocolError, Result};
38use futures::{SinkExt, StreamExt};
39
40#[derive(Debug)]
42struct CertificateFingerprint {
43 fingerprint: Vec<u8>,
44}
45
46impl ServerCertVerifier for CertificateFingerprint {
47 fn verify_server_cert(
48 &self,
49 end_entity: &CertificateDer<'_>,
50 _intermediates: &[CertificateDer<'_>],
51 _server_name: &ServerName,
52 _ocsp_response: &[u8],
53 _now: UnixTime,
54 ) -> std::result::Result<ServerCertVerified, rustls::Error> {
55 use sha2::{Digest, Sha256};
56
57 let mut hasher = Sha256::new();
58 hasher.update(end_entity);
59 let hash = hasher.finalize();
60
61 if hash.as_slice() == self.fingerprint.as_slice() {
62 Ok(ServerCertVerified::assertion())
63 } else {
64 Err(rustls::Error::General(
65 "Pinned certificate hash mismatch".into(),
66 ))
67 }
68 }
69
70 fn verify_tls12_signature(
71 &self,
72 _message: &[u8],
73 _cert: &CertificateDer<'_>,
74 _dss: &DigitallySignedStruct,
75 ) -> std::result::Result<rustls::client::danger::HandshakeSignatureValid, rustls::Error> {
76 Ok(rustls::client::danger::HandshakeSignatureValid::assertion())
77 }
78
79 fn verify_tls13_signature(
80 &self,
81 _message: &[u8],
82 _cert: &CertificateDer<'_>,
83 _dss: &DigitallySignedStruct,
84 ) -> std::result::Result<rustls::client::danger::HandshakeSignatureValid, rustls::Error> {
85 Ok(rustls::client::danger::HandshakeSignatureValid::assertion())
86 }
87
88 fn supported_verify_schemes(&self) -> Vec<rustls::SignatureScheme> {
89 vec![
91 rustls::SignatureScheme::RSA_PKCS1_SHA256,
92 rustls::SignatureScheme::ECDSA_NISTP256_SHA256,
93 rustls::SignatureScheme::ED25519,
94 ]
95 }
96}
97
98#[derive(Debug)]
99struct AcceptAnyServerCert;
100
101impl ServerCertVerifier for AcceptAnyServerCert {
102 fn verify_server_cert(
103 &self,
104 _end_entity: &CertificateDer<'_>,
105 _intermediates: &[CertificateDer<'_>],
106 _server_name: &ServerName,
107 _ocsp_response: &[u8],
108 _now: UnixTime,
109 ) -> std::result::Result<ServerCertVerified, rustls::Error> {
110 Ok(ServerCertVerified::assertion())
111 }
112
113 fn verify_tls12_signature(
114 &self,
115 _message: &[u8],
116 _cert: &CertificateDer<'_>,
117 _dss: &DigitallySignedStruct,
118 ) -> std::result::Result<rustls::client::danger::HandshakeSignatureValid, rustls::Error> {
119 Ok(rustls::client::danger::HandshakeSignatureValid::assertion())
120 }
121
122 fn verify_tls13_signature(
123 &self,
124 _message: &[u8],
125 _cert: &CertificateDer<'_>,
126 _dss: &DigitallySignedStruct,
127 ) -> std::result::Result<rustls::client::danger::HandshakeSignatureValid, rustls::Error> {
128 Ok(rustls::client::danger::HandshakeSignatureValid::assertion())
129 }
130
131 fn supported_verify_schemes(&self) -> Vec<rustls::SignatureScheme> {
132 vec![
133 rustls::SignatureScheme::RSA_PKCS1_SHA256,
134 rustls::SignatureScheme::ECDSA_NISTP256_SHA256,
135 rustls::SignatureScheme::ED25519,
136 ]
137 }
138}
139
140fn load_private_key(reader: &mut BufReader<File>) -> Result<PrivateKeyDer<'static>> {
142 reader
145 .seek(std::io::SeekFrom::Start(0))
146 .map_err(ProtocolError::Io)?;
147
148 let keys: std::result::Result<Vec<_>, _> =
150 PrivatePkcs8KeyDer::pem_reader_iter(reader).collect();
151 let keys =
152 keys.map_err(|_| ProtocolError::TlsError("Failed to parse PKCS8 private key".into()))?;
153
154 if !keys.is_empty() {
155 return Ok(PrivateKeyDer::Pkcs8(keys[0].clone_key()));
156 }
157
158 Err(ProtocolError::TlsError(
161 "No supported private key format found".into(),
162 ))
163}
164
165pub enum TlsVersion {
167 TLS12,
169 TLS13,
171 All,
173}
174
175pub struct TlsServerConfig {
177 cert_path: String,
178 key_path: String,
179 client_ca_path: Option<String>,
181 require_client_auth: bool,
183 tls_versions: Option<Vec<TlsVersion>>,
185 cipher_suites: Option<Vec<rustls::SupportedCipherSuite>>,
187 alpn_protocols: Option<Vec<Vec<u8>>>,
189}
190
191impl TlsServerConfig {
192 pub fn new<P: AsRef<std::path::Path>>(cert_path: P, key_path: P) -> Self {
194 Self {
195 cert_path: cert_path.as_ref().to_string_lossy().to_string(),
196 key_path: key_path.as_ref().to_string_lossy().to_string(),
197 client_ca_path: None,
198 require_client_auth: false,
199 tls_versions: None,
200 cipher_suites: None,
201 alpn_protocols: Some(vec![b"h2".to_vec(), b"http/1.1".to_vec()]),
202 }
203 }
204
205 pub fn with_tls_versions(mut self, versions: Vec<TlsVersion>) -> Self {
207 self.tls_versions = Some(versions);
208 self
209 }
210
211 pub fn with_cipher_suites(mut self, cipher_suites: Vec<rustls::SupportedCipherSuite>) -> Self {
213 self.cipher_suites = Some(cipher_suites);
214 self
215 }
216
217 pub fn with_client_auth<S: Into<String>>(mut self, client_ca_path: S) -> Self {
219 self.client_ca_path = Some(client_ca_path.into());
220 self.require_client_auth = true;
221 self
222 }
223
224 pub fn require_client_auth(mut self, required: bool) -> Self {
226 self.require_client_auth = required;
227 self
228 }
229
230 pub fn with_alpn_protocols(mut self, protocols: Vec<Vec<u8>>) -> Self {
232 self.alpn_protocols = Some(protocols);
233 self
234 }
235
236 pub fn generate_self_signed<P: AsRef<Path>>(cert_path: P, key_path: P) -> io::Result<Self> {
238 let cert = rcgen::generate_simple_self_signed(vec!["localhost".into()])
239 .map_err(|e| io::Error::other(format!("Certificate generation error: {e}")))?;
240
241 let mut cert_file = File::create(&cert_path)?;
243 let pem = cert.cert.pem();
244 cert_file.write_all(pem.as_bytes())?;
245
246 let mut key_file = File::create(&key_path)?;
248 key_file.write_all(cert.signing_key.serialize_pem().as_bytes())?;
249
250 Ok(Self {
251 cert_path: cert_path.as_ref().to_string_lossy().to_string(),
252 key_path: key_path.as_ref().to_string_lossy().to_string(),
253 client_ca_path: None,
254 require_client_auth: false,
255 tls_versions: None,
256 cipher_suites: None,
257 alpn_protocols: Some(vec![b"h2".to_vec(), b"http/1.1".to_vec()]),
258 })
259 }
260
261 pub fn load_server_config(&self) -> Result<ServerConfig> {
263 let cert_file = File::open(&self.cert_path)
265 .map_err(|e| ProtocolError::TlsError(format!("Failed to open cert file: {e}")))?;
266 let mut cert_reader = BufReader::new(cert_file);
267 let cert_chain: std::result::Result<Vec<_>, _> =
268 CertificateDer::pem_reader_iter(&mut cert_reader).collect();
269 let cert_chain: Vec<CertificateDer<'static>> = cert_chain
270 .map_err(|_| ProtocolError::TlsError("Failed to parse certificate".into()))?;
271
272 if cert_chain.is_empty() {
273 return Err(ProtocolError::TlsError("No certificates found".into()));
274 }
275
276 let key_file = File::open(&self.key_path)
278 .map_err(|e| ProtocolError::TlsError(format!("Failed to open key file: {e}")))?;
279 let mut key_reader = BufReader::new(key_file);
280 let private_key = load_private_key(&mut key_reader)?;
281
282 if let Some(versions) = &self.tls_versions {
285 let mut has_tls13 = false;
286 let mut has_tls12 = false;
287 for v in versions {
288 match v {
289 TlsVersion::TLS12 => has_tls12 = true,
290 TlsVersion::TLS13 => has_tls13 = true,
291 TlsVersion::All => {
292 has_tls13 = true;
293 has_tls12 = true;
294 }
295 }
296 }
297 debug!(
299 "TLS versions requested: TLS1.2={}, TLS1.3={}",
300 has_tls12, has_tls13
301 );
302 }
303
304 let config_builder = ServerConfig::builder_with_provider(std::sync::Arc::new(
306 rustls::crypto::ring::default_provider(),
307 ))
308 .with_safe_default_protocol_versions()
309 .map_err(|_| ProtocolError::TlsError("Failed to configure TLS protocol versions".into()))?;
310
311 let cert_builder = config_builder.with_no_client_auth();
312
313 let mut config = cert_builder
315 .with_single_cert(cert_chain.clone(), private_key.clone_key())
316 .map_err(|e| ProtocolError::TlsError(format!("TLS error: {e}")))?;
317
318 if let Some(client_ca_path) = &self.client_ca_path {
320 let client_ca_file = File::open(client_ca_path).map_err(|e| {
322 ProtocolError::TlsError(format!("Failed to open client CA file: {e}"))
323 })?;
324 let mut client_ca_reader = BufReader::new(client_ca_file);
325 let client_ca_certs: std::result::Result<Vec<_>, _> =
326 CertificateDer::pem_reader_iter(&mut client_ca_reader).collect();
327 let client_ca_certs: Vec<CertificateDer<'static>> = client_ca_certs.map_err(|_| {
328 ProtocolError::TlsError("Failed to parse client CA certificate".into())
329 })?;
330
331 if client_ca_certs.is_empty() {
332 return Err(ProtocolError::TlsError(
333 "No client CA certificates found".into(),
334 ));
335 }
336
337 let mut client_root_store = RootCertStore::empty();
339 for cert in client_ca_certs {
340 client_root_store.add(cert).map_err(|e| {
341 ProtocolError::TlsError(format!("Failed to add client CA cert: {e}"))
342 })?;
343 }
344
345 let client_auth = rustls::server::WebPkiClientVerifier::builder(std::sync::Arc::new(
347 client_root_store,
348 ))
349 .build()
350 .map_err(|e| {
351 ProtocolError::TlsError(format!("Failed to build client verifier: {e}"))
352 })?;
353
354 let new_builder = ServerConfig::builder_with_provider(std::sync::Arc::new(
356 rustls::crypto::ring::default_provider(),
357 ))
358 .with_safe_default_protocol_versions()
359 .map_err(|_| {
360 ProtocolError::TlsError("Failed to configure TLS protocol versions".into())
361 })?;
362 let new_cert_builder = new_builder.with_client_cert_verifier(client_auth);
363
364 config = new_cert_builder
366 .with_single_cert(cert_chain, private_key.clone_key())
367 .map_err(|e| ProtocolError::TlsError(format!("TLS error with client auth: {e}")))?;
368
369 debug!("mTLS enabled with client certificate verification required");
370 }
371
372 if let Some(protocols) = &self.alpn_protocols {
374 config.alpn_protocols = protocols.clone();
375 debug!(
376 protocol_count = protocols.len(),
377 "ALPN protocols configured"
378 );
379 }
380
381 Ok(config)
382 }
383
384 pub fn calculate_cert_hash(cert: &CertificateDer<'_>) -> Vec<u8> {
386 use sha2::{Digest, Sha256};
387 let mut hasher = Sha256::new();
388 hasher.update(cert.as_ref());
389 hasher.finalize().to_vec()
390 }
391}
392
393pub struct TlsClientConfig {
395 server_name: String,
396 insecure: bool,
397 pinned_cert_hash: Option<Vec<u8>>,
399 client_cert_path: Option<String>,
401 client_key_path: Option<String>,
403 tls_versions: Option<Vec<TlsVersion>>,
405 cipher_suites: Option<Vec<rustls::SupportedCipherSuite>>,
407}
408
409impl TlsClientConfig {
410 pub fn new<S: Into<String>>(server_name: S) -> Self {
412 Self {
413 server_name: server_name.into(),
414 insecure: false,
415 pinned_cert_hash: None,
416 client_cert_path: None,
417 client_key_path: None,
418 tls_versions: None,
419 cipher_suites: None,
420 }
421 }
422
423 pub fn with_tls_versions(mut self, versions: Vec<TlsVersion>) -> Self {
425 self.tls_versions = Some(versions);
426 self
427 }
428
429 pub fn with_cipher_suites(mut self, cipher_suites: Vec<rustls::SupportedCipherSuite>) -> Self {
431 self.cipher_suites = Some(cipher_suites);
432 self
433 }
434
435 pub fn with_client_certificate<S: Into<String>>(mut self, cert_path: S, key_path: S) -> Self {
437 self.client_cert_path = Some(cert_path.into());
438 self.client_key_path = Some(key_path.into());
439 self
440 }
441
442 pub fn insecure(mut self) -> Self {
455 warn!("INSECURE MODE ENABLED: Certificate verification is disabled. This should only be used for development/testing.");
456 self.insecure = true;
457 self
458 }
459
460 pub fn with_pinned_cert_hash(mut self, hash: Vec<u8>) -> Self {
467 if hash.len() != 32 {
468 warn!(
469 "Certificate hash has unexpected length: {} (expected 32 bytes for SHA-256)",
470 hash.len()
471 );
472 }
473 self.pinned_cert_hash = Some(hash);
474 self
475 }
476
477 pub fn load_client_config(&self) -> Result<ClientConfig> {
479 self.log_tls_version_info();
480
481 if self.insecure {
482 self.build_insecure_client_config()
483 } else {
484 self.build_secure_client_config()
485 }
486 }
487
488 fn log_tls_version_info(&self) {
490 if let Some(versions) = &self.tls_versions {
491 let mut has_tls13 = false;
492 let mut has_tls12 = false;
493 for v in versions {
494 match v {
495 TlsVersion::TLS12 => has_tls12 = true,
496 TlsVersion::TLS13 => has_tls13 = true,
497 TlsVersion::All => {
498 has_tls13 = true;
499 has_tls12 = true;
500 }
501 }
502 }
503 debug!(
504 "TLS client versions requested: TLS1.2={}, TLS1.3={}",
505 has_tls12, has_tls13
506 );
507 }
508 }
509
510 fn build_secure_client_config(&self) -> Result<ClientConfig> {
512 let root_store = self.load_system_root_certificates()?;
513 let builder = ClientConfig::builder_with_provider(std::sync::Arc::new(
514 rustls::crypto::ring::default_provider(),
515 ))
516 .with_safe_default_protocol_versions()
517 .map_err(|_| ProtocolError::TlsError("Failed to configure TLS protocol versions".into()))?
518 .with_root_certificates(root_store);
519
520 if let (Some(client_cert_path), Some(client_key_path)) =
522 (&self.client_cert_path, &self.client_key_path)
523 {
524 let (cert_chain, key) =
525 self.load_client_credentials(client_cert_path, client_key_path)?;
526 builder.with_client_auth_cert(cert_chain, key).map_err(|e| {
527 ProtocolError::TlsError(format!("Failed to set client certificate: {e}"))
528 })
529 } else {
530 Ok(builder.with_no_client_auth())
531 }
532 }
533
534 fn build_insecure_client_config(&self) -> Result<ClientConfig> {
536 let builder = ClientConfig::builder_with_provider(std::sync::Arc::new(
537 rustls::crypto::ring::default_provider(),
538 ))
539 .with_safe_default_protocol_versions()
540 .map_err(|_| ProtocolError::TlsError("Failed to configure TLS protocol versions".into()))?;
541 let verifier = self.create_custom_verifier();
542 let custom_builder = builder
543 .dangerous()
544 .with_custom_certificate_verifier(verifier);
545
546 if let (Some(client_cert_path), Some(client_key_path)) =
548 (&self.client_cert_path, &self.client_key_path)
549 {
550 let (cert_chain, key) =
551 self.load_client_credentials(client_cert_path, client_key_path)?;
552 custom_builder
553 .with_client_auth_cert(cert_chain, key)
554 .map_err(|e| {
555 ProtocolError::TlsError(format!("Failed to set client certificate: {e}"))
556 })
557 } else {
558 Ok(custom_builder.with_no_client_auth())
559 }
560 }
561
562 fn load_system_root_certificates(&self) -> Result<RootCertStore> {
564 let mut root_store = RootCertStore::empty();
565 let native_certs = rustls_native_certs::load_native_certs();
566 if let Some(e) = native_certs.errors.first() {
567 return Err(ProtocolError::TlsError(format!(
568 "Failed to load native certs: {e}"
569 )));
570 }
571
572 for cert in native_certs.certs {
573 root_store.add(cert).map_err(|e| {
574 ProtocolError::TlsError(format!("Failed to add cert to root store: {e}"))
575 })?;
576 }
577
578 Ok(root_store)
579 }
580
581 fn create_custom_verifier(&self) -> Arc<dyn ServerCertVerifier> {
583 if let Some(hash) = &self.pinned_cert_hash {
584 Arc::new(CertificateFingerprint {
585 fingerprint: hash.clone(),
586 })
587 } else {
588 Arc::new(AcceptAnyServerCert)
589 }
590 }
591
592 fn load_client_credentials(
594 &self,
595 cert_path: &str,
596 key_path: &str,
597 ) -> Result<(Vec<CertificateDer<'static>>, PrivateKeyDer<'static>)> {
598 let cert_file = File::open(cert_path).map_err(ProtocolError::Io)?;
600 let mut cert_reader = BufReader::new(cert_file);
601 let certs_result: std::result::Result<Vec<_>, _> =
602 CertificateDer::pem_reader_iter(&mut cert_reader).collect();
603 let certs: Vec<CertificateDer<'static>> = certs_result
604 .map_err(|_| ProtocolError::TlsError("Failed to parse client certificate".into()))?;
605
606 if certs.is_empty() {
607 return Err(ProtocolError::TlsError(
608 "No client certificates found".into(),
609 ));
610 }
611
612 let key_file = File::open(key_path).map_err(ProtocolError::Io)?;
614 let mut key_reader = BufReader::new(key_file);
615 let key = load_private_key(&mut key_reader)?;
616
617 Ok((certs, key))
618 }
619
620 pub fn server_name(&self) -> Result<ServerName<'_>> {
622 ServerName::try_from(self.server_name.as_str())
623 .map_err(|_| ProtocolError::TlsError("Invalid server name".into()))
624 }
625
626 pub fn server_name_string(&self) -> String {
628 self.server_name.clone()
629 }
630}
631
632#[instrument(skip(config))]
634pub async fn start_server(addr: &str, config: TlsServerConfig) -> Result<()> {
635 let tls_config = config.load_server_config()?;
636 let acceptor = TlsAcceptor::from(Arc::new(tls_config));
637 let listener = TcpListener::bind(addr).await?;
638
639 info!(address=%addr, "TLS server listening");
640
641 loop {
642 let (stream, peer) = listener.accept().await?;
643 let acceptor = acceptor.clone();
644
645 tokio::spawn(async move {
646 match acceptor.accept(stream).await {
647 Ok(tls_stream) => {
648 if let Err(e) = handle_tls_connection(tls_stream, peer).await {
649 error!(%peer, error=%e, "Connection error");
650 }
651 }
652 Err(e) => {
653 error!(%peer, error=%e, "TLS handshake failed");
654 }
655 }
656 });
657 }
658}
659
660#[instrument(skip(tls_stream), fields(peer=%peer))]
662async fn handle_tls_connection(
663 tls_stream: ServerTlsStream<TcpStream>,
664 peer: SocketAddr,
665) -> Result<()> {
666 let mut framed = Framed::new(tls_stream, PacketCodec);
667
668 info!("TLS connection established");
669
670 while let Some(packet) = framed.next().await {
671 match packet {
672 Ok(pkt) => {
673 debug!(bytes = pkt.payload.len(), "Received data");
674 on_packet(pkt, &mut framed).await?;
675 }
676 Err(e) => {
677 error!(error=%e, "Protocol error");
678 break;
679 }
680 }
681 }
682
683 info!("TLS connection closed");
684 Ok(())
685}
686
687#[instrument(skip(framed), fields(packet_version=pkt.version, payload_size=pkt.payload.len()))]
689async fn on_packet<T>(pkt: Packet, framed: &mut Framed<T, PacketCodec>) -> Result<()>
690where
691 T: tokio::io::AsyncRead + tokio::io::AsyncWrite + Unpin,
692{
693 let response = Packet {
695 version: pkt.version,
696 payload: pkt.payload,
697 };
698
699 framed.send(response).await?;
700 Ok(())
701}
702
703pub async fn connect(
705 addr: &str,
706 config: TlsClientConfig,
707) -> Result<Framed<ClientTlsStream<TcpStream>, PacketCodec>> {
708 let tls_config = Arc::new(config.load_client_config()?);
709 let connector = TlsConnector::from(tls_config);
710
711 let stream = TcpStream::connect(addr).await?;
712
713 let server_name_str = config.server_name_string();
716 let domain_static: &'static str = Box::leak(server_name_str.into_boxed_str());
717 let domain = ServerName::try_from(domain_static)
718 .map_err(|_| ProtocolError::TlsError("Invalid server name".into()))?;
719
720 let tls_stream = connector
721 .connect(domain, stream)
722 .await
723 .map_err(|e| ProtocolError::TlsError(format!("TLS connection failed: {e}")))?;
724
725 let framed = Framed::new(tls_stream, PacketCodec);
726 Ok(framed)
727}