Collect label and metadata from header (#79)

trueleo · web-flow · commit c48c1638fc83 · 2022-08-30T21:47:28.000+05:30
This PR allows for metadata and tags to be fetched from headers instead of body of the event. This gives some flexibility in schema and multiple sources with different schema for available metadata can send logs to the same stream. Client needs to use prefix x-p-meta- and x-p-tags- in headers to send key value pairs in a server post event. These headers are parsed and validated before being merged into body with p_metadata and p_tags field respectively Changes: - introduced header_parsing::collect_labelled which parses headers key value pairs and gives a string formatted as "k1=v1,k2=v2" - New Error type ParseHeaderError which implements ResponseError. - Use above types in post event handler Fixes #75
diff --git a/server/src/handlers/event.rs b/server/src/handlers/event.rs
@@ -16,8 +16,10 @@
  *
  */
 
+use std::collections::HashMap;
+
 use actix_web::http::StatusCode;
-use actix_web::{web, HttpRequest, HttpResponse};
+use actix_web::{web, HttpRequest, HttpResponse, ResponseError};
 use serde_json::Value;
 
 use crate::event;
@@ -26,7 +28,12 @@ use crate::query::Query;
 use crate::response::{self, EventResponse};
 use crate::s3::S3;
 use crate::storage::ObjectStorage;
-use crate::utils;
+use crate::utils::header_parsing::collect_labelled_headers;
+use crate::utils::{self, merge};
+
+const PREFIX_TAGS: &str = "x-p-tags-";
+const PREFIX_META: &str = "x-p-meta-";
+const SEPARATOR: char = ',';
 
 pub async fn query(_req: HttpRequest, json: web::Json<Value>) -> HttpResponse {
     let json = json.into_inner();
@@ -74,7 +81,16 @@ pub async fn query(_req: HttpRequest, json: web::Json<Value>) -> HttpResponse {
 
 pub async fn post_event(req: HttpRequest, body: web::Json<serde_json::Value>) -> HttpResponse {
     let stream_name: String = req.match_info().get("logstream").unwrap().parse().unwrap();
-    let labels = utils::collect_labels(&req);
+
+    let tags = match collect_labelled_headers(&req, PREFIX_TAGS, SEPARATOR) {
+        Ok(tags) => HashMap::from([("p_tags".to_string(), tags)]),
+        Err(e) => return e.error_response(),
+    };
+
+    let metadata = match collect_labelled_headers(&req, PREFIX_META, SEPARATOR) {
+        Ok(metadata) => HashMap::from([("p_metadata".to_string(), metadata)]),
+        Err(e) => return e.error_response(),
+    };
 
     if let Err(e) = metadata::STREAM_INFO.schema(&stream_name) {
         // if stream doesn't exist, fail to post data
@@ -94,7 +110,10 @@ pub async fn post_event(req: HttpRequest, body: web::Json<serde_json::Value>) ->
         let mut i = 0;
 
         for body in array {
-            let body = utils::flatten_json_body(web::Json(body.clone()), labels.clone()).unwrap();
+            let body = merge(body.clone(), metadata.clone());
+            let body = merge(body, tags.clone());
+            let body = utils::flatten_json_body(web::Json(body)).unwrap();
+
             let e = event::Event {
                 body,
                 stream_name: stream_name.clone(),
@@ -118,8 +137,11 @@ pub async fn post_event(req: HttpRequest, body: web::Json<serde_json::Value>) ->
         .to_http();
     }
 
+    let body = merge(body.clone(), metadata);
+    let body = merge(body, tags);
+
     let event = event::Event {
-        body: utils::flatten_json_body(body, labels).unwrap(),
+        body: utils::flatten_json_body(web::Json(body)).unwrap(),
         stream_name,
     };
 
diff --git a/server/src/utils.rs b/server/src/utils.rs
@@ -17,52 +17,104 @@
  */
 
 use actix_web::web;
-use actix_web::HttpRequest;
 use chrono::{Date, DateTime, Timelike, Utc};
 use serde_json::{json, Value};
 use std::collections::HashMap;
 
 use crate::Error;
 
-const META_LABEL: &str = "x-p-meta";
-
-pub fn flatten_json_body(
-    body: web::Json<serde_json::Value>,
-    labels: Option<String>,
-) -> Result<String, Error> {
-    let mut collector_labels = HashMap::new();
-
-    collector_labels.insert("labels".to_string(), labels.unwrap());
-
+pub fn flatten_json_body(body: web::Json<serde_json::Value>) -> Result<String, Error> {
     let mut flat_value: Value = json!({});
-    let new_body = merge(&body, &collector_labels);
-    flatten_json::flatten(&new_body, &mut flat_value, None, true, Some("_")).unwrap();
+    flatten_json::flatten(&body, &mut flat_value, None, true, Some("_")).unwrap();
     let flattened = serde_json::to_string(&flat_value)?;
 
     Ok(flattened)
 }
 
-fn merge(v: &Value, fields: &HashMap<String, String>) -> Value {
-    match v {
-        Value::Object(m) => {
-            let mut m = m.clone();
+pub fn merge(value: Value, fields: HashMap<String, String>) -> Value {
+    match value {
+        Value::Object(mut m) => {
             for (k, v) in fields {
-                match m.get(k) {
-                    Some(curr_val) => {
+                match m.get_mut(&k) {
+                    Some(val) => {
                         let mut final_val = String::new();
-                        final_val.push_str(curr_val.as_str().unwrap());
+                        final_val.push_str(val.as_str().unwrap());
                         final_val.push(',');
-                        final_val.push_str(v);
-                        m.insert(k.clone(), Value::String(final_val));
+                        final_val.push_str(&v);
+                        *val = Value::String(final_val);
                     }
                     None => {
-                        m.insert(k.clone(), Value::String(v.to_string()));
+                        m.insert(k, Value::String(v));
                     }
                 }
             }
             Value::Object(m)
         }
-        v => v.clone(),
+        value => value,
+    }
+}
+
+pub mod header_parsing {
+    const MAX_HEADERS_ALLOWED: usize = 5;
+    use actix_web::{HttpRequest, HttpResponse, ResponseError};
+
+    pub fn collect_labelled_headers(
+        req: &HttpRequest,
+        prefix: &str,
+        kv_separator: char,
+    ) -> Result<String, ParseHeaderError> {
+        // filter out headers which has right prefix label and convert them into str;
+        let headers = req.headers().iter().filter_map(|(key, value)| {
+            let key = key.as_str().strip_prefix(prefix)?;
+            Some((key, value))
+        });
+
+        let mut labels: Vec<String> = Vec::new();
+
+        for (key, value) in headers {
+            let value = value.to_str().map_err(|_| ParseHeaderError::InvalidValue)?;
+            if key.is_empty() {
+                return Err(ParseHeaderError::Emptykey);
+            }
+            if key.contains(kv_separator) {
+                return Err(ParseHeaderError::SeperatorInKey(kv_separator));
+            }
+            if value.contains(kv_separator) {
+                return Err(ParseHeaderError::SeperatorInValue(kv_separator));
+            }
+
+            labels.push(format!("{}={}", key, value));
+        }
+
+        if labels.len() > MAX_HEADERS_ALLOWED {
+            return Err(ParseHeaderError::MaxHeadersLimitExceeded);
+        }
+
+        Ok(labels.join(","))
+    }
+
+    #[derive(Debug, thiserror::Error)]
+    pub enum ParseHeaderError {
+        #[error("Too many headers received. Limit is of 5 headers")]
+        MaxHeadersLimitExceeded,
+        #[error("A value passed in header is not formattable to plain visible ASCII")]
+        InvalidValue,
+        #[error("Invalid Key was passed which terminated just after the end of prefix")]
+        Emptykey,
+        #[error("A key passed in header contains reserved char {0}")]
+        SeperatorInKey(char),
+        #[error("A value passed in header contains reserved char {0}")]
+        SeperatorInValue(char),
+    }
+
+    impl ResponseError for ParseHeaderError {
+        fn status_code(&self) -> http::StatusCode {
+            http::StatusCode::BAD_REQUEST
+        }
+
+        fn error_response(&self) -> HttpResponse {
+            HttpResponse::build(self.status_code()).body(self.to_string())
+        }
     }
 }
 
@@ -110,25 +162,6 @@ pub fn minute_to_prefix(minute: u32, data_granularity: u32) -> Option<String> {
     ))
 }
 
-/// collect labels passed from http headers
-/// format: labels = "app=k8s, cloud=gcp"
-pub fn collect_labels(req: &HttpRequest) -> Option<String> {
-    let keys = req.headers().keys().cloned().collect::<Vec<_>>();
-
-    let mut labels_vec = Vec::with_capacity(100);
-    for key in keys {
-        if key.to_string().to_lowercase().starts_with(META_LABEL) {
-            let value = req.headers().get(&key)?.to_str().ok();
-            let remove_meta_char = format!("{}-", META_LABEL);
-            let kv =
-                format! {"{}={}", key.to_string().replace(&remove_meta_char, ""), value.unwrap()};
-            labels_vec.push(kv);
-        }
-    }
-
-    Some(labels_vec.join(","))
-}
-
 pub struct TimePeriod {
     start: DateTime<Utc>,
     end: DateTime<Utc>,
