path: root/src/lib.rs

use hyper::{Request, Body, Method, Uri};
use http::request;
use serde::{Serialize, Deserialize};
use std::collections::HashMap;


#[derive(Default, Deserialize, Debug, Clone)]
pub struct ProxyConfig {
    pub bind_addr: Option<String>,      // 127.0.0.1:9292
    pub upstream_addr: Option<String>,  // 127.0.0.1:9200
    pub index: Vec<IndexConfig>
}

#[derive(Deserialize, Debug, Clone)]
pub struct IndexConfig {
    pub name: String,
}

impl ProxyConfig {

    pub fn allow_index(&self, name: &str) -> bool {
        for index in &self.index {
            if index.name == name {
                return true
            }
        }
        false
    }
}

#[derive(Serialize, Deserialize, Debug)]
pub struct ApiRequest {
    pub method: String,
    pub path_and_query: String,
    pub body: Option<SearchBody>,
}

pub struct UrlQueryParams {
    pub allow_no_indices: Option<bool>,
    pub allow_partial_search_results: Option<bool>,
    pub batched_reduce_size: Option<u32>,
    pub ccs_minimize_roundtrips: Option<bool>,
    pub docvalue_fields: Option<String>, // array of strings, comma-separated
    pub expand_wildcards: Option<String>,
    pub explain: Option<bool>,
    pub from: Option<u32>,
    pub ignore_throttled: Option<bool>,
    pub ignore_unavailable: Option<bool>,
    pub max_concurrent_shard_requests: Option<u32>,
    pub pre_filter_shard_size: Option<u32>,
    pub preference: Option<String>,
    pub q: Option<String>,
    pub request_cache: Option<bool>,
    pub rest_total_hits_as_int: Option<bool>,
    pub routing: Option<String>,
    pub scroll: Option<String>, // string is "time value"
    pub search_type: Option<String>,
    pub seq_no_primary_term: Option<bool>,
    pub size: Option<u32>,
    pub sort: Option<String>, // array of strings, comma-separated
    pub _source: Option<bool>, // TODO: bool or string
    pub _source_excludes: Option<String>, // array of strings, comma-separated
    pub _source_includes: Option<String>, // array of strings, comma-separated
    pub stats: Option<String>,
    pub stored_fields: Option<String>, // array of strings, comma-separated
    pub suggest_field: Option<String>,
    pub suggest_text: Option<String>,
    pub terminate_after: Option<u32>,
    pub timeout: Option<String>, // string is "time units"
    pub track_scores: Option<bool>,
    pub track_total_hits: Option<bool>, // XXX: bool or integer
    pub typed_keys: Option<bool>,
    pub version: Option<bool>,
}

// https://www.elastic.co/guide/en/elasticsearch/reference/current/search-search.html
#[derive(Serialize, Deserialize, Debug)]
#[serde(deny_unknown_fields)]
pub struct SearchBody {
    pub query: Option<ApiQuery>,
    pub highlight: Option<ApiHighlight>,
    pub collapse: Option<ApiCollapse>,
    pub post_filter: Option<ApiQuery>, // TODO: leaf query only?
    pub rescore: Option<ApiRescore>, // TODO: single or an array of rescore objects
    // script_fields disabled

    // https://www.elastic.co/guide/en/elasticsearch/reference/current/sort-search-results.html
    pub sort: Option<Vec<SortElement>>,

    pub slice: Option<ApiSlice>,
    pub stored_fields: Option<String>, // array of strings, or "_none_"

    // overlap with URL query parameters
    pub docvalue_fields: Option<Vec<DocValOrString>>,
    pub explain: Option<bool>,
    pub from: Option<u32>,
    pub min_score: Option<f64>,
    pub seq_no_primary_term: Option<bool>,
    pub size: Option<u32>,
    pub _source: Option<bool>, // XXX: bool, string, or object
    pub terminate_after: Option<u32>,
    pub timeout: Option<String>, // string is "time units"
}

#[derive(Serialize, Deserialize, Debug)]
#[serde(deny_unknown_fields)]
pub struct ScrollBody {
    pub scroll_id: String,
    pub scroll: Option<String>,
}

