world_id_primitives/

session.rs

use crate::FieldElement;
use embed_doc_image::embed_doc_image;
use ruint::aliases::U256;
use serde::{Deserialize, Deserializer, Serialize, Serializer, de::Error as _};

#[expect(unused_imports, reason = "used in doc comments")]
use crate::circuit_inputs::QueryProofCircuitInput;

const SESSION_FIELD_ELEMENT_PREFIX: u8 = 0x01;

/// Allows field element generation for Session Proofs
pub trait SessionFieldElement {
    /// Generate a randomized field element with a specific prefix used
    /// only for Session Proofs.
    ///
    /// This is generally used to:
    /// - Generate the [`SessionId::oprf_seed`] used as input for the OPRF derivation of `session_id_r_seed`
    /// - Generate a randomized action for the nullifier input of a Session Proof
    fn random_for_session<R: rand::CryptoRng + rand::RngCore>(rng: &mut R) -> FieldElement;
    /// Returns whether a Field Element is valid for Session Proof use, i.e. it has
    /// the right prefix
    fn is_valid_for_session(&self) -> bool;
}

impl SessionFieldElement for FieldElement {
    fn random_for_session<R: rand::CryptoRng + rand::RngCore>(rng: &mut R) -> FieldElement {
        let mut bytes = [0u8; 32];
        rng.fill_bytes(&mut bytes);
        bytes[0] = SESSION_FIELD_ELEMENT_PREFIX;
        let seed = U256::from_be_bytes(bytes);
        Self::try_from(seed).expect(
            "should always fit in the field because with 0x01 as the MSB, the field element < babyjubjub modulus",
        )
    }

    fn is_valid_for_session(&self) -> bool {
        self.to_be_bytes()[0] == SESSION_FIELD_ELEMENT_PREFIX
    }
}

/// An identifier for a session (can be re-used).
///
/// A session allows RPs to ensure that it's still the same World ID
/// interacting with them across multiple interactions.
///
/// A `SessionId` is obtained after creating an initial session.
///
/// See the diagram below on how Session Proofs work, the [`SessionId`] and the `r` seed
/// ![Session Proofs Diagram][session-proofs.png]
///
/// Note that the `action` stored here is unrelated to the randomized action used
/// internally by [`SessionNullifier`]s — that randomized action exists only to ensure
/// the circuit's nullifier output is unique per Session Proof.
#[embed_doc_image("session-proofs.png", "assets/session-proofs.png")]
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct SessionId {
    /// The actual commitment being verified in the ZK-circuit.
    ///
    /// It is computed as H(DS_C || leaf_index || session_id_r_seed), see
    /// `signal computed_id_commitment` in `oprf_nullifier.circom`.
    commitment: FieldElement,
    /// A random seed generated by the authenticator in the initial Uniqueness Proof.
    ///
    /// This seed is the input to the OPRF Query to derive `session_id_r_seed` (`r`). It
    /// is part of the `session_id` so the RP can provide it when requesting a Session Proof.
    ///
    /// # Important: Prefix
    /// To ensure there are no collisions between the generated `r`s and the nullifiers
    /// for Uniqueness Proofs (as they use the same OPRF Key and query structure), the
    /// `oprf_seed`s, which are plugged as `action` in the Query Proof (see [`QueryProofCircuitInput`]),
    /// MUST be prefixed with an explicit byte of `0x01`. All other actions have a `0x00` byte prefix. This
    /// collision avoidance is important because it ensures that any requests for nullifiers meant
    /// for Uniqueness Proofs are always signed by the RP (otherwise, an RP signature for a Session Proof
    /// could be used for requesting computation of _any_ nullifier).
    ///
    /// # Re-derivation
    ///
    /// The Authenticator can deterministically re-derive `r` from the OPRF nodes without
    /// needing to cache `r` locally as:
    /// ```text
    /// r = OPRF(pk_rpId, DS_C || leafIndex || oprf_seed)
    /// ```
    oprf_seed: FieldElement,
}

impl SessionId {
    const JSON_PREFIX: &str = "session_";

