tokio-gemini/examples/main.rs

use std::{io::Write, os::fd::AsFd, sync::Mutex};

use dashmap::DashMap;
use tokio::io::AsyncBufReadExt;
use tokio_gemini::certs::{
    fingerprint::{self, generate_fingerprint},
    insecure::AllowAllCertVerifier,
    verifier::{SelfsignedCert, SelfsignedCertVerifier},
};

#[tokio::main]
async fn main() -> Result<(), tokio_gemini::LibError> {
    let mut args = std::env::args();
    let mut insecure = false;
    let mut url = "gemini://geminiprotocol.net/".to_owned();
    _ = args.next(); // skip exe path
    if let Some(arg) = args.next() {
        if arg == "-k" {
            insecure = true;
            if let Some(arg) = args.next() {
                url = arg;
            }
        } else {
            url = arg;
        }
    }

    let client = if insecure {
        tokio_gemini::Client::builder()
            .with_custom_verifier(AllowAllCertVerifier::yes_i_know_what_i_am_doing())
            .build()
    } else {
        tokio_gemini::Client::builder()
            .with_webpki_roots() // TODO: do we really need them?
            .with_selfsigned_cert_verifier(CertVerifier::init().await?)
            .build()
    };

    let mut resp = client.request(&url).await?;

    {
        let status_code = resp.status().status_code();
        let status_num: u8 = status_code.into();
        eprintln!("{} {:?}", status_num, status_code);
    }

    if resp.status().reply_type() == tokio_gemini::ReplyType::Success {
        let mime = resp.mime()?;
        eprintln!("Mime: {}", mime);

        if mime.type_() == mime::TEXT {
            println!("{}", resp.text().await?);
        } else {
            eprintln!("Downloading into content.bin");
            let mut f = tokio::fs::File::create("content.bin").await?;
            tokio::io::copy(&mut resp.stream(), &mut f).await?;
        }
    } else {
        eprintln!("Message: {}", resp.message());
    }

    Ok(())
}

struct CertVerifier {
    f: Mutex<std::fs::File>,
    map: DashMap<String, SelfsignedCert>,
}

impl CertVerifier {
    async fn init() -> Result<Self, tokio_gemini::LibError> {
        let map = DashMap::new();

        if tokio::fs::try_exists("known_hosts").await? {
            let mut f = tokio::fs::OpenOptions::new()
                .read(true)
                .open("known_hosts")
                .await?;

            let mut reader = tokio::io::BufReader::new(&mut f);

            let mut buf = String::new();
            loop {
                buf.clear();
                let n = reader.read_line(&mut buf).await?;
                if n == 0 {
                    break;
                }

                // Format:
                // host <space> expires <space> hash-algo <space> fingerprint
                // Example:
                // dc09.ru 1722930541 sha512 dGVzdHRlc3R0ZXN0Cg
                if let [host, expires, algo, fp] = buf
                    .split_whitespace()
                    .take(4)
                    .collect::<Vec<&str>>()
                    .as_slice()
                {
                    let expires = if let Ok(num) = expires.parse::<u64>() {
                        num
                    } else {
                        eprintln!("Cannot parse expires = {:?} as u64", expires);
                        continue;
                    };

                    let algo = match algo {
                        &"sha256" => fingerprint::Algorithm::Sha256,
                        &"sha512" => fingerprint::Algorithm::Sha512,
                        &_ => {
                            eprintln!("Unknown hash algorithm {:?}, skipping", algo);
                            continue;
                        }
                    };

                    map.insert(
                        (*host).to_owned(),
                        SelfsignedCert {
                            algo,
                            fingerprint: (*fp).to_owned(),
                            expires,
                        },
                    );
                } else {
                    eprintln!("Cannot parse line: {:?}", buf);
                    continue;
                }
            }
        }

        let f = Mutex::new(
            std::fs::OpenOptions::new()
                .append(true)
                .create(true)
                .open("known_hosts")?,
        );

        Ok(CertVerifier { f, map })
    }
}

impl SelfsignedCertVerifier for CertVerifier {
    fn verify(
        &self,
        cert: &tokio_gemini::certs::verifier::CertificateDer,
        host: &str,
        _now: tokio_gemini::certs::verifier::UnixTime,
    ) -> Result<bool, tokio_rustls::rustls::Error> {
        if let Some(known_cert) = self.map.get(host) {
            // if host is found in known_hosts, compare certs
            let this_fp = generate_fingerprint(cert, known_cert.algo);
            if this_fp == known_cert.fingerprint {
                // current cert hash matches known cert hash
                eprintln!("Cert for {} matched: {}", &host, &this_fp);
                Ok(true)
            } else {
                // TODO (after implementing `expires`) update cert if known is expired
                eprintln!(
                    "Error: certs do not match! Possibly MitM attack.\nKnown FP: {}\nGot:      {}",
                    &known_cert.fingerprint, &this_fp,
                );
                Ok(false)
            }
        } else {
            // host is unknown, generate hash and add to known_hosts
            let this_fp = generate_fingerprint(cert, fingerprint::Algorithm::Sha512);
            eprintln!(
                "Warning: updating known_hosts with cert {} for {}",
                &this_fp, &host,
            );

            (|| {
                // trick with cloning file descriptor
                // because we are not allowed to mutate &self
                let mut f = std::fs::File::from(
                    self.f.lock().unwrap().as_fd().try_clone_to_owned().unwrap(),
                );
                f.write_all(host.as_bytes())?;
                f.write_all(b" 0 sha512 ")?; // TODO after implementing `expires`
                f.write_all(&this_fp.as_bytes())?;
                f.write(b"\n")?;
                Ok::<(), std::io::Error>(())
            })()
            .unwrap_or_else(|e| {
                eprintln!("Could not add cert to file: {:?}", e);
            });

            self.map.insert(
                host.to_owned(),
                SelfsignedCert {
                    algo: fingerprint::Algorithm::Sha512,
                    fingerprint: this_fp,
                    expires: 0, // TODO after implementing cert parsing in tokio-gemini
                },
            );

            Ok(true)
        }
    }
}