use std::prelude::v1::*;
use std::sync::Arc;
use parking_lot::Mutex;
use axum::extract::Extension;
#[cfg(feature = "controller")]
use axum_extra::typed_header::TypedHeader;
use axum::Json;
use axum::body::Body;
use axum::http::StatusCode;
use axum::response::{ Response };
use serde::{ Deserialize, Serialize };
#[cfg(feature = "controller")]
use product_os_command_control::{
ProductOSController,
authentication::AuthExchangeKeyData,
authentication::XProductOSVerifyHeader,
authentication::XProductOSCommandHeader,
Node,
};
use crate::CONTROLLER_LOCK_TIMEOUT;
#[derive(Debug, Deserialize, Serialize)]
#[serde(rename_all = "camelCase")]
struct StatusResponse {
status: Status,
address: String,
process_id: u32
}
#[derive(Debug, Deserialize, Serialize)]
#[serde(rename_all = "camelCase")]
enum Status {
Ok,
Error
}
fn build_response(status: StatusCode, body: String) -> Response<Body> {
Response::builder()
.status(status)
.header("content-type", "application/json")
.body(Body::from(body))
.unwrap_or_else(|_| {
let mut resp = Response::new(Body::from("{}"));
*resp.status_mut() = status;
resp
})
}
#[cfg(feature = "controller")]
#[allow(dead_code)]
fn with_controller_lock<F, R>(
controller_mutex: Arc<Mutex<ProductOSController>>,
operation: F
) -> Response<Body>
where
F: FnOnce(&mut ProductOSController) -> R,
R: Serialize,
{
match controller_mutex.try_lock_for(CONTROLLER_LOCK_TIMEOUT) {
None => build_response(StatusCode::GATEWAY_TIMEOUT, "{}".to_string()),
Some(mut controller) => {
let result = operation(&mut controller);
let response_body = serde_json::to_string(&result).unwrap_or_else(|_| "{}".to_string());
build_response(StatusCode::OK, response_body)
}
}
}
#[cfg(feature = "controller")]
pub fn command_responder<S>(router: &mut product_os_router::ProductOSRouter<S>)
where
S: Clone + Send + Sync + 'static
{
router.add_handler("/command/auth/exchange", product_os_router::Method::POST, exchange_handler);
router.add_handler("/command/status/alive", product_os_router::Method::POST, status_alive);
router.add_handler("/command/status/ping", product_os_router::Method::POST, check_handler);
router.add_handler("/command/status/announce", product_os_router::Method::POST, announce_handler);
}
#[cfg(feature = "controller")]
async fn exchange_handler(Extension(controller_mutex): Extension<Arc<Mutex<ProductOSController>>>,
Json(payload): Json<AuthExchangeKeyData>) -> Response<Body> {
match controller_mutex.try_lock_for(CONTROLLER_LOCK_TIMEOUT) {
None => build_response(StatusCode::GATEWAY_TIMEOUT, "{}".to_string()),
Some(mut controller) => {
let remote_identifier = payload.identifier;
let remote_session = payload.session;
let remote_public_key = payload.public_key;
match controller.get_registry().get_key(remote_identifier.as_str()) {
None => {
let (session, public_key) = controller.create_key_session();
controller.generate_key(session.as_str(), remote_public_key.as_slice(), remote_identifier, Some(remote_session.clone()));
let response = serde_json::to_string(&AuthExchangeKeyData {
identifier: controller.get_registry().get_me().get_identifier(),
session: remote_session,
public_key: public_key.to_vec()
}).unwrap_or_else(|_| "{}".to_string());
tracing::info!("Responding to Auth::Exchange request {:?}", response);
build_response(StatusCode::OK, response)
},
Some(_) => {
tracing::info!("Responding to Auth::Exchange request with Conflict status");
build_response(StatusCode::CONFLICT, "{}".to_string())
}
}
}
}
}
#[cfg(feature = "controller")]
async fn status_alive(TypedHeader(x_product_os_command): TypedHeader<XProductOSCommandHeader>,
TypedHeader(x_product_os_verify): TypedHeader<XProductOSVerifyHeader>,
Extension(controller_mutex): Extension<Arc<Mutex<ProductOSController>>>) -> Response<Body> {
match controller_mutex.try_lock_for(CONTROLLER_LOCK_TIMEOUT) {
None => build_response(StatusCode::GATEWAY_TIMEOUT, "{}".to_string()),
Some(controller) => {
match controller.authenticate_command_control::<&Node>(None, None, None, Some(x_product_os_command.value().as_str()), Some(x_product_os_verify.value().as_str())) {
Ok(_) => {
tracing::info!("Responding to Status::Alive request");
build_response(StatusCode::OK, "{ \"status\": \"ok\" }".to_string())
},
Err(e) => {
let error = &e.error;
tracing::error!("Error from Status::Alive request {:?}", error);
build_response(StatusCode::UNAUTHORIZED, serde_json::to_string(&e).unwrap_or_else(|_| "{}".to_string()))
}
}
}
}
}
#[cfg(feature = "controller")]
async fn check_handler(TypedHeader(x_product_os_command): TypedHeader<XProductOSCommandHeader>,
TypedHeader(x_product_os_verify): TypedHeader<XProductOSVerifyHeader>,
Extension(controller_mutex): Extension<Arc<Mutex<ProductOSController>>>) -> Response<Body> {
match controller_mutex.try_lock_for(CONTROLLER_LOCK_TIMEOUT) {
None => build_response(StatusCode::GATEWAY_TIMEOUT, "{}".to_string()),
Some(controller) => {
match controller.authenticate_command_control::<&Node>(None, None, None, Some(x_product_os_command.value().as_str()), Some(x_product_os_verify.value().as_str())) {
Ok(_) => {
let response = serde_json::to_string(&StatusResponse {
status: Status::Ok,
address: controller.get_registry().get_me().try_get_address().map(|a| a.to_string()).unwrap_or_default(),
process_id: controller.get_registry().get_me().get_process_id()
}).unwrap_or_else(|_| "{}".to_string());
tracing::info!("Responding to Status::Ping request {:?}", response);
build_response(StatusCode::OK, response)
},
Err(e) => {
let error = &e.error;
tracing::error!("Error from Status::Ping request {:?}", error);
build_response(StatusCode::UNAUTHORIZED, serde_json::to_string(&e).unwrap_or_else(|_| "{}".to_string()))
}
}
}
}
}
#[cfg(feature = "controller")]
async fn announce_handler(TypedHeader(x_product_os_command): TypedHeader<XProductOSCommandHeader>,
TypedHeader(x_product_os_verify): TypedHeader<XProductOSVerifyHeader>,
Extension(controller_mutex): Extension<Arc<Mutex<ProductOSController>>>,
Json(payload): Json<product_os_command_control::Node>) -> Response<Body> {
match controller_mutex.try_lock_for(CONTROLLER_LOCK_TIMEOUT) {
None => build_response(StatusCode::GATEWAY_TIMEOUT, "{}".to_string()),
Some(mut controller) => {
match controller.authenticate_command_control::<&Node>(None, Some(&payload), None, Some(x_product_os_command.value().as_str()), Some(x_product_os_verify.value().as_str())) {
Ok(_) => {
let identifier = payload.get_identifier();
let node = payload;
controller.upsert_node_local(identifier, node);
let response = serde_json::to_string(&StatusResponse {
status: Status::Ok,
address: controller.get_registry().get_me().try_get_address().map(|a| a.to_string()).unwrap_or_default(),
process_id: controller.get_registry().get_me().get_process_id()
}).unwrap_or_else(|_| "{}".to_string());
tracing::info!("Responding to Status::Announce request {:?}", response);
build_response(StatusCode::OK, response)
},
Err(e) => {
let error = &e.error;
tracing::error!("Error from Status::Announce request {:?}", error);
build_response(StatusCode::UNAUTHORIZED, serde_json::to_string(&e).unwrap_or_else(|_| "{}".to_string()))
}
}
}
}
}