use core::fmt;
use core::str::FromStr;
use std::collections::{HashMap, HashSet};
use chrono::{DateTime, Utc};
use indexmap::IndexMap;
use serde::{Deserialize, Deserializer, Serialize, Serializer};
use serde_json::{Map, Value};
use super::{ClaimPath, Disclosure, HashAlgorithm, SdJwtBuilder, SdJwtError};
use crate::jose::{JwsHeader, Signer};
use crate::vc::{Jwt, JwtClaims};
pub const SD_KEY: &str = "_sd";
pub const ELLIPSIS_KEY: &str = "...";
pub const KB_JWT_TYP: &str = "kb+jwt";
#[derive(Clone, PartialEq, Eq)]
pub struct SdJwt {
jwt: Jwt,
disclosures: Vec<Disclosure>,
key_binding: Option<Jwt>,
}
impl SdJwt {
pub fn new(jwt: Jwt, disclosures: Vec<Disclosure>) -> Self {
Self {
jwt,
disclosures,
key_binding: None,
}
}
pub fn parse(s: &str) -> Result<Self, SdJwtError> {
let parts: Vec<&str> = s.split('~').collect();
let [jwt, middle @ .., last] = parts.as_slice() else {
return Err(SdJwtError::Format(
"no `~` after the issuer-signed JWT".into(),
));
};
let jwt = Jwt::new(*jwt)?;
let key_binding = if last.is_empty() {
None
} else {
Some(Jwt::new(*last)?)
};
let mut disclosures = Vec::with_capacity(middle.len());
let mut seen = HashSet::new();
for part in middle {
if part.is_empty() {
return Err(SdJwtError::Format("empty disclosure".into()));
}
if !seen.insert(*part) {
return Err(SdJwtError::DuplicateDisclosure((*part).to_owned()));
}
disclosures.push(Disclosure::parse(part)?);
}
Ok(Self {
jwt,
disclosures,
key_binding,
})
}
pub fn jwt(&self) -> &Jwt {
&self.jwt
}
pub fn disclosures(&self) -> &[Disclosure] {
&self.disclosures
}
pub fn key_binding_jwt(&self) -> Option<&Jwt> {
self.key_binding.as_ref()
}
#[must_use]
pub fn with_key_binding(mut self, key_binding: Jwt) -> Self {
self.key_binding = Some(key_binding);
self
}
#[must_use]
pub fn without_key_binding(mut self) -> Self {
self.key_binding = None;
self
}
pub fn payload(&self) -> Result<Map<String, Value>, SdJwtError> {
let jws = self.jwt.decode()?;
serde_json::from_slice(jws.payload())
.map_err(|e| SdJwtError::Payload(format!("not a JSON object: {e}")))
}
pub fn hash_algorithm(&self) -> Result<HashAlgorithm, SdJwtError> {
hash_algorithm(&self.payload()?)
}
pub fn to_string_without_key_binding(&self) -> String {
let mut s = String::from(self.jwt.as_str());
s.push('~');
for disclosure in &self.disclosures {
s.push_str(disclosure.as_str());
s.push('~');
}
s
}
pub fn sd_hash(&self, alg: HashAlgorithm) -> String {
alg.digest(self.to_string_without_key_binding().as_bytes())
}
pub fn process(&self) -> Result<Processed, SdJwtError> {
let mut payload = self.payload()?;
let alg = hash_algorithm(&payload)?;
payload.shift_remove(HashAlgorithm::CLAIM);
let mut processor = Processor::new(alg, &self.disclosures);
let issued = processor.object(payload, &ClaimPath::root(), None, 1)?;
if let Some(unused) = processor.used.iter().position(|used| !used) {
return Err(SdJwtError::UnusedDisclosure(
self.disclosures[unused].as_str().to_owned(),
));
}
let mut payload = Value::Object(issued.clone());
strip_placeholders(&mut payload);
let Value::Object(payload) = payload else {
unreachable!("an object stays an object")
};
Ok(Processed {
payload,
disclosed: processor.disclosed,
undisclosed: processor.undisclosed,
issued,
parents: processor.parents,
})
}
pub fn select(&self, mut keep: impl FnMut(&Disclosure) -> bool) -> Result<Self, SdJwtError> {
let processed = self.process()?;
let kept: Vec<bool> = self.disclosures.iter().map(&mut keep).collect();
self.selected(&processed, &kept)
}
pub fn disclose<'a>(
&self,
paths: impl IntoIterator<Item = &'a ClaimPath>,
) -> Result<Self, SdJwtError> {
let processed = self.process()?;
let mut kept = vec![false; self.disclosures.len()];
for path in paths {
let mut found = false;
for claim in &processed.disclosed {
if path.is_prefix_of(&claim.path) {
found = true;
let mut index = Some(claim.disclosure);
while let Some(i) = index {
kept[i] = true;
index = processed.parents[i];
}
}
}
if !found {
return Err(SdJwtError::NoSuchClaim(path.clone()));
}
}
self.selected(&processed, &kept)
}
fn selected(&self, processed: &Processed, kept: &[bool]) -> Result<Self, SdJwtError> {
for (i, disclosure) in self.disclosures.iter().enumerate() {
if !kept[i] {
continue;
}
if let Some(parent) = processed.parents[i]
&& !kept[parent]
{
return Err(SdJwtError::MissingParentDisclosure {
disclosure: disclosure.as_str().to_owned(),
parent: self.disclosures[parent].as_str().to_owned(),
});
}
}
let disclosures = self
.disclosures
.iter()
.zip(kept)
.filter(|(_, kept)| **kept)
.map(|(d, _)| d.clone())
.collect();
Ok(Self::new(self.jwt.clone(), disclosures))
}
pub async fn sign_key_binding(
self,
binding: &KeyBinding,
signer: &dyn Signer,
) -> Result<Self, SdJwtError> {
let sd_jwt = self.without_key_binding();
let alg = sd_jwt.hash_algorithm()?;
let claims = binding.claims(sd_jwt.sd_hash(alg));
let header = JwsHeader::new(signer.alg()).with_typ(KB_JWT_TYP);
let key_binding = Jwt::sign(&header, &claims, signer).await?;
Ok(sd_jwt.with_key_binding(key_binding))
}
}
fn hash_algorithm(payload: &Map<String, Value>) -> Result<HashAlgorithm, SdJwtError> {
match payload.get(HashAlgorithm::CLAIM) {
None => Ok(HashAlgorithm::default()),
Some(Value::String(name)) => HashAlgorithm::parse(name),
Some(_) => Err(SdJwtError::Payload("`_sd_alg` is not a string".into())),
}
}
impl fmt::Debug for SdJwt {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("SdJwt")
.field("jwt", &self.jwt)
.field("disclosures", &self.disclosures)
.field("key_binding", &self.key_binding)
.finish()
}
}
impl fmt::Display for SdJwt {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.write_str(&self.to_string_without_key_binding())?;
if let Some(key_binding) = &self.key_binding {
f.write_str(key_binding.as_str())?;
}
Ok(())
}
}
impl FromStr for SdJwt {
type Err = SdJwtError;
fn from_str(s: &str) -> Result<Self, Self::Err> {
Self::parse(s)
}
}
impl Serialize for SdJwt {
fn serialize<S: Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
serializer.collect_str(self)
}
}
impl<'de> Deserialize<'de> for SdJwt {
fn deserialize<D: Deserializer<'de>>(deserializer: D) -> Result<Self, D::Error> {
let s = String::deserialize(deserializer)?;
Self::parse(&s).map_err(serde::de::Error::custom)
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
#[non_exhaustive]
pub struct Processed {
pub payload: Map<String, Value>,
pub disclosed: Vec<DisclosedClaim>,
pub undisclosed: Vec<ClaimPath>,
issued: Map<String, Value>,
parents: Vec<Option<usize>>,
}
impl Processed {
pub fn as_issued(&self) -> &Map<String, Value> {
&self.issued
}
pub fn disclosed_paths(&self) -> impl Iterator<Item = &ClaimPath> {
self.disclosed.iter().map(|claim| &claim.path)
}
pub fn is_disclosed(&self, path: &ClaimPath) -> bool {
self.disclosed.iter().any(|claim| claim.path == *path)
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
#[non_exhaustive]
pub struct DisclosedClaim {
pub path: ClaimPath,
pub disclosure: usize,
}
fn digest_placeholder(value: &Value) -> Option<&str> {
let object = value.as_object()?;
if object.len() != 1 {
return None;
}
object.get(ELLIPSIS_KEY)?.as_str()
}
fn strip_placeholders(value: &mut Value) {
match value {
Value::Object(object) => object.values_mut().for_each(strip_placeholders),
Value::Array(array) => {
array.retain(|element| digest_placeholder(element).is_none());
array.iter_mut().for_each(strip_placeholders);
}
_ => {}
}
}
struct Processor<'a> {
disclosures: &'a [Disclosure],
digests: HashMap<String, usize>,
used: Vec<bool>,
parents: Vec<Option<usize>>,
seen: HashSet<String>,
disclosed: Vec<DisclosedClaim>,
undisclosed: Vec<ClaimPath>,
}
impl<'a> Processor<'a> {
fn new(alg: HashAlgorithm, disclosures: &'a [Disclosure]) -> Self {
let digests = disclosures
.iter()
.enumerate()
.map(|(i, d)| (d.digest(alg), i))
.collect();
Self {
disclosures,
digests,
used: vec![false; disclosures.len()],
parents: vec![None; disclosures.len()],
seen: HashSet::new(),
disclosed: Vec::new(),
undisclosed: Vec::new(),
}
}
fn lookup(&mut self, digest: &str, owner: Option<usize>) -> Result<Option<usize>, SdJwtError> {
if !self.seen.insert(digest.to_owned()) {
return Err(SdJwtError::DuplicateDigest(digest.to_owned()));
}
let Some(&index) = self.digests.get(digest) else {
return Ok(None);
};
self.used[index] = true;
self.parents[index] = owner;
Ok(Some(index))
}
fn value(
&mut self,
value: Value,
path: &ClaimPath,
owner: Option<usize>,
depth: usize,
) -> Result<Value, SdJwtError> {
match value {
Value::Object(_) | Value::Array(_) if depth > SdJwtBuilder::MAX_DEPTH => {
Err(SdJwtError::TooDeep {
max: SdJwtBuilder::MAX_DEPTH,
})
}
Value::Object(object) => self.object(object, path, owner, depth).map(Value::Object),
Value::Array(array) => self.array(array, path, owner, depth).map(Value::Array),
other => Ok(other),
}
}
fn object(
&mut self,
object: Map<String, Value>,
path: &ClaimPath,
owner: Option<usize>,
depth: usize,
) -> Result<Map<String, Value>, SdJwtError> {
let mut out = Map::new();
let mut digests = None;
for (name, value) in object {
if name == SD_KEY {
let Value::Array(array) = value else {
return Err(SdJwtError::Payload(format!(
"`_sd` at {path} is not an array"
)));
};
digests = Some(array);
continue;
}
let value = self.value(value, &path.join(name.as_str()), owner, depth + 1)?;
out.insert(name, value);
}
for digest in digests.unwrap_or_default() {
let Value::String(digest) = digest else {
return Err(SdJwtError::Payload(format!(
"`_sd` at {path} holds a non-string"
)));
};
let Some(index) = self.lookup(&digest, owner)? else {
continue;
};
let disclosure = &self.disclosures[index];
let Some(name) = disclosure.name() else {
return Err(SdJwtError::WrongDisclosureKind {
disclosure: disclosure.as_str().to_owned(),
expected: "an object property disclosure",
});
};
if name == SD_KEY || name == ELLIPSIS_KEY {
return Err(SdJwtError::ReservedClaimName(name.to_owned()));
}
if out.contains_key(name) {
return Err(SdJwtError::DuplicateClaim {
name: name.to_owned(),
path: path.clone(),
});
}
let name = name.to_owned();
let claim_path = path.join(name.as_str());
let value = self.value(
disclosure.value().clone(),
&claim_path,
Some(index),
depth + 1,
)?;
out.insert(name, value);
self.disclosed.push(DisclosedClaim {
path: claim_path,
disclosure: index,
});
}
Ok(out)
}
fn array(
&mut self,
array: Vec<Value>,
path: &ClaimPath,
owner: Option<usize>,
depth: usize,
) -> Result<Vec<Value>, SdJwtError> {
let mut out = Vec::with_capacity(array.len());
for (i, element) in array.into_iter().enumerate() {
let element_path = path.join(i);
let Some(digest) = digest_placeholder(&element).map(str::to_owned) else {
out.push(self.value(element, &element_path, owner, depth + 1)?);
continue;
};
let Some(index) = self.lookup(&digest, owner)? else {
self.undisclosed.push(element_path);
out.push(element);
continue;
};
let disclosure = &self.disclosures[index];
if !disclosure.is_element() {
return Err(SdJwtError::WrongDisclosureKind {
disclosure: disclosure.as_str().to_owned(),
expected: "an array element disclosure",
});
}
let value = self.value(
disclosure.value().clone(),
&element_path,
Some(index),
depth + 1,
)?;
out.push(value);
self.disclosed.push(DisclosedClaim {
path: element_path,
disclosure: index,
});
}
Ok(out)
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
#[non_exhaustive]
pub struct KeyBinding {
pub audience: String,
pub nonce: String,
pub issued_at: i64,
pub additional: IndexMap<String, Value>,
}
impl KeyBinding {
pub fn new(
audience: impl Into<String>,
nonce: impl Into<String>,
issued_at: DateTime<Utc>,
) -> Self {
Self {
audience: audience.into(),
nonce: nonce.into(),
issued_at: JwtClaims::numeric_date(issued_at),
additional: IndexMap::new(),
}
}
#[must_use]
pub fn with_claim(mut self, name: impl Into<String>, value: impl Into<Value>) -> Self {
self.additional.insert(name.into(), value.into());
self
}
pub fn claims(&self, sd_hash: String) -> JwtClaims {
let mut claims = JwtClaims::new().with_aud(self.audience.clone());
claims.iat = Some(self.issued_at);
claims
.additional
.insert("nonce".to_owned(), Value::String(self.nonce.clone()));
claims
.additional
.insert("sd_hash".to_owned(), Value::String(sd_hash));
claims.additional.extend(self.additional.clone());
claims
}
}