#[derive(Serialize, Deserialize, Debug)]
pub struct ApiSlice {
    id: u32,
    max: u32,
    field: Option<String>,
}

#[derive(Serialize, Deserialize, Debug)]
pub struct ApiRescore{
    pub query: Option<ApiQuery>,
    pub window_size: Option<u32>,
}

// TODO: could revert to having query types as an enum, with flattening
#[derive(Serialize, Deserialize, Debug)]
pub struct ApiQuery {
    // compound queries
    #[serde(rename = "bool")]
    bool_query: Option<BoolQuery>,
    boosting: Option<BoostingQuery>,
    constant_score: Option<ConstantScoreQuery>,

    // fulltext (leaf) queries

    // term-level (leaf) queries
    #[serde(rename = "match")]
    match_query: Option<HashMap<String, QueryFieldOrString>>,
    match_phrase: Option<HashMap<String, QueryFieldOrString>>,
    query_string: Option<QueryField>,

    // other
    nested: Option<NestedQuery>,
    rescore_query: Option<Box<ApiQuery>>,
}

#[derive(Serialize, Deserialize, Debug)]
pub struct ApiHighlight{
    // TODO: fields could also be an array of strings?
    fields: HashMap<String, HighlightField>,

    #[serde(flatten)]
    settings: HighlightField,
}

#[derive(Serialize, Deserialize, Debug)]
#[serde(untagged)]
pub enum SortMapValue {
    String(String),
    Object { order: String, mode: Option<String> },
}

#[derive(Serialize, Deserialize, Debug)]
#[serde(untagged)]
pub enum SortElement{
    String(String),
    Object(HashMap<String, SortMapValue>),
}

#[derive(Serialize, Deserialize, Debug)]
#[serde(untagged)]
pub enum DocValOrString {
    String(String),
    Object {field: String, format: Option<String>},
}

#[derive(Serialize, Deserialize, Debug)]
#[serde(untagged)]
pub enum QueryFieldOrString {
    Object(QueryField),
    String(String),
}

#[derive(Serialize, Deserialize, Debug)]
pub struct QueryField{
    query: String,
    fuzziness: Option<String>,
    slop: Option<u32>,
    boost: Option<f64>,
}

#[derive(Serialize, Deserialize, Debug)]
pub struct BoolQuery {
    must: Option<Box<ApiQuery>>,
    filter: Option<Box<ApiQuery>>,
    should: Option<Box<ApiQuery>>,
    must_not: Option<Box<ApiQuery>>,
    minimum_should_match: Option<u32>,
    boost: Option<f64>,
}

#[derive(Serialize, Deserialize, Debug)]
pub struct NestedQuery {
    path: String,
    query: Box<ApiQuery>,
    score_mode: Option<String>,
    ignore_unmapped: Option<bool>,
}

#[derive(Serialize, Deserialize, Debug)]
pub struct BoostingQuery {
    positive: Box<ApiQuery>,
    negative: Box<ApiQuery>,
    negative_boost: f64,
}

#[derive(Serialize, Deserialize, Debug)]
pub struct ConstantScoreQuery {
    filter: Box<ApiQuery>,
    boost: Option<f64>,
}

// https://www.elastic.co/guide/en/elasticsearch/reference/current/highlighting.html
#[derive(Serialize, Deserialize, Debug)]
pub struct HighlightField{
    boundary_chars: Option<String>,
    boundary_max_scan: Option<u32>,
    boundary_scanner: Option<String>,
    boundary_scanner_locale: Option<String>,
    encoder: Option<String>,
    force_source: Option<bool>,
    fragmenter: Option<String>,
    fragment_offset: Option<u32>,
    fragment_size: Option<u32>,
    highlight_query: Option<ApiQuery>,
    matched_fields: Option<Vec<String>>,
    no_match_size: Option<u32>,
    number_of_fragments: Option<u32>,
    order: Option<String>,
    phrase_limit: Option<u32>,
    pre_tags: Option<Vec<String>>,
    post_tags: Option<Vec<String>>,
    require_field_match: Option<bool>,
    tags_schema: Option<String>,
    #[serde(rename = "type")]
    highlight_type: Option<String>,
}