    /// Creates a new session id. Most uses should default to `from_r_seed` instead.
    ///
    /// # Panics
    /// If the provided `oprf_seed` is not prefixed with a `0x01` byte.
    #[must_use]
    pub fn new(commitment: FieldElement, oprf_seed: FieldElement) -> Self {
        // OPRF Seeds must always start with a byte of `0x01`. See [`Self::oprf_seed`]
        // for details. Panic is acceptable as `oprf_seed` generation should
        // generally be done with `Self::from_r_seed`
        assert!(oprf_seed.is_valid_for_session());
        Self {
            commitment,
            oprf_seed,
        }
    }

    /// Initializes a `SessionId` from the OPRF-output seed (`r`), and the `oprf_seed`
    /// used as input for the OPRF computation.
    ///
    /// This matches the logic in `oprf_nullifier.circom` for computing the `commitment` from the OPRF seed.
    ///
    /// # Seed (`session_id_r_seed`)
    /// - The seed is and MUST be computationally indistinguishable from random,
    ///   i.e. uniformly distributed because it uses OPRF.
    /// - When computed, the OPRF nodes will use the same `oprfKeyId` for the RP, with a different domain separator.
    /// - Requesting this seed requires a properly signed request from the RP and a complete query proof.
    /// - The seed generation is based on a randomly generated seed used as an "action" in a Query Proof. Note
    ///   this `action` is different than the randomized action used internally by [`SessionNullifier`]s.
    pub fn from_r_seed<R: rand::CryptoRng + rand::RngCore>(
        leaf_index: u64,
        session_id_r_seed: FieldElement,
        oprf_seed: Option<FieldElement>,
        rng: &mut R,
    ) -> Result<Self, &str> {
        let sub_ds = FieldElement::from_be_bytes_mod_order(b"H(id, r)");

        let oprf_seed = if let Some(seed) = oprf_seed {
            if !seed.is_valid_for_session() {
                return Err("oprf_seed is not valid for session");
            }
            seed
        } else {
            FieldElement::random_for_session(rng)
        };

        let mut input = [*sub_ds, leaf_index.into(), *session_id_r_seed];
        poseidon2::bn254::t3::permutation_in_place(&mut input);
        let commitment = input[1].into();
        Ok(Self {
            commitment,
            oprf_seed,
        })
    }

    /// Returns the `commitment` value.
    #[must_use]
    pub const fn commitment(&self) -> FieldElement {
        self.commitment
    }

    /// Returns the `oprf_seed` value.
    #[must_use]
    pub const fn oprf_seed(&self) -> FieldElement {
        self.oprf_seed
    }

    /// Returns the 64-byte big-endian representation (2 x 32-byte field elements).
    #[must_use]
    pub fn to_compressed_bytes(&self) -> [u8; 64] {
        let mut bytes = [0u8; 64];
        bytes[..32].copy_from_slice(&self.commitment.to_be_bytes());
        bytes[32..].copy_from_slice(&self.oprf_seed.to_be_bytes());
        bytes
    }

    /// Constructs from compressed bytes (must be exactly 64 bytes).
    ///
    /// # Errors
    /// Returns an error if the input is not exactly 64 bytes or if values are not valid field elements.
    pub fn from_compressed_bytes(bytes: &[u8]) -> Result<Self, String> {
        if bytes.len() != 64 {
            return Err(format!(
                "Invalid length: expected 64 bytes, got {}",
                bytes.len()
            ));
        }

        let commitment = FieldElement::from_be_bytes(bytes[..32].try_into().unwrap())
            .map_err(|e| format!("invalid commitment: {e}"))?;
        let oprf_seed = FieldElement::from_be_bytes(bytes[32..].try_into().unwrap())
            .map_err(|e| format!("invalid oprf_seed: {e}"))?;

        if bytes[32] != SESSION_FIELD_ELEMENT_PREFIX {
            return Err("invalid prefix for oprf_seed".to_string());
        }

        Ok(Self {
            commitment,
            oprf_seed,
        })
    }
}

