keyfork/keyforkd/src/main.rs

use std::{collections::HashMap, path::PathBuf};

use keyfork_mnemonic_util::Mnemonic;

use tokio::io::{self, AsyncBufReadExt, BufReader};
use tower::ServiceBuilder;

#[cfg(feature = "tracing")]
use tracing::debug;
#[cfg(feature = "tracing")]
use tracing_subscriber::{
    filter::{EnvFilter, LevelFilter},
    fmt::{format::FmtSpan, layer},
    prelude::*,
    registry,
};

mod error;
mod server;
mod service;
mod middleware;
use error::KeyforkdError;
use server::UnixServer;
use service::Keyforkd;

type Result<T, E = Box<dyn std::error::Error>> = std::result::Result<T, E>;

async fn load_mnemonic() -> Result<Mnemonic> {
    let mut stdin = BufReader::new(io::stdin());
    let mut line = String::new();
    stdin.read_line(&mut line).await?;
    Ok(line.parse()?)
}

#[cfg(feature = "tracing")]
fn setup_registry() {
    let envfilter = EnvFilter::builder()
        .with_default_directive(LevelFilter::DEBUG.into())
        .from_env_lossy();
    registry()
        .with(envfilter)
        .with(layer().with_span_events(FmtSpan::CLOSE))
        .with(tracing_error::ErrorLayer::default())
        .init();
}

#[cfg_attr(feature = "multithread", tokio::main)]
#[cfg_attr(not(feature = "multithread"), tokio::main(flavor = "current_thread"))]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    #[cfg(feature = "tracing")]
    setup_registry();

    #[cfg(feature = "tracing")]
    debug!("reading mnemonic from standard input");
    let mnemonic = load_mnemonic().await?;

    let service = ServiceBuilder::new()
        .layer(middleware::SerdeLayer::new())
        .service(Keyforkd::new(mnemonic));

    let runtime_vars = std::env::vars()
        .filter(|(key, _)| ["XDG_RUNTIME_DIR", "KEYFORKD_SOCKET_PATH"].contains(&key.as_str()))
        .collect::<HashMap<String, String>>();
    let mut runtime_path: PathBuf;
    #[allow(clippy::single_match_else)]
    match runtime_vars.get("KEYFORKD_SOCKET_PATH") {
        Some(occupied) => {
            runtime_path = PathBuf::from(occupied);
        }
        None => {
            runtime_path = PathBuf::from(
                runtime_vars
                    .get("XDG_RUNTIME_DIR")
                    .ok_or(KeyforkdError::NoSocketPath)?,
            );
            runtime_path.push("keyforkd");
            #[cfg(feature = "tracing")]
            debug!("ensuring directory exists: {}", runtime_path.display());
            if !runtime_path.is_dir() {
                tokio::fs::create_dir(&runtime_path).await?;
            }
            runtime_path.push("keyforkd.sock");
        }
    }

    #[cfg(feature = "tracing")]
    debug!(
        "binding UNIX socket in runtime dir: {}",
        runtime_path.display()
    );

    let mut server = match UnixServer::bind(&runtime_path) {
        Ok(s) => s,
        Err(e) => {
            #[cfg(feature = "tracing")]
            debug!(%e, "Encountered error attempting to bind socket: {}", runtime_path.display());
            return Err(e.into());
        }
    };
    match server.run(service).await {
        Ok(_) => (),
        Err(e) => {
            #[cfg(feature = "tracing")]
            debug!(%e, "Encountered error while running");
        }
    }
    Ok(())
}