kong_rs 0.2.0

Plugin Development Kit for developing Kong plugins in Rust
Documentation
use std::collections::BTreeMap;

use http::HeaderMap;
use kong_rs_protos::Kv;
use prost_types::ListValue;
use strum::{EnumString, IntoStaticStr};

use crate::{pdk::Value, stream::Stream, KongResult};

#[derive(Debug, PartialEq, IntoStaticStr, EnumString)]
pub(crate) enum Methods {
  #[strum(serialize = "kong.service.request.set_scheme")]
  SetScheme,
  #[strum(serialize = "kong.service.request.set_path")]
  SetPath,
  #[strum(serialize = "kong.service.request.set_raw_query")]
  SetRawQuery,
  #[strum(serialize = "kong.service.request.set_method")]
  SetMethod,
  #[strum(serialize = "kong.service.request.set_query")]
  SetQuery,
  #[strum(serialize = "kong.service.request.set_header")]
  SetHeader,
  #[strum(serialize = "kong.service.request.add_header")]
  AddHeader,
  #[strum(serialize = "kong.service.request.clear_header")]
  ClearHeader,
  #[strum(serialize = "kong.service.request.set_headers")]
  SetHeaders,
  #[strum(serialize = "kong.service.request.set_raw_body")]
  SetRawBody,
}

#[derive(Clone)]
pub struct ServiceRequestPDK {
  stream: Stream,
}

impl ServiceRequestPDK {
  pub fn new(stream: Stream) -> Self {
    Self { stream }
  }

  pub async fn set_scheme<S: Into<String>>(&self, scheme: S) -> KongResult<()> {
    self.stream.send_string(Methods::SetScheme.into(), scheme.into()).await
  }

  pub async fn set_path<S: Into<String>>(&self, path: S) -> KongResult<()> {
    self.stream.send_string(Methods::SetPath.into(), path.into()).await
  }

  pub async fn set_raw_query<S: Into<String>>(&self, query: S) -> KongResult<()> {
    self.stream.send_string(Methods::SetRawQuery.into(), query.into()).await
  }

  pub async fn set_method<S: Into<String>>(&self, method: S) -> KongResult<()> {
    self.stream.send_string(Methods::SetMethod.into(), method.into()).await
  }

  pub async fn set_query<S: Into<String>>(&self, query: BTreeMap<String, Value>) -> KongResult<()> {
    self.stream.ask_message_with_args(Methods::SetQuery.into(), &prost_types::Struct {
      fields: query.into_iter().map(|(k, v)| (k, prost_types::Value { kind: Some(v.into()) })).collect()
    }).await
  }

  pub async fn set_header(&self, name: &str, value: &str) -> KongResult<()> {
    self.stream.ask(Methods::SetHeader.into(), &Kv {
      k: name.to_owned(),
      v: Some(prost_types::Value { kind: Some(prost_types::value::Kind::StringValue(value.to_owned())) })
    }).await
  }

  pub async fn add_header(&self, name: &str, value: &str) -> KongResult<()> {
    self.stream.ask(Methods::AddHeader.into(), &Kv {
      k: name.to_owned(),
      v: Some(prost_types::Value { kind: Some(prost_types::value::Kind::StringValue(value.to_owned())) })
    }).await
  }

  pub async fn clear_header(&self, name: &str) -> KongResult<()> {
    self.stream.ask(Methods::ClearHeader.into(), &kong_rs_protos::String { v: name.to_owned() }).await
  }

  fn headers_to_struct(headers: HeaderMap) -> prost_types::Struct {
    let mut s = prost_types::Struct { ..Default::default() };

    for key in headers.keys() {
      let values = headers.get_all(key);
      s.fields.insert(
        key.to_string(),
        prost_types::Value {
          kind: Some(prost_types::value::Kind::ListValue(ListValue {
            values: values.into_iter().map(|x| prost_types::Value {
              kind: Some(prost_types::value::Kind::StringValue(std::str::from_utf8(x.as_bytes()).unwrap_or("").to_owned()))
            }).collect()
          }))
        }
      );
    }

    s
  }

  pub async fn set_headers(&self, headers: HeaderMap) -> KongResult<()> {
    let s = Self::headers_to_struct(headers);
    self.stream.ask(Methods::SetHeaders.into(), &s).await
  }

  pub async fn set_body(&self, body: Vec<u8>) -> KongResult<()> {
    let bs = kong_rs_protos::ByteString { v: body };
    self.stream.ask(Methods::SetRawBody.into(), &bs).await
  }
}