impl Default for SessionId {
    fn default() -> Self {
        let mut oprf_seed = [0u8; 32];
        oprf_seed[0] = SESSION_FIELD_ELEMENT_PREFIX;
        let oprf_seed = U256::from_be_bytes(oprf_seed)
            .try_into()
            .expect("always fits in the field");
        Self {
            commitment: FieldElement::ZERO,
            oprf_seed,
        }
    }
}

impl Serialize for SessionId {
    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
    where
        S: Serializer,
    {
        let bytes = self.to_compressed_bytes();
        if serializer.is_human_readable() {
            // JSON: prefixed hex-encoded compressed bytes for explicit typing.
            serializer.serialize_str(&format!("{}{}", Self::JSON_PREFIX, hex::encode(bytes)))
        } else {
            // Binary: compressed bytes
            serializer.serialize_bytes(&bytes)
        }
    }
}

impl<'de> Deserialize<'de> for SessionId {
    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
    where
        D: Deserializer<'de>,
    {
        let bytes = if deserializer.is_human_readable() {
            let value = String::deserialize(deserializer)?;
            let hex_str = value.strip_prefix(Self::JSON_PREFIX).ok_or_else(|| {
                D::Error::custom(format!(
                    "session id must start with '{}'",
                    Self::JSON_PREFIX
                ))
            })?;
            hex::decode(hex_str).map_err(D::Error::custom)?
        } else {
            Vec::deserialize(deserializer)?
        };

        Self::from_compressed_bytes(&bytes).map_err(D::Error::custom)
    }
}

/// A session nullifier for World ID Session proofs. It is analogous to a request nonce,
/// it **does NOT guarantee uniqueness of a World ID** as a `Nullifier` does.
///
/// This type is intended to be opaque for RPs. For an RP context, they should only
/// be concerned of this needing to be passthrough to the `verifySession()` contract function.
///
/// This type exists as an adaptation to be able to use the same ZK-circuit for
/// both Uniqueness Proofs and Session Proofs, and it encompasses:
/// - the nullifier used as the proof output.
/// - a random action bound to the same proof.
///
/// The `WorldIDVerifier.sol` contract expects this as a `uint256[2]` array
/// use `as_ethereum_representation()` for conversion.
///
/// # Future
///
/// Note the session nullifier exists **only** to support the same ZK-circuit than for Uniqueness Proofs; as
/// World ID evolves to a different proving system which won't require circuit precompiles, a new circuit MUST
/// be created which does not generate a nullifier at all, and the input randomized action will not be required either.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct SessionNullifier {
    /// The nullifier value for this proof.
    nullifier: FieldElement,
    /// The random action value bound to this session proof.
    action: FieldElement,
}

impl SessionNullifier {
    const JSON_PREFIX: &str = "snil_";

    /// Creates a new session nullifier.
    #[must_use]
    pub const fn new(nullifier: FieldElement, action: FieldElement) -> Self {
        Self { nullifier, action }
    }

    /// Returns the nullifier value.
    #[must_use]
    pub const fn nullifier(&self) -> FieldElement {
        self.nullifier
    }

    /// Returns the action value.
    #[must_use]
    pub const fn action(&self) -> FieldElement {
        self.action
    }

    /// Returns the session nullifier as an Ethereum-compatible array for `verifySession()`.
    ///
    /// Format: `[nullifier, action]` matching the contract's `uint256[2] sessionNullifier`.
    #[must_use]
    pub fn as_ethereum_representation(&self) -> [U256; 2] {
        [self.nullifier.into(), self.action.into()]
    }

    /// Creates a session nullifier from an Ethereum representation.
    ///
    /// # Errors
    /// Returns an error if the U256 values are not valid field elements.
    pub fn from_ethereum_representation(value: [U256; 2]) -> Result<Self, String> {
        let nullifier =
            FieldElement::try_from(value[0]).map_err(|e| format!("invalid nullifier: {e}"))?;
        let action =
            FieldElement::try_from(value[1]).map_err(|e| format!("invalid action: {e}"))?;
        Ok(Self { nullifier, action })
    }

    /// Returns the 64-byte big-endian representation (2 x 32-byte field elements).
    #[must_use]
    pub fn to_compressed_bytes(&self) -> [u8; 64] {
        let mut bytes = [0u8; 64];
        bytes[..32].copy_from_slice(&self.nullifier.to_be_bytes());
        bytes[32..].copy_from_slice(&self.action.to_be_bytes());
        bytes
    }

