open-standards/joule-code/joule-code-rs/crates/joule-code-receipt/src/lib.rs

//! The Joule Code receipt (SPEC §6, `spec/receipt.md`).
//!
//! A [`Receipt`] is the signed, two-part energy record (`measured + estimated`,
//! integer microjoules) carried in a Joule Wire Protocol frame. Every reading
//! carries a [`Provenance`] tag; the spec never claims an accuracy the platform
//! cannot deliver, so a `measured` figure MUST come from hardware.

use ed25519_dalek::SigningKey;
use joule_code_schema::{ContentId, Tier};
use serde::{Deserialize, Serialize};

pub mod jcr1;
pub mod intoto;

/// Reference estimator coefficients (non-normative, SPEC §6 / `spec/receipt.md`
/// §4): joules per token, in microjoules. Output dominates prefill. Tunable;
/// the standard is that the figure is reported and honestly tagged, not the
/// constant.
pub const J_IN_UJ: u64 = 200_000;
pub const J_OUT_UJ: u64 = 2_000_000;

/// How honest an energy reading is (SPEC §6).
#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub enum Provenance {
    /// Read from a hardware energy interface: an external shunt, or a vendor
    /// energy counter (RAPL, NVML, SMC/IOReport) only while that counter is
    /// demonstrably live for the domain the figure claims (`spec/receipt.md`
    /// §2) — a present counter that is not live falls back to the estimator.
    /// Requires a present `measured_uj`.
    HwShunt,
    /// Computed from a calibrated per-operation model. `measured_uj` may be absent.
    ModelBased,
    /// A coarse constant estimator. `measured_uj` MUST be absent.
    Estimator,
}

/// A two-part energy record, all integer microjoules (SPEC §6).
#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct Energy {
    /// Present ONLY when obtained from a hardware interface.
    pub measured_uj: Option<u64>,
    /// Always present.
    pub estimated_uj: u64,
    pub provenance: Provenance,
}

impl Energy {
    /// The estimator of last resort: from token counts only (SPEC §6 §4).
    /// Tags the reading `Estimator` and leaves `measured_uj` absent.
    pub fn from_tokens(input_tokens: u64, output_tokens: u64) -> Self {
        Energy {
            measured_uj: None,
            estimated_uj: input_tokens * J_IN_UJ + output_tokens * J_OUT_UJ,
            provenance: Provenance::Estimator,
        }
    }

    /// Enforce the provenance ↔ measured coupling (SPEC §6):
    /// `HwShunt` requires `measured_uj`; `Estimator` forbids it.
    pub fn check_provenance(&self) -> Result<(), ReceiptError> {
        match (self.provenance, self.measured_uj) {
            (Provenance::HwShunt, None) => Err(ReceiptError::HwShuntRequiresMeasured),
            (Provenance::Estimator, Some(_)) => Err(ReceiptError::EstimatorMustNotCarryMeasured),
            _ => Ok(()),
        }
    }
}

/// Errors from building, signing, or verifying a receipt.
#[derive(Debug, thiserror::Error)]
pub enum ReceiptError {
    /// `Provenance::HwShunt` with no `measured_uj`.
    #[error("HwShunt provenance requires a measured figure")]
    HwShuntRequiresMeasured,
    /// `Provenance::Estimator` with a `measured_uj` present (dishonest).
    #[error("Estimator provenance must not carry a measured figure")]
    EstimatorMustNotCarryMeasured,
    #[error("bad key or signature: {0}")]
    Crypto(String),
    #[error("receipt signature did not verify")]
    BadSignature,
    /// The envelope verified but its payload did not match this receipt's core.
    #[error("envelope payload does not match receipt")]
    PayloadMismatch,
    #[error(transparent)]
    Jcr1(#[from] jcr1::Jcr1Error),
}

impl Provenance {
    /// The JCR-1 wire label for this provenance (`"HwShunt"`, `"ModelBased"`,
    /// `"Estimator"`), matching the family's receipt encoding.
    pub fn jcr1_label(self) -> &'static str {
        match self {
            Provenance::HwShunt => "HwShunt",
            Provenance::ModelBased => "ModelBased",
            Provenance::Estimator => "Estimator",
        }
    }
}

/// A signed energy receipt for one coding operation (SPEC §6).
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct Receipt {
    pub op_id: ContentId,
    pub tier: Tier,
    pub energy: Energy,
    pub sequence: u64,
    pub nonce: [u8; 8],
}

