Rust 基于tokio-rustls的简单的HTTPS Client/Server实现

use std::io;
use tokio::io::{copy, stdout as tokio_stdout, AsyncWriteExt};

use s04_async_https_connection_via_tcp::utils::tls;

#[tokio::main]
async fn main() -> io::Result<()> {
    let allow_insecure = true;
    // let sni = "www.baidu.com";
    // let dst_addr = "www.baidu.com";
    let sni = "baidu.com";
    let dst_addr = "baidu.com";
    let dst_port = 443;
    let content = format!("GET / HTTP/1.1\r\nHost: {}\r\n\r\n", sni);

    let (mut reader, mut writer) = tls::connect(dst_addr, dst_port, sni, allow_insecure).await?;
    writer.write_all(content.as_bytes()).await?;
    let mut stdout = tokio_stdout();
    copy(&mut reader, &mut stdout).await?;
    Ok(())
}

// utils/tls.rs
use std::convert::TryFrom;
use std::io;
use std::net::ToSocketAddrs;
use std::sync::Arc;
use tokio::io::split;
use tokio::io::{ReadHalf, WriteHalf};
use tokio::net::TcpStream;
use tokio_rustls::TlsConnector;

use tokio_rustls::rustls::{self, ClientConfig, OwnedTrustAnchor, RootCertStore};

struct NoCertVerifier {}

impl rustls::client::ServerCertVerifier for NoCertVerifier {
    fn verify_server_cert(
        &self,
        _end_entity: &rustls::Certificate,
        _intermediates: &[rustls::Certificate],
        _server_name: &rustls::ServerName,
        _scts: &mut dyn Iterator<Item = &[u8]>,
        _ocsp_response: &[u8],
        _now: std::time::SystemTime,
    ) -> Result<rustls::client::ServerCertVerified, rustls::Error> {
        Ok(rustls::client::ServerCertVerified::assertion())
    }
}

pub async fn connect(
    dst_addr: &str,
    dst_port: u16,
    sni: &str,
    allow_insecure: bool,
) -> io::Result<(
    ReadHalf<tokio_rustls::client::TlsStream<tokio::net::TcpStream>>,
    WriteHalf<tokio_rustls::client::TlsStream<tokio::net::TcpStream>>,
)> {
    let addr = (dst_addr, dst_port)
        .to_socket_addrs()?
        .next()
        .ok_or_else(|| io::Error::from(io::ErrorKind::NotFound))?;

    let mut root_store = RootCertStore::empty();
    root_store.add_server_trust_anchors(webpki_roots::TLS_SERVER_ROOTS.0.iter().map(|ta| {
        OwnedTrustAnchor::from_subject_spki_name_constraints(
            ta.subject,
            ta.spki,
            ta.name_constraints,
        )
    }));
    let mut config = ClientConfig::builder()
        .with_safe_defaults()
        .with_root_certificates(root_store)
        .with_no_client_auth();

    if allow_insecure {
        config
            .dangerous()
            .set_certificate_verifier(Arc::new(NoCertVerifier {}));
    }

    let connector = TlsConnector::from(Arc::new(config));
    let stream = TcpStream::connect(&addr).await?;

    let domain = rustls::ServerName::try_from(sni)
        .map_err(|_| io::Error::new(io::ErrorKind::InvalidInput, "invalid dnsname"))?;

    let stream = connector.connect(domain, stream).await?;
    // stream.write_all(content.as_bytes()).await?;

    // let (mut reader, mut writer) = split(stream);

    Ok(split(stream))
}

use s04_async_https_connection_via_tcp::{handle, init, utils::tls};
use std::{io, net::ToSocketAddrs};
use tokio::io::{copy, stdout as tokio_stdout, AsyncWriteExt};
use tokio::net::TcpListener;
use tokio::time::{sleep, Duration};