    /// Constructs from compressed bytes (must be exactly 64 bytes).
    ///
    /// # Errors
    /// Returns an error if the input is not exactly 64 bytes or if values are not valid field elements.
    pub fn from_compressed_bytes(bytes: &[u8]) -> Result<Self, String> {
        if bytes.len() != 64 {
            return Err(format!(
                "Invalid length: expected 64 bytes, got {}",
                bytes.len()
            ));
        }

        let nullifier = FieldElement::from_be_bytes(bytes[..32].try_into().unwrap())
            .map_err(|e| format!("invalid nullifier: {e}"))?;
        let action = FieldElement::from_be_bytes(bytes[32..].try_into().unwrap())
            .map_err(|e| format!("invalid action: {e}"))?;

        Ok(Self { nullifier, action })
    }
}

impl Default for SessionNullifier {
    fn default() -> Self {
        Self {
            nullifier: FieldElement::ZERO,
            action: FieldElement::ZERO,
        }
    }
}

impl Serialize for SessionNullifier {
    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
    where
        S: Serializer,
    {
        let bytes = self.to_compressed_bytes();
        if serializer.is_human_readable() {
            // JSON: prefixed hex-encoded compressed bytes for explicit typing.
            serializer.serialize_str(&format!("{}{}", Self::JSON_PREFIX, hex::encode(bytes)))
        } else {
            // Binary: compressed bytes
            serializer.serialize_bytes(&bytes)
        }
    }
}

impl<'de> Deserialize<'de> for SessionNullifier {
    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
    where
        D: Deserializer<'de>,
    {
        let bytes = if deserializer.is_human_readable() {
            let value = String::deserialize(deserializer)?;
            let hex_str = value.strip_prefix(Self::JSON_PREFIX).ok_or_else(|| {
                D::Error::custom(format!(
                    "session nullifier must start with '{}'",
                    Self::JSON_PREFIX
                ))
            })?;
            hex::decode(hex_str).map_err(D::Error::custom)?
        } else {
            Vec::deserialize(deserializer)?
        };

        Self::from_compressed_bytes(&bytes).map_err(D::Error::custom)
    }
}

impl From<SessionNullifier> for [U256; 2] {
    fn from(value: SessionNullifier) -> Self {
        value.as_ethereum_representation()
    }
}

impl From<(FieldElement, FieldElement)> for SessionNullifier {
    fn from((nullifier, action): (FieldElement, FieldElement)) -> Self {
        Self::new(nullifier, action)
    }
}

#[cfg(test)]
mod session_id_tests {
    use super::*;
    use ruint::uint;

    fn test_field_element(value: u64) -> FieldElement {
        FieldElement::from(value)
    }

    /// Creates an oprf_seed with the right prefix
    fn test_oprf_seed(value: u64) -> FieldElement {
        // set the first byte to 0x01; no need to clear the first bits as the input is u64
        let n = U256::from(value)
            | uint!(0x0100000000000000000000000000000000000000000000000000000000000000_U256);
        FieldElement::try_from(n).expect("test value fits in field")
    }

    #[test]
    fn test_new_and_accessors() {
        let commitment = test_field_element(1001);
        let seed = test_oprf_seed(42);
        let id = SessionId::new(commitment, seed);

        assert_eq!(id.commitment(), commitment);
        assert_eq!(id.oprf_seed(), seed);
    }

    #[test]
    fn test_default() {
        let id = SessionId::default();
        assert_eq!(id.commitment(), FieldElement::ZERO);
        assert_eq!(
            id.oprf_seed(),
            uint!(0x0100000000000000000000000000000000000000000000000000000000000000_U256)
                .try_into()
                .unwrap()
        );
    }

    #[test]
    fn test_bytes_roundtrip() {
        let id = SessionId::new(test_field_element(1001), test_oprf_seed(42));
        let bytes = id.to_compressed_bytes();

        assert_eq!(bytes.len(), 64);

        let decoded = SessionId::from_compressed_bytes(&bytes).unwrap();
        assert_eq!(id, decoded);
    }

