use base64::Engine;
use flate2::read::ZlibDecoder;
use flate2::write::ZlibEncoder;
use flate2::Compression;
use serde_json::{json, Value};
use std::io::{Read, Write};
pub fn encode_deny(
explicit: bool,
action: Option<&str>,
principal_arn: Option<&str>,
matched_statements: Vec<Value>,
context: Option<Value>,
) -> String {
let mut payload = json!({
"allowed": false,
"explicitDeny": explicit,
"matchedStatements": { "items": matched_statements },
});
if let Some(a) = action {
payload["action"] = json!(a);
}
if let Some(p) = principal_arn {
payload["principal"] = json!(p);
}
if let Some(c) = context {
payload["context"] = c;
}
let json_bytes = serde_json::to_vec(&payload).unwrap_or_default();
let mut encoder = ZlibEncoder::new(Vec::new(), Compression::default());
encoder.write_all(&json_bytes).ok();
let compressed = encoder.finish().unwrap_or_default();
base64::engine::general_purpose::STANDARD.encode(compressed)
}
pub fn decode_message(encoded: &str) -> Result<String, &'static str> {
let compressed = base64::engine::general_purpose::STANDARD
.decode(encoded.as_bytes())
.map_err(|_| "EncodedMessage is not valid base64")?;
let mut decoder = ZlibDecoder::new(&compressed[..]);
let mut json_bytes = Vec::new();
decoder
.read_to_end(&mut json_bytes)
.map_err(|_| "EncodedMessage is not a valid deny token")?;
String::from_utf8(json_bytes).map_err(|_| "EncodedMessage payload is not valid UTF-8")
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn round_trip_explicit_deny() {
let token = encode_deny(
true,
Some("s3:GetObject"),
Some("arn:aws:iam::111122223333:user/alice"),
vec![json!({"sourcePolicyId": "PolicyA"})],
Some(json!({"aws:SourceIp": "1.2.3.4"})),
);
let decoded = decode_message(&token).unwrap();
let parsed: Value = serde_json::from_str(&decoded).unwrap();
assert_eq!(parsed["allowed"], false);
assert_eq!(parsed["explicitDeny"], true);
assert_eq!(parsed["action"], "s3:GetObject");
assert_eq!(parsed["principal"], "arn:aws:iam::111122223333:user/alice");
assert_eq!(
parsed["matchedStatements"]["items"][0]["sourcePolicyId"],
"PolicyA"
);
}
#[test]
fn round_trip_implicit_deny_with_no_extras() {
let token = encode_deny(false, None, None, Vec::new(), None);
let decoded = decode_message(&token).unwrap();
let parsed: Value = serde_json::from_str(&decoded).unwrap();
assert_eq!(parsed["allowed"], false);
assert_eq!(parsed["explicitDeny"], false);
assert!(parsed["matchedStatements"]["items"]
.as_array()
.unwrap()
.is_empty());
assert!(parsed.get("action").is_none());
}
#[test]
fn rejects_garbage_base64() {
assert!(decode_message("not!!!base64!!").is_err());
}
#[test]
fn rejects_base64_that_is_not_zlib() {
let token = base64::engine::general_purpose::STANDARD.encode(b"not zlib data");
assert!(decode_message(&token).is_err());
}
}