#[tokio::main]
async fn main() -> io::Result<()> {
    let addr = ("127.0.0.1", 443u16)
        .to_socket_addrs()?
        .next()
        .ok_or_else(|| io::Error::from(io::ErrorKind::NotFound))?;

    // 监听TCP连接
    let listener = TcpListener::bind(&addr).await?;
    let acceptor = init("pixiv.net.crt", "pixiv.net.key")?;
    // 等待3s(此时本地服务已经建立)， 然后发送一个HTTPS 请求
    tokio::spawn(async move {
        sleep(Duration::from_secs(3)).await;
        send_a_https_request().await
    });
    // 处理请求连接
    loop {
        match listener.accept().await {
            Ok((stream, _peer_addr)) => {
                let acceptor = acceptor.clone();
                match acceptor.accept(stream).await {
                    Ok(stream) => {
                        tokio::spawn(async move {
                            if let Err(_err) = handle(stream).await {
                                eprintln!("TLS Handler err: {:?}", _err);
                            }
                        });
                    }
                    Err(_err) => {
                        eprintln!("Tls err: {:?}", _err);
                    }
                }
            }
            Err(_err) => {
                eprintln!("Tcp err: {:?}", _err);
            }
        }
    }
}

async fn send_a_https_request() -> io::Result<()> {
    let allow_insecure = true;
    let sni = "pixiv.net";
    let dst_addr = "127.0.0.1";
    let dst_port = 443;
    let content = format!("GET / HTTP/1.1\r\nHost: {}\r\n\r\n", sni);

    let (mut reader, mut writer) = tls::connect(dst_addr, dst_port, sni, allow_insecure).await?;

    writer.write_all(content.as_bytes()).await?;

    let mut stdout = tokio_stdout();

    // stdout.write_all("\r\nReceived response from server: ".as_bytes()).await?;
    copy(&mut reader, &mut stdout).await?;

    Ok(())
}

pub mod utils;

use rustls_pemfile::{certs, rsa_private_keys};
use std::fs::File;
use std::io::{self, BufReader};
use std::path::Path;
use tokio::io::{split, AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt};
use tokio_rustls::rustls::{Certificate, PrivateKey};
use tokio_rustls::TlsAcceptor;

lazy_static::lazy_static! {
    static ref RESPONSE_403:&'static str = concat!(
        "HTTP/1.1 403 Forbidden\r\n" ,"Content-Length: 0\r\n" ,"Connection: closed\r\n\r\n"
    );
    static ref RESPONSE_200:&'static str = concat!(
        "HTTP/1.1 200 OK\r\n" ,
        "Content-Length: 11\r\n" ,
        "Connection: closed\r\n\r\n",
        "Are you OK?",
    );
}

pub fn load_certs(path: &Path) -> io::Result<Vec<Certificate>> {
    certs(&mut BufReader::new(File::open(path)?))
        .map_err(|_| io::Error::new(io::ErrorKind::InvalidInput, "invalid cert"))
        .map(|mut certs| certs.drain(..).map(Certificate).collect())
}

pub fn load_keys(path: &Path) -> io::Result<Vec<PrivateKey>> {
    rsa_private_keys(&mut BufReader::new(File::open(path)?))
        .map_err(|_| io::Error::new(io::ErrorKind::InvalidInput, "invalid key"))
        .map(|mut keys| keys.drain(..).map(PrivateKey).collect())
}

pub fn init(cert_path: &str, key_path: &str) -> io::Result<TlsAcceptor> {
    let certs = load_certs(Path::new(cert_path))?;
    let mut keys = load_keys(Path::new(key_path))?;

    let server_conf = tokio_rustls::rustls::ServerConfig::builder()
        .with_safe_defaults()
        .with_no_client_auth()
        .with_single_cert(certs, keys.remove(0))
        .map_err(|_err| io::Error::new(io::ErrorKind::InvalidInput, "TLS cert loading error"))?;
    Ok(TlsAcceptor::from(std::sync::Arc::new(server_conf)))
}

pub async fn handle<IO>(stream: IO) -> io::Result<()>
where
    IO: AsyncRead + AsyncWrite + Unpin + AsyncWriteExt,
{
    let (mut local_reader, mut local_writer) = split(stream);

    // 从头部读取信息
    let mut head = [0u8; 2048];
    let n = local_reader.read(&mut head[..]).await?;
    if n == 2048 {
        return Err(io::Error::new(
            io::ErrorKind::Other,
            "Receive a unexpected big size of header!!",
        ));
    }
    let head_str = std::str::from_utf8(&head[..n])
        .map_err(|x| io::Error::new(io::ErrorKind::Interrupted, x))?;
    println!("\r\nReceived request from client: \r\n{}\r\n", head_str);

    // 回复200OK
    local_writer.write_all(RESPONSE_200.as_bytes()).await?;
    local_writer.shutdown().await?;
    Ok(()) as io::Result<()>
}