use chrono::{DateTime, Utc};
use serde::Deserialize;
use std::path::PathBuf;
use std::process::Command;
use crate::auth::Role;
fn default_project() -> String {
"default".to_string()
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum CredentialSource {
Env { var: String },
File { path: PathBuf },
Command { command: String },
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct KeyPolicy {
pub id: String,
pub role: Role,
pub project: String,
pub benchmark_regex: Option<String>,
pub expires_at: Option<DateTime<Utc>>,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct LoadedCredential {
pub policy: KeyPolicy,
pub secret: String,
}
#[derive(Debug, thiserror::Error)]
pub enum CredentialSourceError {
#[error("missing environment variable {0}")]
MissingEnvVar(String),
#[error("failed to read key material from {source_name}: {message}")]
ReadFailure {
source_name: String,
message: String,
},
#[error("command source failed (status {status})")]
CommandFailure { status: i32 },
#[error("key material document parse failed: {0}")]
ParseFailure(String),
#[error("invalid key policy document: {0}")]
InvalidDocument(String),
}
#[derive(Debug, Clone, Deserialize)]
struct RawCredential {
id: String,
role: Role,
#[serde(default = "default_project")]
project: String,
#[serde(default)]
benchmark_regex: Option<String>,
#[serde(default)]
expires_at: Option<DateTime<Utc>>,
secret: String,
}
#[derive(Debug, Deserialize)]
struct JsonDocument {
#[serde(default)]
keys: Vec<RawCredential>,
#[serde(default)]
api_keys: Vec<RawCredential>,
}
#[derive(Debug, Deserialize)]
struct TomlDocument {
#[serde(default)]
keys: Vec<RawCredential>,
#[serde(default)]
api_keys: Vec<RawCredential>,
}
impl CredentialSource {
pub fn load(&self) -> Result<Vec<LoadedCredential>, CredentialSourceError> {
let content = match self {
Self::Env { var } => {
std::env::var(var).map_err(|_| CredentialSourceError::MissingEnvVar(var.clone()))?
}
Self::File { path } => {
std::fs::read_to_string(path).map_err(|e| CredentialSourceError::ReadFailure {
source_name: format!("file {}", path.display()),
message: e.to_string(),
})?
}
Self::Command { command } => run_command(command)?,
};
parse_credentials_document(&content)
}
}
fn run_command(command: &str) -> Result<String, CredentialSourceError> {
if command.trim().is_empty() {
return Err(CredentialSourceError::ReadFailure {
source_name: "command".to_string(),
message: "command is empty".to_string(),
});
}
let output =
shell_command(command)
.output()
.map_err(|e| CredentialSourceError::ReadFailure {
source_name: "command".to_string(),
message: e.to_string(),
})?;
if !output.status.success() {
let status = output.status.code().unwrap_or(-1);
return Err(CredentialSourceError::CommandFailure { status });
}
Ok(String::from_utf8_lossy(&output.stdout).to_string())
}
#[cfg(windows)]
fn shell_command(command: &str) -> Command {
let mut cmd = Command::new("powershell");
cmd.arg("-NoProfile").arg("-Command").arg(command);
cmd
}
#[cfg(not(windows))]
fn shell_command(command: &str) -> Command {
let mut cmd = Command::new("sh");
cmd.arg("-c").arg(command);
cmd
}
fn unwrap_wrapper_credentials(
keys: Vec<RawCredential>,
api_keys: Vec<RawCredential>,
) -> Result<Vec<RawCredential>, CredentialSourceError> {
match (keys.is_empty(), api_keys.is_empty()) {
(false, true) => Ok(keys),
(true, false) => Ok(api_keys),
(false, false) => Err(CredentialSourceError::InvalidDocument(
"document must not contain both 'keys' and 'api_keys'".to_string(),
)),
(true, true) => Err(CredentialSourceError::InvalidDocument(
"document must contain either 'keys' or 'api_keys'".to_string(),
)),
}
}
fn validate_loaded_credentials(
parsed: Vec<RawCredential>,
) -> Result<Vec<LoadedCredential>, CredentialSourceError> {
let mut out = Vec::with_capacity(parsed.len());
for entry in parsed {
if entry.id.trim().is_empty() {
return Err(CredentialSourceError::InvalidDocument(
"entry id must not be empty".to_string(),
));
}
if entry.secret.trim().is_empty() {
return Err(CredentialSourceError::InvalidDocument(format!(
"entry '{}' has empty secret",
entry.id
)));
}
out.push(LoadedCredential {
policy: KeyPolicy {
id: entry.id,
role: entry.role,
project: entry.project,
benchmark_regex: entry.benchmark_regex,
expires_at: entry.expires_at,
},
secret: entry.secret,
});
}
Ok(out)
}
pub fn parse_credentials_document(
raw: &str,
) -> Result<Vec<LoadedCredential>, CredentialSourceError> {
let trimmed = raw.trim();
if trimmed.is_empty() {
return Ok(Vec::new());
}
let mut parse_errors = Vec::new();
match serde_json::from_str::<Vec<RawCredential>>(trimmed) {
Ok(parsed) => return validate_loaded_credentials(parsed),
Err(err) => parse_errors.push(err.to_string()),
}
match serde_json::from_str::<JsonDocument>(trimmed) {
Ok(doc) => {
return validate_loaded_credentials(unwrap_wrapper_credentials(
doc.keys,
doc.api_keys,
)?);
}
Err(err) => parse_errors.push(err.to_string()),
}
match toml::from_str::<Vec<RawCredential>>(trimmed) {
Ok(parsed) => return validate_loaded_credentials(parsed),
Err(err) => parse_errors.push(err.to_string()),
}
match toml::from_str::<TomlDocument>(trimmed) {
Ok(doc) => {
return validate_loaded_credentials(unwrap_wrapper_credentials(
doc.keys,
doc.api_keys,
)?);
}
Err(err) => parse_errors.push(err.to_string()),
}
Err(CredentialSourceError::ParseFailure(parse_errors.join("; ")))
}
#[cfg(test)]
mod tests {
use super::*;
use std::path::Path;
#[cfg(windows)]
fn read_file_command(path: &Path) -> String {
let path = path.display().to_string().replace('\'', "''");
format!("Get-Content -Raw -LiteralPath '{}'", path)
}
#[cfg(not(windows))]
fn read_file_command(path: &Path) -> String {
format!("cat \"{}\"", path.display())
}
#[cfg(windows)]
fn failing_command() -> String {
"Write-Error 'boom'; exit 9".to_string()
}
#[cfg(not(windows))]
fn failing_command() -> String {
"echo boom >&2; exit 9".to_string()
}
#[test]
fn parse_json_credentials() {
let doc = r#"[
{
"id":"ci-promoter",
"role":"promoter",
"project":"my-project",
"benchmark_regex":".*",
"secret":"pg_live_abcdefghijklmnopqrstuvwxyz123456"
}
]"#;
let creds = parse_credentials_document(doc).unwrap();
assert_eq!(creds.len(), 1);
assert_eq!(creds[0].policy.id, "ci-promoter");
assert_eq!(creds[0].policy.role, Role::Promoter);
}
#[test]
fn parse_toml_credentials_table() {
let doc = r#"
[[keys]]
id = "dev"
role = "admin"
project = "default"
secret = "pg_test_abcdefghijklmnopqrstuvwxyz123456"
"#;
let creds = parse_credentials_document(doc).unwrap();
assert_eq!(creds.len(), 1);
assert_eq!(creds[0].policy.role, Role::Admin);
}
#[test]
fn parse_json_credentials_object_wrapper() {
let doc = r#"{
"keys": [
{
"id":"ci-promoter",
"role":"promoter",
"project":"my-project",
"secret":"pg_live_abcdefghijklmnopqrstuvwxyz123456"
}
]
}"#;
let creds = parse_credentials_document(doc).unwrap();
assert_eq!(creds.len(), 1);
assert_eq!(creds[0].policy.id, "ci-promoter");
assert_eq!(creds[0].policy.role, Role::Promoter);
}
#[test]
fn parse_json_credentials_wrapper_rejects_missing_expected_fields() {
let doc = r#"{
"apiKeys": [
{
"id":"ci-promoter",
"role":"promoter",
"project":"my-project",
"secret":"pg_live_abcdefghijklmnopqrstuvwxyz123456"
}
]
}"#;
let err = parse_credentials_document(doc).unwrap_err().to_string();
assert!(err.contains("document must contain either 'keys' or 'api_keys'"));
}
#[test]
fn parse_toml_credentials_wrapper_rejects_duplicate_wrappers() {
let doc = r#"
[[keys]]
id = "dev"
role = "admin"
project = "default"
secret = "pg_test_abcdefghijklmnopqrstuvwxyz123456"
[[api_keys]]
id = "ci"
role = "viewer"
project = "default"
secret = "pg_test_abcdefghijklmnopqrstuvwxyz123456"
"#;
let err = parse_credentials_document(doc).unwrap_err().to_string();
assert!(err.contains("document must not contain both 'keys' and 'api_keys'"));
}
#[test]
fn command_source_success() {
let temp_dir = tempfile::tempdir().unwrap();
let path = temp_dir.path().join("keys.json");
std::fs::write(
&path,
"[{\"id\":\"k1\",\"role\":\"viewer\",\"project\":\"p\",\"secret\":\"pg_test_abcdefghijklmnopqrstuvwxyz123456\"}]",
)
.unwrap();
let src = CredentialSource::Command {
command: read_file_command(&path),
};
let creds = src.load().unwrap();
assert_eq!(creds.len(), 1);
assert_eq!(creds[0].policy.id, "k1");
}
#[test]
fn command_source_failure_hides_stderr() {
let src = CredentialSource::Command {
command: failing_command(),
};
let err = src.load().unwrap_err().to_string();
assert!(err.contains("status 9"));
assert!(!err.contains("boom"));
}
#[test]
fn file_source_loads_credentials() {
let temp = tempfile::NamedTempFile::new().unwrap();
std::fs::write(
temp.path(),
"[{\"id\":\"f1\",\"role\":\"viewer\",\"project\":\"x\",\"secret\":\"pg_test_abcdefghijklmnopqrstuvwxyz123456\"}]",
)
.unwrap();
let src = CredentialSource::File {
path: temp.path().to_path_buf(),
};
let creds = src.load().unwrap();
assert_eq!(creds[0].policy.id, "f1");
}
#[test]
#[allow(unsafe_code)]
fn env_source_loads_credentials() {
let var = "PERFGATE_AUTH_TEST_KEYS";
unsafe {
std::env::set_var(
var,
"[{\"id\":\"e1\",\"role\":\"viewer\",\"project\":\"x\",\"secret\":\"pg_test_abcdefghijklmnopqrstuvwxyz123456\"}]",
);
}
let src = CredentialSource::Env {
var: var.to_string(),
};
let creds = src.load().unwrap();
assert_eq!(creds[0].policy.id, "e1");
}
}