Skip to main content

aft/
harness.rs

1use serde::de::{self, Visitor};
2use serde::{Deserialize, Deserializer, Serialize, Serializer};
3use std::fmt;
4use std::str::FromStr;
5
6#[derive(Debug, Clone, PartialEq, Eq, Hash)]
7pub enum Harness {
8    Opencode,
9    Pi,
10    Runner,
11    Mcp { client: String },
12    Fed { fingerprint: String },
13}
14
15impl Harness {
16    pub fn storage_segment(&self) -> String {
17        match self {
18            Harness::Opencode => "opencode".to_string(),
19            Harness::Pi => "pi".to_string(),
20            Harness::Runner => "runner".to_string(),
21            Harness::Mcp { client } => format!("mcp--{}", sanitize_client(client)),
22            Harness::Fed { fingerprint } => format!(
23                "fed--{}--{}",
24                &fingerprint[..FED_SLUG_READABLE_HEX_LEN],
25                hash_hex_prefix(fingerprint, FED_SLUG_HASH_HEX_LEN)
26            ),
27        }
28    }
29
30    pub fn wire_label(&self) -> String {
31        match self {
32            Harness::Opencode => "opencode".to_string(),
33            Harness::Pi => "pi".to_string(),
34            Harness::Runner => "runner".to_string(),
35            Harness::Mcp { client } => format!("mcp:{client}"),
36            Harness::Fed { fingerprint } => format!("fed:{fingerprint}"),
37        }
38    }
39}
40
41/// Max length of the readable (pre-hash) slug portion. The full segment is
42/// `mcp--<readable>--<32 hex>`, so the readable part is capped to keep directory
43/// names bounded while the hash guarantees uniqueness.
44const MCP_SLUG_READABLE_MAX: usize = 40;
45const MCP_SLUG_HASH_HEX_LEN: usize = 32;
46const FED_FINGERPRINT_MIN_HEX_LEN: usize = 32;
47const FED_FINGERPRINT_MAX_HEX_LEN: usize = 64;
48const FED_SLUG_READABLE_HEX_LEN: usize = 16;
49const FED_SLUG_HASH_HEX_LEN: usize = 8;
50
51fn hash_hex_prefix(raw: &str, hex_len: usize) -> String {
52    let hash = blake3::hash(raw.as_bytes()).to_hex();
53    hash.as_str()[..hex_len].to_string()
54}
55
56/// Build the storage slug for an MCP client. The readable portion is a
57/// sanitized, length-capped rendering of the raw client; a short hash of the
58/// RAW (un-sanitized) client is appended so that distinct clients that sanitize
59/// to the same readable string (e.g. `a/b`, `a:b`, `a b`, casing variants, or
60/// non-ASCII that collapses to `unknown`) still get distinct directories. The
61/// hash is over the raw bytes, so it is collision-resistant where the readable
62/// slug is not.
63fn sanitize_client(client: &str) -> String {
64    let lower = client.to_ascii_lowercase();
65    let mut out = String::with_capacity(lower.len());
66    let mut last_was_dash = false;
67    for ch in lower.chars() {
68        let keep = ch.is_ascii_alphanumeric() || matches!(ch, '.' | '_' | '-');
69        if keep {
70            out.push(ch);
71            last_was_dash = false;
72        } else if !last_was_dash {
73            out.push('-');
74            last_was_dash = true;
75        }
76    }
77    let trimmed = out.trim_matches(|c| c == '-' || c == '.');
78    let mut readable = if trimmed.is_empty() {
79        "unknown".to_string()
80    } else {
81        trimmed.to_string()
82    };
83    if readable.len() > MCP_SLUG_READABLE_MAX {
84        readable.truncate(MCP_SLUG_READABLE_MAX);
85        // Truncation can leave a trailing separator; trim it for tidiness.
86        readable = readable.trim_end_matches(['-', '.']).to_string();
87        if readable.is_empty() {
88            readable = "unknown".to_string();
89        }
90    }
91
92    // A 128-bit hash suffix prevents hostile same-readable slugs from sharing
93    // storage while keeping directory names short enough for common filesystems.
94    format!(
95        "{readable}--{}",
96        hash_hex_prefix(client, MCP_SLUG_HASH_HEX_LEN)
97    )
98}
99
100fn is_lower_hex(value: &str) -> bool {
101    value
102        .bytes()
103        .all(|byte| matches!(byte, b'0'..=b'9' | b'a'..=b'f'))
104}
105
106fn parse_fed_harness(value: &str) -> Result<Harness, String> {
107    let fingerprint = &value[4..];
108    if !(FED_FINGERPRINT_MIN_HEX_LEN..=FED_FINGERPRINT_MAX_HEX_LEN).contains(&fingerprint.len())
109        || !is_lower_hex(fingerprint)
110    {
111        return Err(format!(
112            "unsupported harness '{value}'; fed fingerprint must be 32-64 lowercase hex characters"
113        ));
114    }
115    Ok(Harness::Fed {
116        fingerprint: fingerprint.to_string(),
117    })
118}
119
120impl Serialize for Harness {
121    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
122    where
123        S: Serializer,
124    {
125        serializer.serialize_str(&self.wire_label())
126    }
127}
128
129impl<'de> Deserialize<'de> for Harness {
130    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
131    where
132        D: Deserializer<'de>,
133    {
134        struct HarnessVisitor;
135
136        impl<'de> Visitor<'de> for HarnessVisitor {
137            type Value = Harness;
138
139            fn expecting(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
140                formatter.write_str(
141                    "a harness string: 'opencode', 'pi', 'runner', 'mcp:<client>', or 'fed:<fingerprint>'",
142                )
143            }
144
145            fn visit_str<E>(self, value: &str) -> Result<Self::Value, E>
146            where
147                E: de::Error,
148            {
149                Harness::from_str(value).map_err(E::custom)
150            }
151        }
152
153        deserializer.deserialize_str(HarnessVisitor)
154    }
155}
156
157impl fmt::Display for Harness {
158    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
159        f.write_str(&self.wire_label())
160    }
161}
162
163impl std::str::FromStr for Harness {
164    type Err = String;
165
166    fn from_str(value: &str) -> Result<Self, Self::Err> {
167        match value {
168            "opencode" => Ok(Harness::Opencode),
169            "pi" => Ok(Harness::Pi),
170            "runner" => Ok(Harness::Runner),
171            other if other.starts_with("mcp:") => {
172                let client = &other[4..];
173                if client.is_empty() {
174                    Err(
175                        "unsupported harness 'mcp:'; mcp client name must be non-empty".to_string(),
176                    )
177                } else {
178                    Ok(Harness::Mcp {
179                        client: client.to_string(),
180                    })
181                }
182            }
183            other if other.starts_with("fed:") => parse_fed_harness(other),
184            other => Err(format!(
185                "unsupported harness '{other}'; expected 'opencode', 'pi', 'runner', 'mcp:<client>', or 'fed:<fingerprint>'"
186            )),
187        }
188    }
189}
190
191#[cfg(test)]
192mod tests {
193    use super::{sanitize_client, Harness};
194    use std::str::FromStr;
195
196    #[test]
197    fn harness_enum_serde_roundtrip() {
198        assert_eq!(
199            serde_json::to_string(&Harness::Opencode).unwrap(),
200            "\"opencode\""
201        );
202        assert_eq!(serde_json::to_string(&Harness::Pi).unwrap(), "\"pi\"");
203
204        assert_eq!(
205            serde_json::from_str::<Harness>("\"opencode\"").unwrap(),
206            Harness::Opencode
207        );
208        assert_eq!(
209            serde_json::from_str::<Harness>("\"pi\"").unwrap(),
210            Harness::Pi
211        );
212        assert!(serde_json::from_str::<Harness>("\"claude_code\"").is_err());
213    }
214
215    #[test]
216    fn opencode_pi_storage_segment_unchanged() {
217        assert_eq!(Harness::Opencode.storage_segment(), "opencode");
218        assert_eq!(Harness::Pi.storage_segment(), "pi");
219    }
220
221    #[test]
222    fn runner_round_trips() {
223        assert_eq!(Harness::from_str("runner").unwrap(), Harness::Runner);
224        assert_eq!(Harness::Runner.storage_segment(), "runner");
225        assert_eq!(
226            serde_json::to_string(&Harness::Runner).unwrap(),
227            "\"runner\""
228        );
229        assert_eq!(
230            serde_json::from_str::<Harness>("\"runner\"").unwrap(),
231            Harness::Runner
232        );
233    }
234
235    #[test]
236    fn mcp_round_trips() {
237        let h = Harness::Mcp {
238            client: "claude-code".to_string(),
239        };
240        assert_eq!(serde_json::to_string(&h).unwrap(), "\"mcp:claude-code\"");
241        assert_eq!(
242            serde_json::from_str::<Harness>("\"mcp:claude-code\"").unwrap(),
243            h
244        );
245        assert_eq!(
246            Harness::from_str("mcp:cursor").unwrap(),
247            Harness::Mcp {
248                client: "cursor".to_string(),
249            }
250        );
251        assert!(Harness::from_str("mcp:").is_err());
252    }
253
254    #[test]
255    fn fed_round_trips_and_rejects_malformed_fingerprints() {
256        let fingerprint64 = "0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef";
257        let h = Harness::Fed {
258            fingerprint: fingerprint64.to_string(),
259        };
260        assert_eq!(
261            serde_json::to_string(&h).unwrap(),
262            format!("\"fed:{fingerprint64}\"")
263        );
264        assert_eq!(
265            serde_json::from_str::<Harness>(&format!("\"fed:{fingerprint64}\"")).unwrap(),
266            h
267        );
268
269        let fingerprint32 = "0123456789abcdef0123456789abcdef";
270        assert_eq!(
271            Harness::from_str(&format!("fed:{fingerprint32}")).unwrap(),
272            Harness::Fed {
273                fingerprint: fingerprint32.to_string(),
274            }
275        );
276
277        for invalid in [
278            "fed:",
279            "fed:0123456789ABCDEF0123456789abcdef0123456789abcdef0123456789abcdef",
280            "fed:0123456789abcdef0123456789abcde",
281            "fed:0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef0",
282            "fed:0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdeg",
283            "fed:0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef:extra",
284        ] {
285            assert!(
286                Harness::from_str(invalid).is_err(),
287                "{invalid:?} must be rejected"
288            );
289        }
290    }
291
292    #[test]
293    fn storage_segment_hostile_clients_are_path_safe() {
294        let cases = ["../../etc", "a/b", r"a\b", "a:b", "", "Claude.Code"];
295        for client in cases {
296            let seg = Harness::Mcp {
297                client: client.to_string(),
298            }
299            .storage_segment();
300            assert!(
301                !seg.is_empty(),
302                "segment must be non-empty for client {client:?}"
303            );
304            assert!(
305                !seg.contains(['/', '\\', ':']),
306                "segment {seg:?} must not contain path separators for client {client:?}"
307            );
308            assert!(
309                !seg.contains(".."),
310                "segment {seg:?} must not contain '..' for client {client:?}"
311            );
312            assert!(
313                seg.starts_with("mcp--"),
314                "segment {seg:?} must use mcp-- prefix"
315            );
316        }
317        // Readable portion preserved, hash suffix appended.
318        let claude = Harness::Mcp {
319            client: "Claude.Code".to_string(),
320        }
321        .storage_segment();
322        assert!(
323            claude.starts_with("mcp--claude.code--"),
324            "expected readable slug with hash suffix, got {claude:?}"
325        );
326        // Empty client → readable "unknown" plus a (stable) hash of empty bytes.
327        let empty = sanitize_client("");
328        assert!(
329            empty.starts_with("unknown--"),
330            "empty client must render unknown-- plus hash, got {empty:?}"
331        );
332    }
333
334    #[test]
335    fn storage_segment_disambiguates_clients_that_sanitize_to_same_slug() {
336        // a/b, a:b, a b, A-B all collapse to the readable slug "a-b" but are
337        // DISTINCT clients — the raw-bytes hash suffix must keep their storage
338        // directories distinct so two different MCP clients never share state.
339        let seg = |c: &str| {
340            Harness::Mcp {
341                client: c.to_string(),
342            }
343            .storage_segment()
344        };
345        let variants = [seg("a/b"), seg("a:b"), seg("a b"), seg("A-B")];
346        for s in &variants {
347            assert!(
348                s.starts_with("mcp--a-b--"),
349                "expected shared readable slug a-b, got {s:?}"
350            );
351            let (_readable, suffix) = s.rsplit_once("--").expect("hash suffix");
352            assert_eq!(
353                suffix.len(),
354                super::MCP_SLUG_HASH_HEX_LEN,
355                "hash suffix must carry 128 bits of disambiguation: {s:?}"
356            );
357            assert!(
358                suffix.chars().all(|ch| ch.is_ascii_hexdigit()),
359                "hash suffix must be hex: {s:?}"
360            );
361        }
362        let unique: std::collections::HashSet<_> = variants.iter().collect();
363        assert_eq!(
364            unique.len(),
365            variants.len(),
366            "distinct clients must get distinct storage segments: {variants:?}"
367        );
368
369        // Same raw client → same segment (deterministic, stable across calls).
370        assert_eq!(seg("cursor"), seg("cursor"));
371
372        // Very long client: readable portion is capped, segment stays bounded.
373        let long = seg(&"x".repeat(500));
374        assert!(
375            long.len()
376                <= "mcp--".len()
377                    + super::MCP_SLUG_READABLE_MAX
378                    + "--".len()
379                    + super::MCP_SLUG_HASH_HEX_LEN,
380            "long client segment must be length-bounded, got len {}",
381            long.len()
382        );
383    }
384
385    #[test]
386    fn fed_storage_segment_is_path_safe_bounded_and_disambiguated() {
387        let prefix = "0123456789abcdef";
388        let fingerprint_a = format!("{prefix}{}", "0".repeat(48));
389        let fingerprint_b = format!("{prefix}{}", "f".repeat(48));
390        let seg = |fingerprint: &str| {
391            Harness::Fed {
392                fingerprint: fingerprint.to_string(),
393            }
394            .storage_segment()
395        };
396        let seg_a = seg(&fingerprint_a);
397        let seg_b = seg(&fingerprint_b);
398        for segment in [&seg_a, &seg_b] {
399            assert!(
400                segment.starts_with(&format!("fed--{prefix}--")),
401                "fed segment must keep the 16-hex readable prefix, got {segment:?}"
402            );
403            assert!(
404                !segment.contains(['/', '\\', ':']),
405                "fed segment must be path-safe, got {segment:?}"
406            );
407            let (_readable, suffix) = segment.rsplit_once("--").expect("hash suffix");
408            assert_eq!(
409                suffix.len(),
410                super::FED_SLUG_HASH_HEX_LEN,
411                "fed hash suffix must use 8 hex chars, got {segment:?}"
412            );
413            assert!(
414                suffix.chars().all(|ch| ch.is_ascii_hexdigit()),
415                "fed hash suffix must be hex, got {segment:?}"
416            );
417            assert!(
418                segment.len()
419                    <= "fed--".len()
420                        + super::FED_SLUG_READABLE_HEX_LEN
421                        + "--".len()
422                        + super::FED_SLUG_HASH_HEX_LEN,
423                "fed segment must be length-bounded, got len {}",
424                segment.len()
425            );
426        }
427        assert_ne!(
428            seg_a, seg_b,
429            "distinct fed fingerprints that share a readable prefix must not share storage"
430        );
431        assert_eq!(seg(&fingerprint_a), seg_a);
432    }
433}