etcd_admin.rs

use std::future::Future;
use std::pin::Pin;
use std::sync::Arc;

use bytes::Bytes;
use etcd_client::{ConnectOptions, DeleteOptions, GetOptions, SortOrder, SortTarget};
use http_body_util::{BodyExt, Full};
use hyper::body::Incoming;
use hyper::http;
use hyper::server::conn::http1::Builder;
use hyper::service::Service;
use hyper_util::rt::TokioIo;
use log::{debug, error, info};
use tokio::net::TcpListener;
use uuid::Uuid;

use crate::config::EtcdConfig;
use crate::router::Route;
use crate::wrapped_etcd_client::WrappedEtcdClient;

const BLACKLIST_PREFIX: &str = "blacklist";
const DOUBLE_WRITE_PREFIX: &str = "double_write";
const BLACKLIST_API_PREFIX: &str = "/api/admin/v1/blacklist";
const DOUBLE_WRITE_API_PREFIX: &str = "/api/admin/v1/double_write";

#[derive(Clone)]
struct AdminSvc {
    etcd_prefix: String,
    admin_key: String,
    etcd_client: Arc<WrappedEtcdClient>
}

impl Service<hyper::Request<Incoming>> for AdminSvc {
    type Response = hyper::Response<Full<Bytes>>;
    type Error = hyper::Error;
    type Future = Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>> + Send>>;

    fn call(&self, req: hyper::Request<Incoming>) -> Self::Future {
        let validate_result = self.validate_req(&req);
        if validate_result.is_err() {
            let err = validate_result.err().unwrap();
            error!("etcd admin server validate request error: {:?}", err);
            return Box::pin(async move {
                mk_resp(401, String::from("Unauthorized!"))
            });
        }

        let url = req.uri().path();
        let method = req.method().clone();
        debug!("etcd admin server recv request: {:?} {:?}", method, url);

        if !url.starts_with(BLACKLIST_API_PREFIX) && !url.starts_with(DOUBLE_WRITE_API_PREFIX) {
            return Box::pin(async move {
                mk_resp(404, String::from("not found"))
            });
        }
        let url_segments: Vec<&str> = url.split('/').collect();
        // url pattern like /api/admin/v1/{admin_type}/{id}
        if url_segments.len() > 6 {
            return Box::pin(async move {
                mk_resp(404, String::from("not found"))
            });
        }
        let admin_type = String::from(url_segments[4]);
        if admin_type != BLACKLIST_PREFIX && admin_type != DOUBLE_WRITE_PREFIX {
            return Box::pin(async move {
                mk_resp(404, String::from("type not supported"))
            });
        }
        let id = if url_segments.len() > 5 {
            Some(String::from(url_segments[5]))
        } else {
            None
        };

        let etcd_prefix = format!("{}/{}/", self.etcd_prefix, admin_type);
        let etcd_client = self.etcd_client.clone();
        match method {
            http::Method::HEAD => {
                Box::pin(async move {
                    mk_resp(200, String::from("ok"))
                })
            }
            http::Method::GET => {
                Box::pin(get(etcd_prefix, etcd_client, id))
            }
            http::Method::PUT => {
                let req_body = req.into_body();
                Box::pin(put(etcd_prefix, etcd_client, req_body, id))
            }
            http::Method::DELETE => {
                Box::pin(delete(etcd_prefix, etcd_client, id))
            }
            _ => {
                Box::pin(async move {
                    mk_resp(405, String::from("method not allowed"))
                })
            }
        }
    }
}

impl AdminSvc {
    fn validate_req(&self, req: &hyper::Request<Incoming>) -> anyhow::Result<()> {
        let headers = req.headers();
        let api_key = headers.get("Authorization").map(|v| v.to_str().unwrap_or_default());

        if api_key != Some(&self.admin_key) {
            return Err(anyhow::Error::msg("UnAuthorized!"));
        }

        Ok(())
    }
}