    #[test]
    fn test_bytes_use_field_element_encoding() {
        let id = SessionId::new(test_field_element(1001), test_oprf_seed(42));
        let bytes = id.to_compressed_bytes();

        let mut expected = [0u8; 64];
        expected[..32].copy_from_slice(&id.commitment().to_be_bytes());
        expected[32..].copy_from_slice(&id.oprf_seed().to_be_bytes());
        assert_eq!(bytes, expected);
    }

    #[test]
    fn test_invalid_bytes_length() {
        let too_short = vec![0u8; 63];
        let result = SessionId::from_compressed_bytes(&too_short);
        assert!(result.is_err());
        assert!(result.unwrap_err().contains("Invalid length"));

        let too_long = vec![0u8; 65];
        let result = SessionId::from_compressed_bytes(&too_long);
        assert!(result.is_err());
        assert!(result.unwrap_err().contains("Invalid length"));
    }

    #[test]
    fn test_from_compressed_bytes_rejects_wrong_oprf_seed_prefix() {
        let mut bytes = [0u8; 64];
        // Valid commitment (zero is a valid field element)
        // oprf_seed with wrong prefix: 0x00 instead of 0x01
        bytes[32] = 0x00;
        let result = SessionId::from_compressed_bytes(&bytes);
        assert!(result.is_err());
        assert!(
            result.unwrap_err().contains("invalid prefix"),
            "should reject oprf_seed without 0x01 prefix"
        );
    }

    #[test]
    fn test_json_roundtrip() {
        let id = SessionId::new(test_field_element(1001), test_oprf_seed(42));
        let json = serde_json::to_string(&id).unwrap();

        assert!(json.starts_with("\"session_"));
        assert!(json.ends_with('"'));

        let decoded: SessionId = serde_json::from_str(&json).unwrap();
        assert_eq!(id, decoded);
    }

    #[test]
    fn test_json_format() {
        let id = SessionId::new(test_field_element(1), test_oprf_seed(2));
        let json = serde_json::to_string(&id).unwrap();

        let parsed: serde_json::Value = serde_json::from_str(&json).unwrap();
        assert!(parsed.is_string());
        let value = parsed.as_str().unwrap();
        assert!(value.starts_with("session_"));
    }

    #[test]
    fn test_json_wrong_prefix_rejected() {
        let result = serde_json::from_str::<SessionId>("\"snil_00\"");
        assert!(result.is_err());
    }

    #[test]
    fn test_from_r_seed_generates_random_seed() {
        let mut rng = rand::rngs::OsRng;
        let r_seed = test_field_element(999);

        let id1 = SessionId::from_r_seed(0, r_seed, None, &mut rng).unwrap();
        let id2 = SessionId::from_r_seed(0, r_seed, None, &mut rng).unwrap();

        assert_ne!(id1.oprf_seed(), id2.oprf_seed());
    }

    #[test]
    fn test_from_r_seed_generated_seed_has_session_prefix() {
        let mut rng = rand::rngs::OsRng;
        let r_seed = test_field_element(999);

        for _ in 0..50 {
            let id = SessionId::from_r_seed(0, r_seed, None, &mut rng).unwrap();
            // Top byte must be exactly 0x01: bit 248 set, bits 249-255 clear
            assert_eq!(id.oprf_seed().to_u256() >> 248, U256::from(1));
        }
    }

    #[test]
    fn test_from_r_seed_commitment_snapshot() {
        let leaf_index = 42u64;
        let r_seed = test_field_element(123);
        let oprf_seed = test_oprf_seed(456);

        let session_id =
            SessionId::from_r_seed(leaf_index, r_seed, Some(oprf_seed), &mut rand::rngs::OsRng)
                .unwrap();

        let expected = "0x1e7853ebd4fc9d9f0232fdcfae116023610bdf66a22e2700445d7a2e0e7e6152"
            .parse::<U256>()
            .unwrap();
        assert_eq!(
            session_id.commitment().to_u256(),
            expected,
            "commitment snapashot for session commitment changed"
        );
    }
}

