Files
local-mini-kms/src/serve.rs
2022-07-29 00:35:05 +08:00

357 lines
13 KiB
Rust

use std::net::SocketAddr;
use std::sync::Mutex;
use clap::{App, Arg, ArgMatches, SubCommand};
use hyper::{Body, Client, Method, Request, Response, Server, StatusCode};
use hyper::body::Buf;
use hyper::client::HttpConnector;
use hyper::server::conn::AddrStream;
use hyper::service::{make_service_fn, service_fn};
use josekit::jwk::alg::rsa::RsaKeyPair;
use josekit::jwk::KeyPair;
use rust_util::{debugging, failure_and_exit, information, opt_result, simple_error, success, XResult};
use rust_util::util_clap::{Command, CommandError};
use seckey::SecBytes;
use serde::{Deserialize, Serialize};
use serde_json::{json, Map, Value};
use zeroize::Zeroize;
use crate::{db, jose, proc};
use crate::db::Key;
type GenericError = Box<dyn std::error::Error + Send + Sync>;
type Result<T> = std::result::Result<T, GenericError>;
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").takes_value(true).default_value("127.0.0.1:5567").help("Listen"))
.arg(Arg::with_name("local-db").long("local-db").takes_value(true).default_value("local-mini-kms.db").help("Local db file"))
}
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 listen = sub_arg_matches.value_of("listen").expect("Get argument listen error");
let rt = tokio::runtime::Runtime::new().expect("Create tokio runtime 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<Body>,
_client: Client<HttpConnector>,
) -> Result<Response<Body>> {
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") => init(req).await,
(&Method::POST, "/update") => update().await,
(&Method::POST, "/decrypt") => decrypt(req).await,
(&Method::POST, "/encrypt") => encrypt(req).await,
(&Method::GET, "/status") => status().await,
(&Method::GET, "/version") => get_version().await,
_ => Ok(Response::builder().status(StatusCode::NOT_FOUND).body(serde_json::to_string_pretty(&json!({ "error": "not_found" }))?.into())?),
}
}
macro_rules! do_response {
($ex: expr) => (
match $ex {
Ok((status_code, body)) => Ok(Response::builder().status(status_code).body(serde_json::to_string_pretty(&body)?.into())?),
Err(e) => Ok(Response::builder().status(StatusCode::INTERNAL_SERVER_ERROR).body(serde_json::to_string_pretty(&json!({
"error": "internal_error",
"error_message": format!("{}", e),
}))?.into())?),
}
)
}
// -------------------------------------------------------------------------------------------------
struct MemoryKey {
database_file: String,
instance_rsa_key_pair: RsaKeyPair,
master_key: Option<SecBytes>,
}
lazy_static::lazy_static! {
static ref STATUP_RW_LOCK: Mutex<Option<MemoryKey>> = Mutex::new(None);
}
fn init_instance(db: &str) -> XResult<bool> {
let conn = db::open_db(db)?;
db::init_db(&conn)?;
let mut startup_rw_lock = 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<bool> {
let mut startup_rw_lock = 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),
}
}
#[derive(Serialize, Deserialize)]
struct MultipleViewValue {
value: Option<String>,
value_hex: Option<String>,
value_base64: Option<String>,
}
impl MultipleViewValue {
fn from(v: &[u8]) -> Self {
Self {
value: Some(String::from_utf8_lossy(v).to_string()),
value_hex: Some(hex::encode(v)),
value_base64: Some(base64::encode(v)),
}
}
fn to_bytes(&self) -> XResult<Vec<u8>> {
if let Some(v) = &self.value {
Ok(v.as_bytes().to_vec())
} else if let Some(v) = &self.value_hex {
let v = opt_result!(hex::decode(v), "Decode hex failed: {}");
Ok(v)
} else if let Some(v) = &self.value_base64 {
let v = opt_result!(base64::decode(v), "Decode base64 failed: {}");
Ok(v)
} else {
simple_error!("Multiple view value is all empty")
}
}
}
#[derive(Serialize, Deserialize)]
struct DecryptRequest {
encrypted_value: String,
}
async fn decrypt(req: Request<Body>) -> Result<Response<Body>> {
do_response!(inner_decrypt(req).await)
}
async fn inner_decrypt(req: Request<Body>) -> XResult<(StatusCode, Value)> {
let whole_body = hyper::body::aggregate(req).await?;
let data: DecryptRequest = serde_json::from_reader(whole_body.reader())?;
debugging!("To be decrypted value: {}", &data.encrypted_value);
let key = match get_master_key() {
None => return Ok((StatusCode::BAD_REQUEST, json!({ "error": "status_not_ready" }))),
Some(key) => key,
};
let decrypted_value = jose::deserialize_jwe_aes(&data.encrypted_value, &*key.read());
drop(key);
decrypted_value.map(|v| {
let v = MultipleViewValue::from(&v.0);
let mut map = Map::new();
if let Some(v) = &v.value {
map.insert("value".to_string(), Value::String(v.to_string()));
}
if let Some(v) = &v.value_hex {
map.insert("value_hex".to_string(), Value::String(v.to_string()));
}
if let Some(v) = &v.value_base64 {
map.insert("value_base64".to_string(), Value::String(v.to_string()));
}
(StatusCode::OK, Value::Object(map))
})
}
async fn encrypt(req: Request<Body>) -> Result<Response<Body>> {
do_response!(inner_encrypt(req).await)
}
async fn inner_encrypt(req: Request<Body>) -> XResult<(StatusCode, Value)> {
let whole_body = hyper::body::aggregate(req).await?;
let data: MultipleViewValue = serde_json::from_reader(whole_body.reader())?;
let value = data.to_bytes()?;
let key = match get_master_key() {
None => return Ok((StatusCode::BAD_REQUEST, json!({ "error": "status_not_ready" }))),
Some(key) => key,
};
let encrypt_result = jose::serialize_jwe_aes(&value, &*key.read());
drop(key);
encrypt_result.map(|e| {
(StatusCode::OK, json!({
"encrypted_value": e,
}))
})
}
async fn update() -> Result<Response<Body>> {
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 init(req: Request<Body>) -> Result<Response<Body>> {
do_response!(inner_init(req).await)
}
#[derive(Serialize, Deserialize)]
struct InitRequest {
clear_master_key_hex: Option<String>,
clear_master_key_base64: Option<String>,
encrypted_master_key: Option<String>,
}
async fn inner_init(req: Request<Body>) -> XResult<(StatusCode, Value)> {
let whole_body = hyper::body::aggregate(req).await?;
let init_request: InitRequest = serde_json::from_reader(whole_body.reader())?;
let mut startup_rw_lock = STATUP_RW_LOCK.lock().expect("Lock read startup rw lock error");
match &*startup_rw_lock {
None => return Ok((StatusCode::INTERNAL_SERVER_ERROR, json!({ "error": "internal_error", "error_message": "not init " }))),
Some(memory_key) => match memory_key.master_key {
Some(_) => return Ok((StatusCode::BAD_REQUEST, json!({ "error": "bad_request", "error_message": "already init " }))),
None => {}
},
}
let clear_master_key = if let Some(clear_master_key_base64) = &init_request.clear_master_key_base64 {
base64::decode(clear_master_key_base64)?
} else if let Some(clear_master_key_hex) = init_request.clear_master_key_hex {
hex::decode(clear_master_key_hex)?
} else if let Some(encrypted_master_key) = init_request.encrypted_master_key {
debugging!("Received encrypted master key: {}", encrypted_master_key);
if let Some(k) = &*startup_rw_lock {
let (clear_master_key, _) = jose::deserialize_jwe_rsa(&encrypted_master_key, &k.instance_rsa_key_pair.to_jwk_private_key())?;
clear_master_key
} else {
return Ok((StatusCode::INTERNAL_SERVER_ERROR, json!({ "error": "internal_error", "error_message": "not init " })));
}
} else {
return Ok((StatusCode::BAD_REQUEST, json!({ "error": "bad_request", "error_message": "master key is not assigned" })));
};
if clear_master_key.len() != 32 {
return Ok((StatusCode::BAD_REQUEST, json!({ "error": "bad_request", "error_message": "bad clear_master_key length" })));
}
if let Some(k) = &mut *startup_rw_lock {
let conn = opt_result!(db::open_db(&k.database_file), "Open db failed: {}");
let default_master_key_verification_key = db::find_key(&conn, db::DEFAULT_MASTER_KEY_VERIFICATION_KEY)?;
match default_master_key_verification_key {
None => {
let key = Key {
name: db::DEFAULT_MASTER_KEY_VERIFICATION_KEY.to_string(),
encrypted_key: jose::serialize_jwe_aes("LOCAL-MINI-KMS:MAGIC-VERIFICATION-KEY".as_bytes(), &clear_master_key)?,
};
db::insert_key(&conn, &key)?;
}
Some(key) => {
debugging!("Found jwe: {}", &key.encrypted_key);
let _ = opt_result!(jose::deserialize_jwe_aes(&key.encrypted_key, &clear_master_key), "Deserialize master key verification key failed: {}");
}
}
information!("Set master key success");
let sec_bytes = SecBytes::with(clear_master_key.len(), |buf| buf.copy_from_slice(&clear_master_key.as_slice()[..]));
let mut clear_master_key = clear_master_key;
clear_master_key.zeroize();
k.master_key = Some(sec_bytes);
k.instance_rsa_key_pair = jose::generate_rsa_key(4096)?;
}
Ok((StatusCode::OK, json!({})))
}
async fn status() -> Result<Response<Body>> {
do_response!(inner_status().await)
}
async fn inner_status() -> XResult<(StatusCode, Value)> {
let startup_rw_lock = STATUP_RW_LOCK.lock().expect("Lock read startup rw lock error");
let body = match &*startup_rw_lock {
None => json!({ "status": "n/a" }),
Some(memory_key) => match memory_key.master_key {
None => json!({
"status": "not-ready",
"instance_public_key_jwk": memory_key.instance_rsa_key_pair.to_jwk_key_pair().to_public_key()?,
"instance_public_key_pem": String::from_utf8_lossy(&memory_key.instance_rsa_key_pair.to_pem_public_key()).to_string(),
}),
Some(_) => json!({
"status": "ready",
}),
}
};
Ok((StatusCode::OK, body))
}
async fn get_version() -> Result<Response<Body>> {
Ok(Response::builder().body(format!(
"{} - {}\n", env!("CARGO_PKG_NAME"), env!("CARGO_PKG_VERSION")
).into())?)
}
fn get_master_key() -> Option<SecBytes> {
let startup_rw_lock = STATUP_RW_LOCK.lock().expect("Lock read startup rw lock error");
match &*startup_rw_lock {
None => None,
Some(k) => match &k.master_key {
None => None,
Some(k) => {
let k = &*k.read();
Some(SecBytes::with(k.len(), |buf| buf.copy_from_slice(k)))
}
},
}
}