Skip to main content

wots_rhdl/
tag.rs

1//! Opaque SHA-lane tag encoding for fused WOTS tasks.
2//!
3//! The compression lane must preserve all 32 bits without interpreting them.
4//! This module is the sole authority for their layout:
5//!
6//! | Bits | Width | Meaning |
7//! |---|---:|---|
8//! | 0..=3 | 4 | scheduler context |
9//! | 4..=6 | 3 | [`TaskKind`] |
10//! | 7..=13 | 7 | WOTS segment or mask index |
11//! | 14..=17 | 4 | chain step |
12//! | 18..=23 | 6 | SHA/public-key block index |
13//! | 24..=31 | 8 | context generation/sequence |
14
15use core::fmt;
16
17use rhdl::prelude::*;
18
19/// Number of bits in the opaque compression-lane tag.
20pub const TASK_TAG_BITS: usize = 32;
21
22/// Shift of the context field in a task tag.
23pub const CONTEXT_SHIFT: u32 = 0;
24/// Shift of the task-kind field in a task tag.
25pub const KIND_SHIFT: u32 = 4;
26/// Shift of the segment field in a task tag.
27pub const SEGMENT_SHIFT: u32 = 7;
28/// Shift of the chain-step field in a task tag.
29pub const CHAIN_STEP_SHIFT: u32 = 14;
30/// Shift of the block-index field in a task tag.
31pub const BLOCK_INDEX_SHIFT: u32 = 18;
32/// Shift of the generation field in a task tag.
33pub const GENERATION_SHIFT: u32 = 24;
34
35/// Compression task classes in their fixed three-bit wire encoding.
36#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
37#[repr(u8)]
38pub enum TaskKind {
39    /// Hash the 32-byte private seed to derive the private key.
40    Seed = 0,
41    /// Derive the public seed with PRF index zero.
42    PublicSeed = 1,
43    /// Derive one of the 16 observable randomization masks.
44    Mask = 2,
45    /// Derive one segment PRF using private-key indices 1 through 67.
46    SegmentPrf = 3,
47    /// Compress the 64-byte secret-segment preimage's data block.
48    SecretData = 4,
49    /// Compress the fixed padding block for a 64-byte secret preimage.
50    SecretPadding = 5,
51    /// Advance one segment by one randomized chain step.
52    Chain = 6,
53    /// Compress one of the 34 serialized public-key endpoint blocks.
54    PublicKey = 7,
55}
56
57impl TaskKind {
58    /// Returns the fixed three-bit numeric encoding.
59    #[must_use]
60    pub const fn code(self) -> u8 {
61        self as u8
62    }
63
64    /// Decodes a three-bit numeric encoding.
65    #[must_use]
66    pub const fn from_code(code: u8) -> Option<Self> {
67        match code {
68            0 => Some(Self::Seed),
69            1 => Some(Self::PublicSeed),
70            2 => Some(Self::Mask),
71            3 => Some(Self::SegmentPrf),
72            4 => Some(Self::SecretData),
73            5 => Some(Self::SecretPadding),
74            6 => Some(Self::Chain),
75            7 => Some(Self::PublicKey),
76            _ => None,
77        }
78    }
79}
80
81/// Unpacked fields consumed and produced by RHDL tag kernels.
82#[derive(Clone, Copy, Debug, Default, Digital, Eq, PartialEq)]
83pub struct TagFields {
84    /// Scheduler context, in `0..16`.
85    pub context: b4,
86    /// Three-bit [`TaskKind`] encoding.
87    pub kind: b3,
88    /// Segment `0..66`, or mask index `0..15` for mask tasks.
89    pub segment: b7,
90    /// Chain step `1..15` for chain tasks and zero otherwise.
91    pub chain_step: b4,
92    /// Public-key block `0..33`, secret padding block one, or zero.
93    pub block_index: b6,
94    /// Eight-bit context generation used to reject stale responses.
95    pub generation: b8,
96}
97
98/// Pack all tag fields into the fixed 32-bit SHA-lane representation.
99#[kernel]
100pub fn encode_tag_kernel(fields: TagFields) -> b32 {
101    let context: b32 = fields.context.resize();
102    let kind: b32 = fields.kind.resize();
103    let segment: b32 = fields.segment.resize();
104    let chain_step: b32 = fields.chain_step.resize();
105    let block_index: b32 = fields.block_index.resize();
106    let generation: b32 = fields.generation.resize();
107    context
108        | (kind << 4)
109        | (segment << 7)
110        | (chain_step << 14)
111        | (block_index << 18)
112        | (generation << 24)
113}
114
115/// Unpack an opaque SHA-lane tag without changing any bit.
116#[kernel]
117pub fn decode_tag_kernel(tag: b32) -> TagFields {
118    TagFields {
119        context: tag.resize(),
120        kind: (tag >> 4).resize(),
121        segment: (tag >> 7).resize(),
122        chain_step: (tag >> 14).resize(),
123        block_index: (tag >> 18).resize(),
124        generation: (tag >> 24).resize(),
125    }
126}
127
128/// Return whether a tag uses the canonical bounds for its task class.
129///
130/// Unused fields must be zero. This prevents semantically ambiguous encodings
131/// from reaching a context store or being accepted as a response.
132#[kernel]
133#[allow(clippy::match_same_arms)] // Separate wire values keep the encoding auditable.
134pub fn tag_is_well_formed_kernel(tag: b32) -> bool {
135    let fields = decode_tag_kernel(tag);
136    let no_segment = fields.segment == b7(0);
137    let no_step = fields.chain_step == b4(0);
138    let no_block = fields.block_index == b6(0);
139    match fields.kind {
140        Bits::<3>(0) => no_segment && no_step && no_block,
141        Bits::<3>(1) => no_segment && no_step && no_block,
142        Bits::<3>(2) => fields.segment < b7(16) && no_step && no_block,
143        Bits::<3>(3) => fields.segment < b7(67) && no_step && no_block,
144        Bits::<3>(4) => fields.segment < b7(67) && no_step && no_block,
145        Bits::<3>(5) => fields.segment < b7(67) && no_step && fields.block_index == b6(1),
146        Bits::<3>(6) => {
147            fields.segment < b7(67)
148                && fields.chain_step > b4(0)
149                && fields.chain_step <= b4(15)
150                && no_block
151        }
152        _ => no_segment && no_step && fields.block_index < b6(34),
153    }
154}
155
156/// Host-side decoded representation of an opaque task tag.
157#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
158pub struct DecodedTaskTag {
159    /// Scheduler context.
160    pub context: u8,
161    /// Task class.
162    pub kind: TaskKind,
163    /// Segment or mask index.
164    pub segment: u8,
165    /// Chain step.
166    pub chain_step: u8,
167    /// SHA/public-key block index.
168    pub block_index: u8,
169    /// Context generation.
170    pub generation: u8,
171}
172
173impl DecodedTaskTag {
174    /// Construct canonical tag fields, rejecting invalid class-specific bounds.
175    ///
176    /// # Errors
177    ///
178    /// Returns [`TaskTagError`] if `context` exceeds four bits or any field is
179    /// noncanonical for `kind`.
180    pub fn new(
181        context: u8,
182        kind: TaskKind,
183        segment: u8,
184        chain_step: u8,
185        block_index: u8,
186        generation: u8,
187    ) -> Result<Self, TaskTagError> {
188        let candidate = Self {
189            context,
190            kind,
191            segment,
192            chain_step,
193            block_index,
194            generation,
195        };
196        if context >= 16 || !candidate.is_well_formed() {
197            return Err(TaskTagError::NonCanonical(candidate.encode_unchecked()));
198        }
199        Ok(candidate)
200    }
201
202    /// Encode these fields into the opaque 32-bit lane tag.
203    #[must_use]
204    pub fn encode(self) -> u32 {
205        self.encode_unchecked()
206    }
207
208    /// Return whether all class-specific fields use their canonical bounds.
209    #[must_use]
210    pub const fn is_well_formed(self) -> bool {
211        let no_segment = self.segment == 0;
212        let no_step = self.chain_step == 0;
213        let no_block = self.block_index == 0;
214        match self.kind {
215            TaskKind::Seed | TaskKind::PublicSeed => no_segment && no_step && no_block,
216            TaskKind::Mask => self.segment < 16 && no_step && no_block,
217            TaskKind::SegmentPrf | TaskKind::SecretData => self.segment < 67 && no_step && no_block,
218            TaskKind::SecretPadding => self.segment < 67 && no_step && self.block_index == 1,
219            TaskKind::Chain => {
220                self.segment < 67 && self.chain_step > 0 && self.chain_step < 16 && no_block
221            }
222            TaskKind::PublicKey => no_segment && no_step && self.block_index < 34,
223        }
224    }
225
226    fn encode_unchecked(self) -> u32 {
227        u32::from(self.context)
228            | u32::from(self.kind.code()) << KIND_SHIFT
229            | u32::from(self.segment) << SEGMENT_SHIFT
230            | u32::from(self.chain_step) << CHAIN_STEP_SHIFT
231            | u32::from(self.block_index) << BLOCK_INDEX_SHIFT
232            | u32::from(self.generation) << GENERATION_SHIFT
233    }
234}
235
236/// Error returned while decoding or constructing an opaque task tag.
237#[derive(Clone, Copy, Debug, Eq, PartialEq)]
238pub enum TaskTagError {
239    /// The raw tag contains fields outside the canonical bounds for its class.
240    NonCanonical(u32),
241}
242
243impl fmt::Display for TaskTagError {
244    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
245        match self {
246            Self::NonCanonical(raw) => write!(formatter, "noncanonical WOTS task tag {raw:#010x}"),
247        }
248    }
249}
250
251impl std::error::Error for TaskTagError {}
252
253/// Decode and validate a host-side opaque task tag.
254///
255/// # Errors
256///
257/// Returns [`TaskTagError`] if any class-specific field is noncanonical.
258pub fn decode_task_tag(raw: u32) -> Result<DecodedTaskTag, TaskTagError> {
259    let kind_code = ((raw >> KIND_SHIFT) & 0x07).to_le_bytes()[0];
260    let decoded = DecodedTaskTag {
261        context: (raw & 0x0f).to_le_bytes()[0],
262        kind: match kind_code {
263            0 => TaskKind::Seed,
264            1 => TaskKind::PublicSeed,
265            2 => TaskKind::Mask,
266            3 => TaskKind::SegmentPrf,
267            4 => TaskKind::SecretData,
268            5 => TaskKind::SecretPadding,
269            6 => TaskKind::Chain,
270            _ => TaskKind::PublicKey,
271        },
272        segment: ((raw >> SEGMENT_SHIFT) & 0x7f).to_le_bytes()[0],
273        chain_step: ((raw >> CHAIN_STEP_SHIFT) & 0x0f).to_le_bytes()[0],
274        block_index: ((raw >> BLOCK_INDEX_SHIFT) & 0x3f).to_le_bytes()[0],
275        generation: (raw >> GENERATION_SHIFT).to_le_bytes()[0],
276    };
277    if decoded.is_well_formed() {
278        Ok(decoded)
279    } else {
280        Err(TaskTagError::NonCanonical(raw))
281    }
282}
283
284/// Convert RHDL tag fields to their host-side numeric representation.
285#[must_use]
286pub fn software_fields(fields: TagFields) -> DecodedTaskTag {
287    let kind = fields.kind.raw().to_le_bytes()[0];
288    DecodedTaskTag {
289        context: fields.context.raw().to_le_bytes()[0],
290        kind: match kind {
291            0 => TaskKind::Seed,
292            1 => TaskKind::PublicSeed,
293            2 => TaskKind::Mask,
294            3 => TaskKind::SegmentPrf,
295            4 => TaskKind::SecretData,
296            5 => TaskKind::SecretPadding,
297            6 => TaskKind::Chain,
298            _ => TaskKind::PublicKey,
299        },
300        segment: fields.segment.raw().to_le_bytes()[0],
301        chain_step: fields.chain_step.raw().to_le_bytes()[0],
302        block_index: fields.block_index.raw().to_le_bytes()[0],
303        generation: fields.generation.raw().to_le_bytes()[0],
304    }
305}
306
307/// Convert host-side decoded fields to the RHDL representation.
308#[must_use]
309pub fn hardware_fields(fields: DecodedTaskTag) -> TagFields {
310    TagFields {
311        context: b4(u128::from(fields.context)),
312        kind: b3(u128::from(fields.kind.code())),
313        segment: b7(u128::from(fields.segment)),
314        chain_step: b4(u128::from(fields.chain_step)),
315        block_index: b6(u128::from(fields.block_index)),
316        generation: b8(u128::from(fields.generation)),
317    }
318}
319
320const _: () = {
321    assert!(CONTEXT_SHIFT == 0);
322    assert!(GENERATION_SHIFT + 8 == 32);
323};