#[cfg(test)]
mod session_nullifier_tests {
    use super::*;

    fn test_field_element(value: u64) -> FieldElement {
        FieldElement::from(value)
    }

    #[test]
    fn test_new_and_accessors() {
        let nullifier = test_field_element(1001);
        let action = test_field_element(42);
        let session = SessionNullifier::new(nullifier, action);

        assert_eq!(session.nullifier(), nullifier);
        assert_eq!(session.action(), action);
    }

    #[test]
    fn test_as_ethereum_representation() {
        let nullifier = test_field_element(100);
        let action = test_field_element(200);
        let session = SessionNullifier::new(nullifier, action);

        let repr = session.as_ethereum_representation();
        assert_eq!(repr[0], U256::from(100));
        assert_eq!(repr[1], U256::from(200));
    }

    #[test]
    fn test_from_ethereum_representation() {
        let repr = [U256::from(100), U256::from(200)];
        let session = SessionNullifier::from_ethereum_representation(repr).unwrap();

        assert_eq!(session.nullifier(), test_field_element(100));
        assert_eq!(session.action(), test_field_element(200));
    }

    #[test]
    fn test_json_roundtrip() {
        let session = SessionNullifier::new(test_field_element(1001), test_field_element(42));
        let json = serde_json::to_string(&session).unwrap();

        // Verify JSON uses the prefixed compact representation
        assert!(json.starts_with("\"snil_"));
        assert!(json.ends_with('"'));

        // Verify roundtrip
        let decoded: SessionNullifier = serde_json::from_str(&json).unwrap();
        assert_eq!(session, decoded);
    }

    #[test]
    fn test_json_format() {
        let session = SessionNullifier::new(test_field_element(1), test_field_element(2));
        let json = serde_json::to_string(&session).unwrap();

        // Should be a prefixed compact string
        let parsed: serde_json::Value = serde_json::from_str(&json).unwrap();
        assert!(parsed.is_string());
        let value = parsed.as_str().unwrap();
        assert!(value.starts_with("snil_"));
    }

    #[test]
    fn test_bytes_roundtrip() {
        let session = SessionNullifier::new(test_field_element(1001), test_field_element(42));
        let bytes = session.to_compressed_bytes();

        assert_eq!(bytes.len(), 64); // 32 + 32 bytes

        let decoded = SessionNullifier::from_compressed_bytes(&bytes).unwrap();
        assert_eq!(session, decoded);
    }

    #[test]
    fn test_bytes_use_field_element_encoding() {
        let session = SessionNullifier::new(test_field_element(1001), test_field_element(42));
        let bytes = session.to_compressed_bytes();

        let mut expected = [0u8; 64];
        expected[..32].copy_from_slice(&session.nullifier().to_be_bytes());
        expected[32..].copy_from_slice(&session.action().to_be_bytes());
        assert_eq!(bytes, expected);
    }

    #[test]
    fn test_invalid_bytes_length() {
        let too_short = vec![0u8; 63];
        let result = SessionNullifier::from_compressed_bytes(&too_short);
        assert!(result.is_err());
        assert!(result.unwrap_err().contains("Invalid length"));

        let too_long = vec![0u8; 65];
        let result = SessionNullifier::from_compressed_bytes(&too_long);
        assert!(result.is_err());
        assert!(result.unwrap_err().contains("Invalid length"));
    }

    #[test]
    fn test_default() {
        let session = SessionNullifier::default();
        assert_eq!(session.nullifier(), FieldElement::ZERO);
        assert_eq!(session.action(), FieldElement::ZERO);
    }

    #[test]
    fn test_from_tuple() {
        let nullifier = test_field_element(100);
        let action = test_field_element(200);
        let session: SessionNullifier = (nullifier, action).into();

        assert_eq!(session.nullifier(), nullifier);
        assert_eq!(session.action(), action);
    }

    #[test]
    fn test_into_u256_array() {
        let session = SessionNullifier::new(test_field_element(100), test_field_element(200));
        let arr: [U256; 2] = session.into();

        assert_eq!(arr[0], U256::from(100));
        assert_eq!(arr[1], U256::from(200));
    }
}