feat: Support for custom headers to be added to every request to each server to support alternative authentication mechanisms [#47]

2025-08-29 10:23:17 -06:00
parent 4ea39f74fe
commit 0b29351366
5 changed files with 235 additions and 5 deletions
@@ -2,6 +2,8 @@
 mod tests {
  use anyhow::anyhow;
  use pretty_assertions::{assert_eq, assert_str_eq};
+  use reqwest::header::{HeaderMap, HeaderName, HeaderValue};
+  use serde_json::Value;
  use serial_test::serial;
  use tokio::sync::mpsc;

@@ -340,6 +342,43 @@ mod tests {
    assert_eq!(servarr_config.api_token, Some(String::new()));
    assert_eq!(servarr_config.api_token_file, None);
    assert_eq!(servarr_config.ssl_cert_path, None);
+    assert_eq!(servarr_config.custom_headers, None);
+  }
+
+  #[test]
+  fn serialize_header_map_basic() {
+    let mut header_map = HeaderMap::new();
+    header_map.insert(
+      HeaderName::from_static("x-api-key"),
+      HeaderValue::from_static("abc123"),
+    );
+    header_map.insert(
+      HeaderName::from_static("header-1"),
+      HeaderValue::from_static("test"),
+    );
+
+    let config = ServarrConfig {
+      custom_headers: Some(header_map),
+      ..ServarrConfig::default()
+    };
+
+    let v: Value = serde_json::to_value(&config).expect("serialize ok");
+    let custom = v.get("custom_headers").unwrap();
+    assert!(custom.is_object());
+    let obj = custom.as_object().unwrap();
+
+    assert_eq!(obj.get("x-api-key").unwrap(), "abc123");
+    assert_eq!(obj.get("header-1").unwrap(), "test");
+
+    assert!(obj.get("X-Api-Key").is_none());
+    assert!(obj.get("HEADER-1").is_none());
+  }
+
+  #[test]
+  fn serialize_header_map_none_is_null() {
+    let config = ServarrConfig::default();
+    let v: Value = serde_json::to_value(&config).expect("serialize ok");
+    assert!(v.get("custom_headers").unwrap().is_null());
  }

  #[test]
@@ -383,6 +422,66 @@ mod tests {
    assert_eq!(config.port, None);
  }

+  #[test]
+  #[serial]
+  fn test_deserialize_optional_env_var_header_map_is_present() {
+    unsafe { std::env::set_var("TEST_VAR_DESERIALIZE_HEADER_OPTION", "localhost") };
+    let expected_custom_headers = {
+      let mut headers = HeaderMap::new();
+      headers.insert("X-Api-Host", "localhost".parse().unwrap());
+      headers.insert("api-token", "test123".parse().unwrap());
+      headers
+    };
+    let yaml_data = r#"
+      custom_headers:
+        X-Api-Host: ${TEST_VAR_DESERIALIZE_HEADER_OPTION}
+        api-token: "test123"
+    "#;
+
+    let config: ServarrConfig = serde_yaml::from_str(yaml_data).unwrap();
+
+    assert_eq!(config.custom_headers, Some(expected_custom_headers));
+    unsafe { std::env::remove_var("TEST_VAR_DESERIALIZE_HEADER_OPTION") };
+  }
+
+  #[test]
+  #[serial]
+  fn test_deserialize_optional_env_var_header_map_does_not_overwrite_non_env_value() {
+    unsafe {
+      std::env::set_var(
+        "TEST_VAR_DESERIALIZE_HEADER_OPTION_NO_OVERWRITE",
+        "localhost",
+      )
+    };
+    let expected_custom_headers = {
+      let mut headers = HeaderMap::new();
+      headers.insert("X-Api-Host", "www.example.com".parse().unwrap());
+      headers.insert("api-token", "test123".parse().unwrap());
+      headers
+    };
+    let yaml_data = r#"
+      custom_headers:
+        X-Api-Host: www.example.com
+        api-token: "test123"
+    "#;
+
+    let config: ServarrConfig = serde_yaml::from_str(yaml_data).unwrap();
+
+    assert_eq!(config.custom_headers, Some(expected_custom_headers));
+    unsafe { std::env::remove_var("TEST_VAR_DESERIALIZE_HEADER_OPTION_NO_OVERWRITE") };
+  }
+
+  #[test]
+  fn test_deserialize_optional_env_var_header_map_empty() {
+    let yaml_data = r#"
+      api_token: "test123"
+    "#;
+
+    let config: ServarrConfig = serde_yaml::from_str(yaml_data).unwrap();
+
+    assert_eq!(config.custom_headers, None);
+  }
+
  #[test]
  #[serial]
  fn test_deserialize_optional_u16_env_var_is_present() {
@@ -496,7 +595,9 @@ mod tests {
    let api_token = "thisisatest".to_owned();
    let api_token_file = "/root/.config/api_token".to_owned();
    let ssl_cert_path = "/some/path".to_owned();
-    let expected_str = format!("ServarrConfig {{ name: Some(\"{name}\"), host: Some(\"{host}\"), port: Some({port}), uri: Some(\"{uri}\"), weight: Some({weight}), api_token: Some(\"***********\"), api_token_file: Some(\"{api_token_file}\"), ssl_cert_path: Some(\"{ssl_cert_path}\") }}");
+    let mut custom_headers = HeaderMap::new();
+    custom_headers.insert("X-Custom-Header", "value".parse().unwrap());
+    let expected_str = format!("ServarrConfig {{ name: Some(\"{name}\"), host: Some(\"{host}\"), port: Some({port}), uri: Some(\"{uri}\"), weight: Some({weight}), api_token: Some(\"***********\"), api_token_file: Some(\"{api_token_file}\"), ssl_cert_path: Some(\"{ssl_cert_path}\"), custom_headers: Some({{\"x-custom-header\": \"value\"}}) }}");
    let servarr_config = ServarrConfig {
      name: Some(name),
      host: Some(host),
@@ -506,6 +607,7 @@ mod tests {
      api_token: Some(api_token),
      api_token_file: Some(api_token_file),
      ssl_cert_path: Some(ssl_cert_path),
+      custom_headers: Some(custom_headers),
    };

    assert_str_eq!(format!("{servarr_config:?}"), expected_str);
@@ -3,7 +3,9 @@ use colored::Colorize;
 use itertools::Itertools;
 use log::{debug, error};
 use regex::Regex;
+use reqwest::header::{HeaderMap, HeaderName, HeaderValue};
 use serde::{Deserialize, Serialize};
+use std::collections::HashMap;
 use std::path::PathBuf;
 use std::{fs, process};
 use tokio::sync::mpsc::Sender;
@@ -335,6 +337,12 @@ pub struct ServarrConfig {
  pub api_token_file: Option<String>,
  #[serde(default, deserialize_with = "deserialize_optional_env_var")]
  pub ssl_cert_path: Option<String>,
+  #[serde(
+    default,
+    deserialize_with = "deserialize_optional_env_var_header_map",
+    serialize_with = "serialize_header_map"
+  )]
+  pub custom_headers: Option<HeaderMap>,
 }

 impl ServarrConfig {
@@ -380,6 +388,7 @@ impl Default for ServarrConfig {
      api_token: Some(String::new()),
      api_token_file: None,
      ssl_cert_path: None,
+      custom_headers: None,
    }
  }
 }
@@ -389,6 +398,27 @@ pub fn log_and_print_error(error: String) {
  eprintln!("error: {}", error.red());
 }

+fn serialize_header_map<S>(headers: &Option<HeaderMap>, serializer: S) -> Result<S::Ok, S::Error>
+where
+  S: serde::Serializer,
+{
+  if let Some(headers) = headers {
+    let mut map = HashMap::new();
+    for (name, value) in headers.iter() {
+      let name_str = name.as_str().to_string();
+      let value_str = value
+        .to_str()
+        .map_err(serde::ser::Error::custom)?
+        .to_string();
+
+      map.insert(name_str, value_str);
+    }
+    map.serialize(serializer)
+  } else {
+    serializer.serialize_none()
+  }
+}
+
 fn deserialize_optional_env_var<'de, D>(deserializer: D) -> Result<Option<String>, D::Error>
 where
  D: serde::Deserializer<'de>,
@@ -403,6 +433,28 @@ where
  }
 }

+fn deserialize_optional_env_var_header_map<'de, D>(
+  deserializer: D,
+) -> Result<Option<HeaderMap>, D::Error>
+where
+  D: serde::Deserializer<'de>,
+{
+  let opt: Option<HashMap<String, String>> = Option::deserialize(deserializer)?;
+  match opt {
+    Some(map) => {
+      let mut header_map = HeaderMap::new();
+      for (k, v) in map.iter() {
+        let name = HeaderName::from_bytes(k.as_bytes()).map_err(serde::de::Error::custom)?;
+        let value_str = interpolate_env_vars(v);
+        let value = HeaderValue::from_str(&value_str).map_err(serde::de::Error::custom)?;
+        header_map.insert(name, value);
+      }
+      Ok(Some(header_map))
+    }
+    None => Ok(None),
+  }
+}
+
 fn deserialize_u16_env_var<'de, D>(deserializer: D) -> Result<Option<u16>, D::Error>
 where
  D: serde::Deserializer<'de>,