use std::net::SocketAddr; use clap::{App, Arg, ArgMatches, SubCommand}; use hyper::{Body, Client, Method, Request, Response, Server, StatusCode}; use hyper::client::HttpConnector; use hyper::server::conn::AddrStream; use hyper::service::{make_service_fn, service_fn}; use rust_util::{failure_and_exit, information, success, warning, XResult}; use rust_util::util_clap::{Command, CommandError}; use serde_json::{json, Value}; use tokio::runtime::Runtime; use crate::{db, jose, proc}; use crate::do_response; use crate::serve_common::{self, GenericError, MemoryKey, Result}; use crate::serve_encrypt_decrypt; use crate::serve_init; use crate::serve_init::InitRequest; use crate::serve_read_write; use crate::serve_status; use crate::yubikey_hmac; pub struct CommandImpl; impl Command for CommandImpl { fn name(&self) -> &str { "serve" } fn subcommand<'a>(&self) -> App<'a, 'a> { SubCommand::with_name(self.name()).about("Local mini KMS serve") .arg(Arg::with_name("listen").long("listen").short("L").takes_value(true).default_value("127.0.0.1:5567").help("Listen")) .arg(Arg::with_name("local-db").long("local-db").short("d").takes_value(true).default_value("local-mini-kms.db").help("Local db file")) .arg(Arg::with_name("yubikey-challenge").long("yubikey-challenge").short("c").takes_value(true).help("Yubikey challenge")) .arg(Arg::with_name("init-encrypted-master-key").long("init-encrypted-master-key").short("k").takes_value(true).help("Init encrypted mater key")) } fn run(&self, _arg_matches: &ArgMatches, sub_arg_matches: &ArgMatches) -> CommandError { let local_mini_kms_db = sub_arg_matches.value_of("local-db").expect("Get local mini kms db error"); match init_instance(local_mini_kms_db) { Ok(true) => success!("Init server success"), Ok(false) => failure_and_exit!("SHOULD NOT HAPPEN, server already init"), Err(e) => failure_and_exit!("Init server failed: {}", e), } let rt = Runtime::new().expect("Create tokio runtime error"); init_with_yubikey_challenge(&rt, sub_arg_matches); let listen = sub_arg_matches.value_of("listen").expect("Get argument listen error"); rt.block_on(async { let addr = listen.parse().expect(&format!("Parse listen error: {}", listen)); let client = Client::new(); let new_service = make_service_fn(move |conn: &AddrStream| { let remote_addr = conn.remote_addr(); let client = client.clone(); async move { Ok::<_, GenericError>(service_fn(move |req| { response_requests(remote_addr, req, client.to_owned()) })) } }); let server = Server::bind(&addr).serve(new_service); information!("Listening on http://{}", addr); match server.await { Err(e) => failure_and_exit!("Server error: {}", e), Ok(_) => success!("Server ended"), } }); Ok(Some(0)) } } // ref: https://github.com/hyperium/hyper/blob/master/examples/web_api.rs // ref: https://crates.io/crates/rusqlite async fn response_requests( remote_addr: SocketAddr, req: Request, _client: Client, ) -> Result> { let process = proc::get_process(remote_addr.port()); information!("Receive request: {}, from: {}, process: {:?}", req.uri(), remote_addr, process); match (req.method(), req.uri().path()) { (&Method::POST, "/init") => serve_init::init(req).await, (&Method::POST, "/update") => update().await, (&Method::POST, "/decrypt") => serve_encrypt_decrypt::decrypt(req).await, (&Method::POST, "/encrypt") => serve_encrypt_decrypt::encrypt(req).await, (&Method::POST, "/read") => serve_read_write::read(req).await, (&Method::POST, "/write") => serve_read_write::write(req).await, (&Method::GET, "/status") => serve_status::status().await, (&Method::GET, "/version") => get_version().await, (&Method::GET, "/") => get_root().await, _ => Ok(Response::builder() .status(StatusCode::NOT_FOUND) .body(format!("{}\n", serde_json::to_string_pretty(&json!({ "error": "not_found", }))?).into())?), } } // ------------------------------------------------------------------------------------------------- fn init_instance(db: &str) -> XResult { let conn = db::open_db(db)?; db::init_db(&conn)?; let mut startup_rw_lock = serve_common::STATUP_RW_LOCK.lock().expect("Lock write startup rw lock error"); match &*startup_rw_lock { Some(_) => Ok(false), None => { let memory_key = MemoryKey { database_file: db.to_string(), instance_rsa_key_pair: jose::generate_rsa_key(4096)?, master_key: None, }; *startup_rw_lock = Some(memory_key); Ok(true) } } } fn update_instance_rsa_key_pair() -> XResult { let mut startup_rw_lock = serve_common::STATUP_RW_LOCK.lock().expect("Lock write startup rw lock error"); match &mut *startup_rw_lock { Some(k) => { k.instance_rsa_key_pair = jose::generate_rsa_key(4096)?; Ok(true) } None => Ok(false), } } async fn update() -> Result> { do_response!(inner_update().await) } async fn inner_update() -> XResult<(StatusCode, Value)> { let update = update_instance_rsa_key_pair()?; Ok((StatusCode::OK, json!({ "update": update, }))) } async fn get_version() -> Result> { Ok(Response::builder().body(format!( "{} - {}\n", env!("CARGO_PKG_NAME"), env!("CARGO_PKG_VERSION") ).into())?) } async fn get_root() -> Result> { if std::env::var("LOCAL_MINI_KMS_HELP").is_ok() { Ok(Response::builder().body(format!( r##"{} - {} Supports commands: - GET /version - GET /status - POST /init - POST /update - POST /encrypt - POST /decrypt - POST /read - POST /write "##, env!("CARGO_PKG_NAME"), env!("CARGO_PKG_VERSION") ).into())?) } else { Ok(Response::builder().body("Root Not Found\n".into())?) } } fn init_with_yubikey_challenge(rt: &Runtime, sub_arg_matches: &ArgMatches) { let mut yubikey_challenge = sub_arg_matches.value_of("yubikey-challenge").map(ToString::to_string); let init_encrypted_master_key = sub_arg_matches.value_of("init-encrypted-master-key"); if init_encrypted_master_key.is_some() && yubikey_challenge.is_none() { yubikey_challenge = rpassword::prompt_password("Yubikey challenge: ").ok(); } let (challenge_key, init_encrypted_master_key) = match (yubikey_challenge, init_encrypted_master_key) { (Some(yubikey_challenge), Some(init_encrypted_master_key)) => { match yubikey_hmac::yubikey_challenge_as_32_bytes(yubikey_challenge.as_bytes()) { Err(e) => { warning!("Yubikey challenge failed: {}", e); return; } Ok(challenge_key) => (challenge_key, init_encrypted_master_key), } } (Some(_), None) | (None, Some(_)) => { warning!("Arguments yubikey-challenge and init-encrypted-master-key should both assigned."); return; } _ => return, }; match jose::deserialize_jwe_aes(&init_encrypted_master_key, &challenge_key) { Err(e) => warning!("Yubikey seal master key failed: {}", e), Ok((key, _)) => { success!("Yubikey un-seal master key success"); let init_master_key_result = rt.block_on(async { serve_init::inner_init_request(InitRequest { yubikey_challenge: None, clear_master_key_hex: Some(hex::encode(&key)), clear_master_key_base64: None, encrypted_master_key: None, }).await }); match init_master_key_result { Err(e) => warning!("Init master key failed: {}", e), Ok((StatusCode::OK, _)) => success!("Init master key success"), Ok((_, response)) => warning!("Init master failed: {}", response), } } } }