dynamo_kv_router/
tracking_hash.rs

1// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2// SPDX-License-Identifier: Apache-2.0
3
4//! Router-owned hashing for derived active-sequence tracking state.
5
6use std::fmt;
7use std::fs;
8use std::str::FromStr;
9
10use anyhow::{Context, Result, bail};
11use dynamo_tokens::SequenceHash;
12use serde::{Deserialize, Serialize};
13use zeroize::Zeroizing;
14
15use crate::config::KvRouterConfig;
16use crate::identity::RoutingPartitionRef;
17use crate::protocols::{
18    BlockHashOptions, LocalBlockHash, complete_block_count, compute_block_hash_for_seq,
19    compute_seq_hash_for_block, compute_seq_hash_for_tokens_with_seeds,
20};
21
22const KEY_SIZE: usize = 32;
23const KEYED_XXH3_V1_DOMAIN: &[u8] = b"dynamo.router.tracking-hash/keyed-xxh3-v1\0";
24
25/// Hash algorithm used only for router-derived active-sequence tracking state.
26#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
27#[serde(rename_all = "kebab-case")]
28pub enum TrackingHashAlgorithm {
29    /// Existing public XXH3 block and sequence hash construction.
30    #[default]
31    PublicXxh3V1,
32    /// Provider-keyed scope derivation with XXH3 block and chain hashing.
33    KeyedXxh3V1,
34}
35
36impl fmt::Display for TrackingHashAlgorithm {
37    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
38        formatter.write_str(match self {
39            Self::PublicXxh3V1 => "public-xxh3-v1",
40            Self::KeyedXxh3V1 => "keyed-xxh3-v1",
41        })
42    }
43}
44
45impl FromStr for TrackingHashAlgorithm {
46    type Err = String;
47
48    fn from_str(value: &str) -> Result<Self, Self::Err> {
49        match value {
50            "public-xxh3-v1" => Ok(Self::PublicXxh3V1),
51            "keyed-xxh3-v1" => Ok(Self::KeyedXxh3V1),
52            _ => Err(format!(
53                "router_tracking_hash must be public-xxh3-v1 or keyed-xxh3-v1, got {value:?}"
54            )),
55        }
56    }
57}
58
59pub(crate) fn validate_tracking_hash_options(
60    algorithm: TrackingHashAlgorithm,
61    has_key_file: bool,
62    key_id: Option<&str>,
63) -> std::result::Result<(), String> {
64    match algorithm {
65        TrackingHashAlgorithm::PublicXxh3V1 if has_key_file || key_id.is_some() => Err(
66            "router tracking key options require router_tracking_hash=keyed-xxh3-v1".to_string(),
67        ),
68        TrackingHashAlgorithm::KeyedXxh3V1 if !has_key_file => {
69            Err("keyed-xxh3-v1 requires router_tracking_key_file".to_string())
70        }
71        TrackingHashAlgorithm::KeyedXxh3V1
72            if !key_id.is_some_and(|value| !value.is_empty() && value.trim() == value) =>
73        {
74            Err("keyed-xxh3-v1 requires a nonempty router_tracking_key_id".to_string())
75        }
76        _ => Ok(()),
77    }
78}
79
80/// Stable, trusted scope shared by router instances that should produce the
81/// same derived tracking identities.
82///
83/// TODO(#11971): Evaluate deriving the static tracking-domain identity from
84/// `IndexerDomainId`. `PoolId` and `DcId` represent placement and must remain
85/// excluded from tracking identity.
86#[derive(Clone, Copy, Debug)]
87pub struct TrackingHashScope<'a> {
88    pub partition: RoutingPartitionRef<'a>,
89    pub block_size: u32,
90}
91
92/// Runtime tracking-hash state. Secret bytes are intentionally excluded from
93/// serialization and debug output.
94pub struct TrackingHashContext {
95    algorithm: TrackingHashAlgorithm,
96    key_id: Option<Box<str>>,
97    provider_key: Option<Zeroizing<[u8; KEY_SIZE]>>,
98}
99
100impl fmt::Debug for TrackingHashContext {
101    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
102        formatter
103            .debug_struct("TrackingHashContext")
104            .field("algorithm", &self.algorithm)
105            .field("key_id", &self.key_id)
106            .field(
107                "provider_key",
108                &self.provider_key.as_ref().map(|_| "[REDACTED]"),
109            )
110            .finish()
111    }
112}
113
114impl TrackingHashContext {
115    /// Load and validate the runtime context from router configuration.
116    pub fn from_config(config: &KvRouterConfig) -> Result<Self> {
117        validate_tracking_hash_options(
118            config.router_tracking_hash,
119            config.router_tracking_key_file.is_some(),
120            config.router_tracking_key_id.as_deref(),
121        )
122        .map_err(anyhow::Error::msg)?;
123
124        match config.router_tracking_hash {
125            TrackingHashAlgorithm::PublicXxh3V1 => Ok(Self {
126                algorithm: TrackingHashAlgorithm::PublicXxh3V1,
127                key_id: None,
128                provider_key: None,
129            }),
130            TrackingHashAlgorithm::KeyedXxh3V1 => {
131                let key_id = config
132                    .router_tracking_key_id
133                    .as_deref()
134                    .expect("validated keyed tracking config must have a key ID");
135                let key_path = config
136                    .router_tracking_key_file
137                    .as_ref()
138                    .expect("validated keyed tracking config must have a key file");
139                let key_bytes = Zeroizing::new(
140                    fs::read(key_path).context("failed to read router tracking key file")?,
141                );
142                let actual_size = key_bytes.len();
143                if actual_size != KEY_SIZE {
144                    bail!(
145                        "router tracking key file must contain exactly {KEY_SIZE} raw bytes; found {actual_size}"
146                    );
147                }
148                let mut provider_key = Zeroizing::new([0_u8; KEY_SIZE]);
149                provider_key.copy_from_slice(&key_bytes);
150                Ok(Self {
151                    algorithm: TrackingHashAlgorithm::KeyedXxh3V1,
152                    key_id: Some(key_id.into()),
153                    provider_key: Some(provider_key),
154                })
155            }
156        }
157    }
158
159    pub fn algorithm(&self) -> TrackingHashAlgorithm {
160        self.algorithm
161    }
162
163    /// Return the provider-managed key epoch without exposing secret material.
164    pub fn key_id(&self) -> Option<&str> {
165        self.key_id.as_deref()
166    }
167
168    /// Compute sequence identities for active tracking. Public mode may reuse
169    /// already-computed public block hashes; keyed mode must hash canonical
170    /// block bytes under its derived block seed.
171    pub(crate) fn compute_sequence_hashes(
172        &self,
173        scope: TrackingHashScope<'_>,
174        tokens: &[u32],
175        options: BlockHashOptions<'_>,
176        precomputed_public_block_hashes: Option<&[LocalBlockHash]>,
177    ) -> Vec<SequenceHash> {
178        match self.algorithm {
179            TrackingHashAlgorithm::PublicXxh3V1 => {
180                if let Some(block_hashes) = precomputed_public_block_hashes {
181                    compute_seq_hash_for_block(block_hashes)
182                } else {
183                    compute_seq_hash_for_block(&compute_block_hash_for_seq(
184                        tokens,
185                        scope.block_size,
186                        options,
187                    ))
188                }
189            }
190            TrackingHashAlgorithm::KeyedXxh3V1 => {
191                let (block_seed, chain_seed) = self.derive_seeds(scope, options);
192                compute_seq_hash_for_tokens_with_seeds(
193                    tokens,
194                    scope.block_size,
195                    options,
196                    block_seed,
197                    chain_seed,
198                )
199            }
200        }
201    }
202
203    /// Compute active-tracking identities using the configured reuse policy.
204    ///
205    /// Callers that start from raw tokens should use this method rather than
206    /// choosing between keyed/public and random identities themselves.
207    pub fn compute_sequence_hashes_for_tracking(
208        &self,
209        scope: TrackingHashScope<'_>,
210        tokens: &[u32],
211        options: BlockHashOptions<'_>,
212        assume_kv_reuse: bool,
213        precomputed_public_block_hashes: Option<&[LocalBlockHash]>,
214    ) -> Vec<SequenceHash> {
215        let num_blocks = complete_block_count(
216            tokens.len(),
217            scope.block_size,
218            options.is_eagle.unwrap_or(false),
219        );
220        if num_blocks == 0 {
221            return Vec::new();
222        }
223
224        if assume_kv_reuse {
225            self.compute_sequence_hashes(scope, tokens, options, precomputed_public_block_hashes)
226        } else {
227            (0..num_blocks).map(|_| fastrand::u64(..)).collect()
228        }
229    }
230
231    fn derive_seeds(
232        &self,
233        scope: TrackingHashScope<'_>,
234        options: BlockHashOptions<'_>,
235    ) -> (u64, u64) {
236        let key = self
237            .provider_key
238            .as_ref()
239            .expect("keyed tracking hash context must contain a provider key");
240        let mut hasher = blake3::Hasher::new_keyed(key);
241        hasher.update(KEYED_XXH3_V1_DOMAIN);
242        frame_string(&mut hasher, 1, self.key_id.as_deref().unwrap_or_default());
243        frame_string(&mut hasher, 2, scope.partition.model_name);
244        frame_string(&mut hasher, 3, scope.partition.routing_group);
245        frame_fixed(&mut hasher, 4, &scope.block_size.to_le_bytes());
246        frame_optional_string(&mut hasher, 5, normalize_optional(options.cache_namespace));
247        frame_optional_string(&mut hasher, 6, normalize_optional(options.lora_name));
248        frame_fixed(
249            &mut hasher,
250            7,
251            &[u8::from(options.is_eagle.unwrap_or(false))],
252        );
253
254        let digest = hasher.finalize();
255        let bytes = digest.as_bytes();
256        let block_seed = u64::from_le_bytes(bytes[..8].try_into().unwrap());
257        let chain_seed = u64::from_le_bytes(bytes[8..16].try_into().unwrap());
258        (block_seed, chain_seed)
259    }
260}
261
262fn normalize_optional(value: Option<&str>) -> Option<&str> {
263    value.filter(|value| !value.is_empty())
264}
265
266fn frame_string(hasher: &mut blake3::Hasher, tag: u8, value: &str) {
267    frame_fixed(hasher, tag, value.as_bytes());
268}
269
270fn frame_optional_string(hasher: &mut blake3::Hasher, tag: u8, value: Option<&str>) {
271    hasher.update(&[tag, u8::from(value.is_some())]);
272    if let Some(value) = value {
273        hasher.update(&(value.len() as u64).to_le_bytes());
274        hasher.update(value.as_bytes());
275    }
276}
277
278fn frame_fixed(hasher: &mut blake3::Hasher, tag: u8, value: &[u8]) {
279    hasher.update(&[tag]);
280    hasher.update(&(value.len() as u64).to_le_bytes());
281    hasher.update(value);
282}
283
284#[cfg(test)]
285mod tests {
286    use std::io::Write;
287
288    use tempfile::NamedTempFile;
289    use zeroize::Zeroize;
290
291    use super::*;
292    use crate::protocols::{BlockExtraInfo, BlockMmObjectInfo};
293
294    fn keyed_config(key_id: &str, key_bytes: &[u8]) -> (NamedTempFile, KvRouterConfig) {
295        let mut key_file = NamedTempFile::new().unwrap();
296        key_file.write_all(key_bytes).unwrap();
297        let config = KvRouterConfig {
298            router_tracking_hash: TrackingHashAlgorithm::KeyedXxh3V1,
299            router_tracking_key_file: Some(key_file.path().to_path_buf()),
300            router_tracking_key_id: Some(key_id.to_string()),
301            ..Default::default()
302        };
303        (key_file, config)
304    }
305
306    fn scope<'a>(model_name: &'a str, routing_group: &'a str) -> TrackingHashScope<'a> {
307        TrackingHashScope {
308            partition: RoutingPartitionRef::new(model_name, routing_group),
309            block_size: 4,
310        }
311    }
312
313    #[test]
314    fn public_mode_is_bit_compatible() {
315        let context = TrackingHashContext::from_config(&KvRouterConfig::default()).unwrap();
316        let tokens: Vec<u32> = (0..8).collect();
317        let options = BlockHashOptions {
318            cache_namespace: Some("tenant-a"),
319            lora_name: Some("adapter-a"),
320            ..Default::default()
321        };
322        let public_blocks = compute_block_hash_for_seq(&tokens, 4, options);
323
324        assert_eq!(
325            context.compute_sequence_hashes(
326                scope("model", "default"),
327                &tokens,
328                options,
329                Some(&public_blocks),
330            ),
331            compute_seq_hash_for_block(&public_blocks)
332        );
333    }
334
335    #[test]
336    fn keyed_hash_vector_pins_scope_and_chain_framing() {
337        let (_key_file, config) = keyed_config("2026-01", &[0x5a; KEY_SIZE]);
338        let context = TrackingHashContext::from_config(&config).unwrap();
339        let tokens: Vec<u32> = (0..12).collect();
340
341        let hashes = context.compute_sequence_hashes(
342            scope("model-a", "tenant-a"),
343            &tokens,
344            BlockHashOptions {
345                cache_namespace: Some("cache-a"),
346                lora_name: Some("adapter-a"),
347                ..Default::default()
348            },
349            None,
350        );
351
352        assert_eq!(
353            hashes,
354            vec![
355                4_363_769_719_052_127_296,
356                14_998_523_962_162_619_427,
357                13_920_914_207_884_994_756,
358            ]
359        );
360    }
361
362    #[test]
363    fn keyed_contexts_are_stable_and_scope_sensitive() {
364        let (_key_file, config) = keyed_config("2026-01", &[0x23; KEY_SIZE]);
365        let first = TrackingHashContext::from_config(&config).unwrap();
366        let second = TrackingHashContext::from_config(&config).unwrap();
367        let (_next_key_file, next_config) = keyed_config("2026-02", &[0x23; KEY_SIZE]);
368        let next_epoch = TrackingHashContext::from_config(&next_config).unwrap();
369        let tokens: Vec<u32> = (0..12).collect();
370        let base_options = BlockHashOptions::default();
371        let base =
372            first.compute_sequence_hashes(scope("model-a", "group-a"), &tokens, base_options, None);
373
374        assert_eq!(
375            base,
376            second.compute_sequence_hashes(
377                scope("model-a", "group-a"),
378                &tokens,
379                base_options,
380                None,
381            )
382        );
383        assert_ne!(
384            base,
385            first
386                .compute_sequence_hashes(scope("model-b", "group-a"), &tokens, base_options, None,)
387        );
388        assert_ne!(
389            base,
390            first
391                .compute_sequence_hashes(scope("model-a", "group-b"), &tokens, base_options, None,)
392        );
393        assert_ne!(
394            base,
395            first.compute_sequence_hashes(
396                TrackingHashScope {
397                    partition: RoutingPartitionRef::new("model-a", "group-a"),
398                    block_size: 3,
399                },
400                &tokens,
401                base_options,
402                None,
403            )
404        );
405        assert_ne!(
406            base,
407            next_epoch.compute_sequence_hashes(
408                scope("model-a", "group-a"),
409                &tokens,
410                base_options,
411                None,
412            )
413        );
414        assert_ne!(
415            base,
416            first.compute_sequence_hashes(
417                scope("model-a", "group-a"),
418                &tokens,
419                BlockHashOptions {
420                    cache_namespace: Some("cache-a"),
421                    ..Default::default()
422                },
423                None,
424            )
425        );
426        assert_ne!(
427            base,
428            first.compute_sequence_hashes(
429                scope("model-a", "group-a"),
430                &tokens,
431                BlockHashOptions {
432                    lora_name: Some("adapter-a"),
433                    ..Default::default()
434                },
435                None,
436            )
437        );
438        assert_ne!(
439            base,
440            first.compute_sequence_hashes(
441                scope("model-a", "group-a"),
442                &tokens,
443                BlockHashOptions {
444                    is_eagle: Some(true),
445                    ..Default::default()
446                },
447                None,
448            )
449        );
450    }
451
452    #[test]
453    fn keyed_vectors_pin_multimodal_eagle_and_partial_block_rules() {
454        let (_key_file, config) = keyed_config("2026-01", &[0x41; KEY_SIZE]);
455        let context = TrackingHashContext::from_config(&config).unwrap();
456        let tokens: Vec<u32> = (0..10).collect();
457        let mm_infos = vec![
458            Some(BlockExtraInfo {
459                mm_objects: vec![BlockMmObjectInfo {
460                    mm_hash: 42,
461                    offsets: vec![(0, 2)],
462                }],
463            }),
464            None,
465        ];
466
467        let without_mm = context.compute_sequence_hashes(
468            scope("model", "default"),
469            &tokens,
470            BlockHashOptions::default(),
471            None,
472        );
473        let with_mm = context.compute_sequence_hashes(
474            scope("model", "default"),
475            &tokens,
476            BlockHashOptions {
477                block_mm_infos: Some(&mm_infos),
478                ..Default::default()
479            },
480            None,
481        );
482        let with_eagle = context.compute_sequence_hashes(
483            scope("model", "default"),
484            &tokens,
485            BlockHashOptions {
486                is_eagle: Some(true),
487                ..Default::default()
488            },
489            None,
490        );
491
492        assert_eq!(without_mm.len(), 2);
493        assert_eq!(
494            with_mm,
495            vec![13_077_030_603_177_067_515, 12_131_634_976_806_651_614]
496        );
497        assert_eq!(
498            with_eagle,
499            vec![18_351_479_723_295_049_348, 8_577_555_336_206_814_019]
500        );
501        assert_ne!(without_mm, with_mm);
502    }
503
504    #[test]
505    fn key_loading_rejects_missing_and_malformed_files_without_exposing_bytes() {
506        let (_short_file, short_config) = keyed_config("2026-01", &[7; KEY_SIZE - 1]);
507        assert!(
508            TrackingHashContext::from_config(&short_config)
509                .unwrap_err()
510                .to_string()
511                .contains("exactly 32 raw bytes")
512        );
513
514        let (_long_file, long_config) = keyed_config("2026-01", &[8; KEY_SIZE + 1]);
515        assert!(
516            TrackingHashContext::from_config(&long_config)
517                .unwrap_err()
518                .to_string()
519                .contains("exactly 32 raw bytes")
520        );
521
522        let missing_config = KvRouterConfig {
523            router_tracking_hash: TrackingHashAlgorithm::KeyedXxh3V1,
524            router_tracking_key_file: Some("/definitely/missing/tracking-key".into()),
525            router_tracking_key_id: Some("2026-01".to_string()),
526            ..Default::default()
527        };
528        let missing_error = TrackingHashContext::from_config(&missing_config)
529            .unwrap_err()
530            .to_string();
531        assert!(missing_error.contains("failed to read router tracking key file"));
532        assert!(!missing_error.contains("/definitely/missing/tracking-key"));
533
534        let unreadable_dir = tempfile::tempdir().unwrap();
535        let unreadable_config = KvRouterConfig {
536            router_tracking_hash: TrackingHashAlgorithm::KeyedXxh3V1,
537            router_tracking_key_file: Some(unreadable_dir.path().to_path_buf()),
538            router_tracking_key_id: Some("2026-01".to_string()),
539            ..Default::default()
540        };
541        let unreadable_error = TrackingHashContext::from_config(&unreadable_config)
542            .unwrap_err()
543            .to_string();
544        assert!(unreadable_error.contains("failed to read router tracking key file"));
545        assert!(!unreadable_error.contains(&unreadable_dir.path().display().to_string()));
546
547        let (_key_file, valid_config) = keyed_config("2026-01", &[0xab; KEY_SIZE]);
548        let debug = format!(
549            "{:?}",
550            TrackingHashContext::from_config(&valid_config).unwrap()
551        );
552        assert!(debug.contains("[REDACTED]"));
553        assert!(!debug.contains("171"));
554    }
555
556    #[test]
557    fn retained_key_is_zeroizable_and_only_epoch_is_exposed() {
558        let (_key_file, config) = keyed_config("2026-01", &[0xab; KEY_SIZE]);
559        let mut context = TrackingHashContext::from_config(&config).unwrap();
560
561        assert_eq!(context.key_id(), Some("2026-01"));
562        let provider_key = context.provider_key.as_mut().unwrap();
563        provider_key.zeroize();
564        assert_eq!(provider_key.as_ref(), &[0_u8; KEY_SIZE]);
565
566        let public = TrackingHashContext::from_config(&KvRouterConfig::default()).unwrap();
567        assert_eq!(public.key_id(), None);
568    }
569
570    #[test]
571    fn config_validation_enforces_mode_specific_options() {
572        let assert_rejected_by_both = |config: &KvRouterConfig, expected: &str| {
573            assert_eq!(config.validate_config().unwrap_err(), expected);
574            assert_eq!(
575                TrackingHashContext::from_config(config)
576                    .unwrap_err()
577                    .to_string(),
578                expected
579            );
580        };
581
582        let public_with_key = KvRouterConfig {
583            router_tracking_key_file: Some("key".into()),
584            ..Default::default()
585        };
586        assert_rejected_by_both(
587            &public_with_key,
588            "router tracking key options require router_tracking_hash=keyed-xxh3-v1",
589        );
590
591        let keyed_without_options = KvRouterConfig {
592            router_tracking_hash: TrackingHashAlgorithm::KeyedXxh3V1,
593            ..Default::default()
594        };
595        assert_rejected_by_both(
596            &keyed_without_options,
597            "keyed-xxh3-v1 requires router_tracking_key_file",
598        );
599
600        let (_key_file, keyed_with_whitespace_id) = keyed_config(" 2026-01", &[1; KEY_SIZE]);
601        assert_rejected_by_both(
602            &keyed_with_whitespace_id,
603            "keyed-xxh3-v1 requires a nonempty router_tracking_key_id",
604        );
605
606        let (_key_file, keyed) = keyed_config("2026-01", &[1; KEY_SIZE]);
607        assert!(keyed.validate_config().is_ok());
608        assert!(TrackingHashContext::from_config(&keyed).is_ok());
609    }
610}
dynamo_kv_router/tracking_hash.rs

dynamo_kv_router/
tracking_hash.rs
