Coverage / crates/spectra-mcp/src/tls.rs 80.4%

pub fn provider() -> Arc<rustls::crypto::CryptoProvider> {

    Arc::new(rustls::crypto::ring::default_provider())

}

pub fn spki_sha256(cert_der: &[u8]) -> Result<[u8; 32]> {

    let (_, cert) = x509_parser::parse_x509_certificate(cert_der).map_err(|e| anyhow::anyhow!("not an X.509 certificate: {e}"))?;

    let d = ring::digest::digest(&ring::digest::SHA256, cert.public_key().raw);

    Ok(d.as_ref().try_into().expect("SHA-256 is 32 bytes"))

}

    pub fn load(data: &Path, source: &CertSource) -> Result<ServerTls> {

        let (cert_pem, key_pem, names) = match source {

            CertSource::Own { names } => {

                let keys = data.join("keys");

                let (crt, key) = (keys.join("tls.crt"), keys.join("tls.key"));

                if !crt.exists() || !key.exists() {

                    make_own(&crt, &key, names)?;

                }

                (std::fs::read_to_string(&crt)?, std::fs::read_to_string(&key)?, names.clone())

        let chain: Vec<CertificateDer<'static>> = rustls_pemfile::certs(&mut cert_pem.as_bytes()).collect::<Result<_, _>>()?;

        let leaf = chain.first().context("no certificate in the PEM")?;

        let spki = spki_sha256(leaf)?;

        let key: PrivateKeyDer<'static> = rustls_pemfile::private_key(&mut key_pem.as_bytes())?.context("no private key in the PEM")?;

        let mut config = rustls::ServerConfig::builder_with_provider(provider())

            .with_safe_default_protocol_versions()?

            .with_no_client_auth()

            .with_single_cert(chain, key)?;

        config.alpn_protocols = vec![b"http/1.1".to_vec()];

        Ok(ServerTls { config: Arc::new(config), spki_sha256: spki, names })

    }

fn make_own(crt: &Path, key: &Path, names: &[String]) -> Result<()> {

    if names.is_empty() {

    }

    let ck = rcgen::generate_simple_self_signed(names.to_vec())?;

    if let Some(dir) = crt.parent() {

        std::fs::create_dir_all(dir)?;

    write_private(key, ck.signing_key.serialize_pem().as_bytes())?;

    std::fs::write(crt, ck.cert.pem())?;

    Ok(())

}

fn write_private(path: &Path, bytes: &[u8]) -> Result<()> {

        let mut f = std::fs::OpenOptions::new().write(true).create(true).truncate(true).mode(0o600).open(path)?;

        f.write_all(bytes)?;

        Ok(())

}

    pub fn new(tcp: TcpListener, config: Arc<rustls::ServerConfig>) -> Result<TlsListener> {

        let addr = tcp.local_addr()?;

        let acceptor = tokio_rustls::TlsAcceptor::from(config);

        let (tx, rx) = tokio::sync::mpsc::channel(64);

        tokio::spawn(async move {

                let (stream, peer) = match tcp.accept().await {

                    Ok(x) => x,

                let (acceptor, out) = (acceptor.clone(), tx.clone());

                tokio::spawn(async move {

                    match tokio::time::timeout(Duration::from_secs(10), acceptor.accept(stream)).await {

                        Ok(Ok(tls)) => {

                            let _ = out.send((tls, peer)).await;

                        Ok(Err(e)) => tracing::debug!(%peer, "TLS handshake failed: {e}"),

                });

                if tx.is_closed() {

                }

        Ok(TlsListener { rx, addr })

    }

    async fn accept(&mut self) -> (Self::Io, Self::Addr) {

        match self.rx.recv().await {

            Some(x) => x,

    }

    pub fn new(pin: Option<[u8; 32]>) -> Arc<Pinned> {

        Arc::new(Pinned { pin, seen: Mutex::new(None), algs: provider().signature_verification_algorithms })

    }

    pub fn seen(&self) -> Option<(Vec<u8>, [u8; 32])> {

        self.seen.lock().unwrap().clone()

    }

    pub fn client_config(self: &Arc<Self>) -> Result<Arc<rustls::ClientConfig>> {

        let mut c = rustls::ClientConfig::builder_with_provider(provider())

            .with_safe_default_protocol_versions()?

            .dangerous()

            .with_custom_certificate_verifier(self.clone())

            .with_no_client_auth();

        c.alpn_protocols = vec![b"http/1.1".to_vec()];

        Ok(Arc::new(c))

    }

    fn verify_server_cert(

        &self,

        end_entity: &CertificateDer<'_>,

        _intermediates: &[CertificateDer<'_>],

        _server_name: &ServerName<'_>,

        _ocsp: &[u8],

        _now: UnixTime,

    ) -> Result<ServerCertVerified, rustls::Error> {

        let spki = spki_sha256(end_entity).map_err(|_| rustls::Error::InvalidCertificate(rustls::CertificateError::BadEncoding))?;

        *self.seen.lock().unwrap() = Some((end_entity.to_vec(), spki));

        match self.pin {

            Some(pin) if pin != spki => Err(rustls::Error::InvalidCertificate(rustls::CertificateError::ApplicationVerificationFailure)),

            _ => Ok(ServerCertVerified::assertion()),

    }

    fn verify_tls13_signature(

        &self,

        message: &[u8],

        cert: &CertificateDer<'_>,

        dss: &DigitallySignedStruct,

    ) -> Result<HandshakeSignatureValid, rustls::Error> {

        rustls::crypto::verify_tls13_signature(message, cert, dss, &self.algs)

    }

    fn supported_verify_schemes(&self) -> Vec<SignatureScheme> {

        self.algs.supported_schemes()

    }

pub async fn connect(addr: &str, config: Arc<rustls::ClientConfig>) -> Result<tokio_rustls::client::TlsStream<TcpStream>> {

    let tcp = tokio::time::timeout(Duration::from_secs(10), TcpStream::connect(addr))

        .await

        .with_context(|| format!("connecting to {addr}: timed out"))?

        .with_context(|| format!("connecting to {addr}"))?;

    tcp.set_nodelay(true).ok();

    let host = addr.rsplit_once(':').map(|(h, _)| h.trim_start_matches('[').trim_end_matches(']')).unwrap_or(addr);

    let name = ServerName::try_from(host.to_string()).unwrap_or_else(|_| ServerName::try_from("spectra").unwrap());

    let tls = tokio_rustls::TlsConnector::from(config).connect(name, tcp).await.with_context(|| format!("TLS with {addr}"))?;

    Ok(tls)

}

    fn own_certificate_is_made_once_and_pinned() {

        let dir = tempfile::tempdir().unwrap();

        let src = CertSource::Own { names: vec!["localhost".into()] };

        let a = ServerTls::load(dir.path(), &src).unwrap();

        let b = ServerTls::load(dir.path(), &src).unwrap();

        assert_eq!(a.spki_sha256, b.spki_sha256, "the certificate is kept between starts");

            let mode = std::fs::metadata(dir.path().join("keys/tls.key")).unwrap().permissions().mode();

            assert_eq!(mode & 0o777, 0o600);

    }