#[derive(Serialize, Deserialize, Debug)]
pub struct ApiCollapse{
    field: String,
    inner_hits: Option<InnerHitsOneOrMore>,
}

#[derive(Serialize, Deserialize, Debug)]
#[serde(untagged)]
pub enum InnerHitsOneOrMore {
    Single(InnerHits),
    Multiple(Vec<InnerHits>),
}

#[derive(Serialize, Deserialize, Debug)]
pub struct InnerHits {
    from: Option<u32>,
    size: Option<u32>,
    sort: Option<Vec<SortElement>>,
    name: Option<String>,
}

#[derive(Debug)]
pub enum ParsedRequest {
    Malformed(String),
    ParseError(String),
    NotAllowed(String),
    NotSupported(String),
    NotFound(String),
    Allowed(Request<Body>),
}

pub fn parse_request(req: Request<Body>, config: &ProxyConfig) -> ParsedRequest {
    let (parts, body) = req.into_parts();

    // split path into at most 3 chunks
    let mut req_path = parts.uri.path();
    if req_path.starts_with("/") {
        req_path = &req_path[1..];
    }
    let path_chunks: Vec<&str> = req_path.split("/").collect();
    if path_chunks.len() > 3 {
        return ParsedRequest::NotSupported("only request paths with up to three segments allowed".to_string())
    }

    println!("{:?}", path_chunks);

    // this is sort of like a router
    match (&parts.method, path_chunks.as_slice()) {
        (&Method::GET, [""]) | (&Method::HEAD, [""]) => {
            parse_request_basic("", &parts, config)
        },
        (&Method::POST, ["_search", "scroll"]) | (&Method::DELETE, ["_search", "scroll"]) => {
            parse_request_scroll(None, &parts, body, config)
        },
        (&Method::POST, ["_search", "scroll", key]) | (&Method::DELETE, ["_search", "scroll", key]) => {
            parse_request_scroll(Some(key), &parts, body, config)
        },
        (&Method::GET, [index, "_search"]) | (&Method::POST, [index, "_search"]) => {
            parse_request_search(index, &parts, body, config)
        },
        //(Method::GET, [index, "_count"]) => {
        //    parse_request_count(index, "_count", None, &parts, body, config)
        //},
        (&Method::GET, [index, "_doc", key]) | (&Method::GET, [index, "_source", key]) => {
            parse_request_read(index, path_chunks[1], key, &parts, body, config)
        },
        _ => ParsedRequest::NotSupported("unknown endpoint".to_string()),
    }
}

pub fn parse_request_basic(endpoint: &str, parts: &request::Parts, config: &ProxyConfig) -> ParsedRequest {
    // XXX: partial
    let upstream_uri = Uri::builder()
        .scheme("http")
        .authority(config.upstream_addr.as_ref().unwrap_or(&"localhost:9200".to_string()).as_str())
        .path_and_query(format!("/{}", endpoint).as_str())
        .build()
        .unwrap();
    println!("{:?}", upstream_uri);
    let upstream_req = Request::builder()
        .uri(upstream_uri)
        .method(&parts.method)
        .body(Body::empty())
        .unwrap();
    ParsedRequest::Allowed(upstream_req)
}

pub fn parse_request_scroll(key: Option<&str>, parts: &request::Parts, body: Body, config: &ProxyConfig) -> ParsedRequest {
    ParsedRequest::NotSupported("not yet implemented".to_string())
}

pub fn parse_request_search(index: &str, parts: &request::Parts, body: Body, config: &ProxyConfig) -> ParsedRequest {
    ParsedRequest::NotSupported("not yet implemented".to_string())
}

pub fn parse_request_read(index: &str, endpoint: &str, key: &str, parts: &request::Parts, body: Body, config: &ProxyConfig) -> ParsedRequest {
    ParsedRequest::NotSupported("not yet implemented".to_string())
}