impl Receipt {
    /// Project this receipt to its JCR-1 wire JSON — the `joule-code/1` profile
    /// of the family's receipt format (SPEC §6). `steward` is the endpoint
    /// identity placed in the COSE `kid`. Null fields (e.g. an absent
    /// `measured_uj`) are dropped from the signed core.
    pub fn to_jcr1_wire(&self, steward: &str) -> serde_json::Value {
        use serde_json::json;
        let mut obj = serde_json::Map::new();
        obj.insert("jc_receipt".into(), json!("joule-code/1"));
        obj.insert("op_id".into(), json!(self.op_id.as_str()));
        obj.insert("tier".into(), json!(self.tier.wire_tag()));
        obj.insert("estimated_uj".into(), json!(self.energy.estimated_uj));
        if let Some(m) = self.energy.measured_uj {
            obj.insert("measured_uj".into(), json!(m));
        }
        obj.insert(
            "energy_provenance".into(),
            json!(self.energy.provenance.jcr1_label()),
        );
        obj.insert("sequence".into(), json!(self.sequence));
        obj.insert("nonce".into(), json!(hex_lower(&self.nonce)));
        obj.insert("steward".into(), json!(steward));
        serde_json::Value::Object(obj)
    }

    /// The canonical signed-core bytes for this receipt (JCR-1 det-CBOR).
    pub fn signed_core(&self, steward: &str) -> Result<Vec<u8>, ReceiptError> {
        self.energy.check_provenance()?;
        Ok(jcr1::signed_core_bytes(&self.to_jcr1_wire(steward))?)
    }

    /// Seal the receipt into a JCR-1 COSE_Sign1 envelope (SPEC §6). Validates
    /// the provenance coupling first so a dishonest receipt cannot be sealed.
    pub fn seal(&self, key: &SigningKey, steward: &str) -> Result<Vec<u8>, ReceiptError> {
        let core = self.signed_core(steward)?;
        Ok(jcr1::cose_sign1_envelope(&core, key, steward))
    }

    /// Verify a JCR-1 COSE_Sign1 envelope against an endpoint key and confirm
    /// its payload is exactly this receipt's canonical signed core.
    pub fn verify_envelope(
        &self,
        envelope: &[u8],
        steward: &str,
        endpoint: &[u8; 32],
    ) -> Result<(), ReceiptError> {
        let payload = jcr1::verify_cose_sign1(envelope, steward, endpoint)?;
        if payload != self.signed_core(steward)? {
            return Err(ReceiptError::PayloadMismatch);
        }
        Ok(())
    }
}

/// Lowercase hex without external deps (nonce rendering for the wire form).
pub(crate) fn hex_lower(bytes: &[u8]) -> String {
    let mut s = String::with_capacity(bytes.len() * 2);
    for b in bytes {
        s.push(char::from_digit((b >> 4) as u32, 16).unwrap());
        s.push(char::from_digit((b & 0x0f) as u32, 16).unwrap());
    }
    s
}

#[cfg(test)]
mod tests {
    use super::*;
    use ed25519_dalek::SigningKey;

    fn rcpt(energy: Energy) -> Receipt {
        Receipt {
            op_id: ContentId("b3:AAAA".into()),
            tier: Tier::L4,
            energy,
            sequence: 1,
            nonce: [0, 0, 0, 0, 0, 0, 0, 1],
        }
    }

    #[test]
    fn estimator_seal_verify_roundtrips() {
        let sk = SigningKey::from_bytes(&[9u8; 32]);
        let steward = "did:example:endpoint1";
        let r = rcpt(Energy::from_tokens(2000, 2000));
        // 2000*200_000 (in) + 2000*2_000_000 (out) = 4_400_000_000 µJ (4400 J).
        assert_eq!(r.energy.estimated_uj, 4_400_000_000);
        let env = r.seal(&sk, steward).unwrap();
        assert!(r.verify_envelope(&env, steward, &sk.verifying_key().to_bytes()).is_ok());
        // A different steward must not verify.
        assert!(r.verify_envelope(&env, "did:example:other", &sk.verifying_key().to_bytes()).is_err());
    }

    #[test]
    fn hwshunt_without_measured_is_rejected() {
        let r = rcpt(Energy {
            measured_uj: None,
            estimated_uj: 2_000_000,
            provenance: Provenance::HwShunt,
        });
        assert!(matches!(
            r.energy.check_provenance(),
            Err(ReceiptError::HwShuntRequiresMeasured)
        ));
    }

    #[test]
    fn estimator_with_measured_is_rejected() {
        let r = rcpt(Energy {
            measured_uj: Some(1_000_000),
            estimated_uj: 4_400_000,
            provenance: Provenance::Estimator,
        });
        assert!(matches!(
            r.energy.check_provenance(),
            Err(ReceiptError::EstimatorMustNotCarryMeasured)
        ));
    }
}