kong_rs 0.2.0

Plugin Development Kit for developing Kong plugins in Rust
Documentation
use http::HeaderMap;
use kong_rs_protos::{ExitArgs, Kv};
use prost_types::ListValue;
use strum::{EnumString, IntoStaticStr};

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

#[derive(Debug, PartialEq, IntoStaticStr, EnumString)]
pub(crate) enum Methods {
  #[strum(serialize = "kong.response.get_status")]
  GetStatus,
  #[strum(serialize = "kong.response.get_header")]
  GetHeader,
  #[strum(serialize = "kong.response.get_headers")]
  GetHeaders,
  #[strum(serialize = "kong.response.get_source")]
  GetSource,
  #[strum(serialize = "kong.response.set_status")]
  SetStatus,
  #[strum(serialize = "kong.response.set_header")]
  SetHeader,
  #[strum(serialize = "kong.response.add_header")]
  AddHeader,
  #[strum(serialize = "kong.response.clear_header")]
  ClearHeader,
  #[strum(serialize = "kong.response.set_headers")]
  SetHeaders,
  #[strum(serialize = "kong.response.exit")]
  Exit,
}

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

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

  pub async fn get_status(&self) -> KongResult<usize> {
    self.stream
      .ask_int(Methods::GetStatus.into())
      .await
      .map(|port| port as usize)
  }
  
  pub async fn get_header(&self, name: String) -> KongResult<String> {
    self.stream
      .ask_string_with_args(Methods::GetHeader.into(), &kong_rs_protos::String { v: name })
      .await
  }

  pub async fn get_headers(&self, max_headers: Option<usize>) -> KongResult<HeaderMap> {
    let max_headers = max_headers.unwrap_or(100);
    let headers: prost_types::Struct = self.stream.ask_message_with_args(
      Methods::GetHeaders.into(),
      &kong_rs_protos::Int { v: max_headers as i32 }
    ).await?;
    self.stream.unwrap_headers(headers)
  }

  pub async fn get_source(&self) -> KongResult<String> {
    self.stream.ask_string(Methods::GetStatus.into()).await
  }

  pub async fn set_status(&self, status: usize) -> KongResult<()> {
    self.stream.send_int(Methods::SetStatus.into(), status as i32).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 exit(&self, status: usize, body: Vec<u8>, headers: Option<HeaderMap>) -> KongResult<()> {
    let exit_args = ExitArgs { status: status as i32, body, headers: headers.map(Self::headers_to_struct) };
    self.stream.ask(Methods::Exit.into(), &exit_args).await
  }
}