path: root/src/consul.rs

use std::collections::HashMap;

use anyhow::{anyhow, Result};
use serde::{Deserialize, Serialize};

use crate::config::RuntimeConfigConsul;

#[derive(Serialize, Deserialize, Debug)]
pub struct ServiceEntry {
  #[serde(rename = "Tags")]
  pub tags: Vec<String>,
}

#[derive(Serialize, Deserialize, Debug, Default)]
pub struct CatalogNode {
  #[serde(rename = "Services")]
  pub services: HashMap<String, ServiceEntry>,
}

pub struct Consul {
  client: reqwest::Client,
  url: String,
  idx: Option<u64>,
}

impl Consul {
  pub fn new(config: &RuntimeConfigConsul) -> Self {
    let client = if let Some((ca, skip_verify, ident)) = config.tls.clone() {
      if skip_verify {
        reqwest::Client::builder()
          .use_rustls_tls()
          .danger_accept_invalid_certs(true)
          .identity(ident)
          .build()
          .expect("Unable to build reqwest client")
      } else if let Some(ca) = ca {
        reqwest::Client::builder()
          .use_rustls_tls()
          .add_root_certificate(ca)
          .identity(ident)
          .build()
          .expect("Unable to build reqwest client")
      } else {
        reqwest::Client::builder()
          .use_rustls_tls()
          .identity(ident)
          .build()
          .expect("Unable to build reqwest client")
      }
    } else {
      reqwest::Client::new()
    };
    return Self {
      client,
      url: config.url.clone(),
      idx: None,
    };
  }

  pub fn watch_node_reset(&mut self) -> () {
    self.idx = None;
  }

  pub async fn watch_node(&mut self, host: &str) -> Result<CatalogNode> {
    let url = match self.idx {
      Some(i) => format!("{}/v1/catalog/node/{}?index={}", self.url, host, i),
      None => format!("{}/v1/catalog/node/{}", self.url, host),
    };

    let http = self.client.get(&url).send().await?;
    self.idx = match http.headers().get("X-Consul-Index") {
      Some(v) => Some(v.to_str()?.parse::<u64>()?),
      None => return Err(anyhow!("X-Consul-Index header not found")),
    };

    let resp: Option<CatalogNode> = http.json().await?;
    return Ok(resp.unwrap_or_default());
  }

  pub async fn kv_put(&self, key: &str, bytes: Vec<u8>) -> Result<()> {
    let url = format!("{}/v1/kv/{}", self.url, key);
    let http = self.client.put(&url).body(bytes).send().await?;
    http.error_for_status()?;
    Ok(())
  }
}