pub(crate) async fn start_admin_server(etcd_config: EtcdConfig) -> anyhow::Result<()> {
    if !etcd_config.admin_enabled {
        info!("etcd admin server not start, admin_enabled is false");
        return Ok(());
    }
    
    let etcd_client = Arc::new(WrappedEtcdClient::connect(
        etcd_config.endpoints,
        Some(ConnectOptions::new()
            .with_user(etcd_config.username, etcd_config.password)
            .with_timeout(std::time::Duration::from_millis(etcd_config.timeout))
        ),
        true
    ).await?);

    let admin_svc = AdminSvc { etcd_prefix: etcd_config.prefix, admin_key: etcd_config.admin_api_key, etcd_client };
    let admin_listener = TcpListener::bind(&etcd_config.admin_addr).await?;

    tokio::spawn(async move {
        while let Ok((client_stream, addr)) = admin_listener.accept().await {
            debug!("etcd admin server accept connection from: {:?}", addr);
            let io = TokioIo::new(client_stream);
            let admin_svc = admin_svc.clone();
            tokio::spawn(async move {
                if let Err(err) = Builder::new().serve_connection(io, admin_svc).await {
                    error!("etcd admin server serve connection error: {:?}", err);
                }
            });
        }
    });
    info!("etcd admin server start at: {}", etcd_config.admin_addr);
    Ok(())
}

fn mk_resp(status: u16, s: String) -> Result<hyper::Response<Full<Bytes>>, hyper::Error> {
    Ok(hyper::Response::builder().status(status).body(Full::new(Bytes::from(s))).unwrap())
}

async fn put(
    etcd_prefix: String,
    etcd_client: Arc<WrappedEtcdClient>,
    req_body: Incoming,
    id: Option<String>
) -> Result<hyper::Response<Full<Bytes>>, hyper::Error> {

    let id = if id.is_none() {
        Uuid::new_v4().to_string()
    } else {
        id.unwrap()
    };

    let req_body = req_body.collect().await?.to_bytes();
    let route = serde_json::from_slice::<Route>(req_body.as_ref());
    if route.is_err() {
        return mk_resp(400, String::from("request body is not correct"))
    }
    let route = route.unwrap();
    debug!("etcd admin server put request body: {:?}", route);

    if route.keys.is_empty() {
        return mk_resp(400, String::from("empty commands"))
    }

    let key = format!("{}{}", etcd_prefix, id);
    let put_resp = etcd_client.put(key, req_body, None).await;
    if put_resp.is_err() {
        let err = put_resp.err().unwrap();
        return mk_resp(500, String::from("server error: etcd put error: ") + err.to_string().as_str())
    }
    mk_resp(200, String::from("ok"))
}

async fn delete(
    etcd_prefix: String,
    etcd_client: Arc<WrappedEtcdClient>,
    id: Option<String>
) -> Result<hyper::Response<Full<Bytes>>, hyper::Error> {
    if id.is_none() {
        return mk_resp(400, String::from("missing id"))
    }
    let id = id.unwrap();

    let key = format!("{}{}", etcd_prefix, id);
    let delete_resp = etcd_client.delete(key, Some(DeleteOptions::new())).await;
    if delete_resp.is_err() {
        let err = delete_resp.err().unwrap();
        return mk_resp(500, String::from("server error: etcd delete error: ") + err.to_string().as_str())
    }
    debug!("etcd admin server deleted from etcd, keys num: {}", delete_resp.unwrap().deleted());
    mk_resp(200, String::from("ok"))
}

async fn get(etcd_prefix: String,
             etcd_client: Arc<WrappedEtcdClient>,
             id: Option<String>
) -> Result<hyper::Response<Full<Bytes>>, hyper::Error> {
    let key = match id {
        Some(id) => { //query with id prefix
            format!("{}{}", etcd_prefix, id)
        }
        _ => {
            etcd_prefix
        }
    };

    let option = GetOptions::new()
        .with_prefix()
        .with_sort(SortTarget::Key, SortOrder::Ascend);
    let list_resp = etcd_client.get(key, Some(option)).await;
    if list_resp.is_err() {
        let err = list_resp.err().unwrap();
        return mk_resp(500, String::from("server error: etcd get error: ") + err.to_string().as_str())
    }
    let list_resp = list_resp.unwrap();

    let mut result = Vec::with_capacity(list_resp.kvs().len());
    list_resp.kvs().iter().for_each(|kv| {
        let key = kv.key_str();
        let value = kv.value_str();
        if key.is_err() || value.is_err() {
            error!("etcd admin server get from etcd is not correct, will ignore it!");
            return
        }
        let key = key.unwrap();
        let value = value.unwrap();

        let route: serde_json::error::Result<Route> = serde_json::from_str(value);
        if route.is_err() {
            error!("etcd admin server get from etcd is not correct, will ignore it!");
            return
        }
        debug!("etcd admin server get from etcd, key: {}, value: {}", key, value);
        result.push(route.unwrap());
    });

    let result = serde_json::to_string(&result);
    if result.is_err() {
        return mk_resp(500, String::from("server error: serde_json to string error!"))
    }
    mk_resp(200, result.unwrap())
}