Skip to main content

iicp_client/
client.rs

1// SPDX-License-Identifier: Apache-2.0
2use std::sync::LazyLock;
3use std::time::Duration;
4
5use regex::Regex;
6
7use crate::confidentiality::encrypt_payload;
8use crate::errors::{IicpError, Result};
9use crate::http::{make_traceparent, HttpClient};
10use crate::types::*;
11
12// Compiled once at first use — avoid per-call allocation (fix: rust#3).
13static INTENT_RE: LazyLock<Regex> =
14    LazyLock::new(|| Regex::new(r"^urn:iicp:intent:[a-z0-9_:/-]+$").unwrap());
15
16const MAX_TIMEOUT_MS: u64 = 120_000;
17const MAX_RETRIES: u32 = 3;
18
19/// SSRF guard: return true only if url is safe to use as a node endpoint (#388).
20fn is_ssrf_safe(url: &str) -> bool {
21    let lower = url.to_lowercase();
22    let rest = if let Some(s) = lower.strip_prefix("https://") {
23        s
24    } else if let Some(s) = lower.strip_prefix("http://") {
25        s
26    } else {
27        return false;
28    };
29
30    // Extract host — handles IPv6 [addr]:port and plain host:port/path
31    let host = if rest.starts_with('[') {
32        rest.split(']')
33            .next()
34            .map(|s| s.trim_start_matches('['))
35            .unwrap_or("")
36    } else {
37        rest.split('/')
38            .next()
39            .unwrap_or("")
40            .split(':')
41            .next()
42            .unwrap_or("")
43    };
44
45    if host.is_empty() {
46        return false;
47    }
48    if matches!(host, "localhost" | "0.0.0.0" | "::1" | "::") {
49        return false;
50    }
51    const BLOCKED_SUFFIXES: &[&str] = &[
52        ".local",
53        ".internal",
54        ".lan",
55        ".test",
56        ".invalid",
57        ".localhost",
58    ];
59    if BLOCKED_SUFFIXES.iter().any(|s| host.ends_with(s)) {
60        return false;
61    }
62    // Bare hostname (no dot and no colon = Docker service name; IPv6 has colons)
63    if !host.contains('.') && !host.contains(':') {
64        return false;
65    }
66    if let Ok(addr) = host.parse::<std::net::IpAddr>() {
67        match addr {
68            std::net::IpAddr::V4(v4) => {
69                if v4.is_loopback()
70                    || v4.is_private()
71                    || v4.is_link_local()
72                    || v4.is_broadcast()
73                    || v4.is_unspecified()
74                {
75                    return false;
76                }
77                let o = v4.octets();
78                if o[0] == 100 && (64..=127).contains(&o[1]) {
79                    return false; // CGNAT 100.64/10
80                }
81            }
82            std::net::IpAddr::V6(v6) => {
83                if v6.is_loopback() || v6.is_multicast() || v6.is_unspecified() {
84                    return false;
85                }
86            }
87        }
88    }
89    true
90}
91
92/// Reject model/region strings containing query-separator or newline characters (#388 §FINDING-4-5).
93fn is_safe_query_param(s: &str) -> bool {
94    !s.contains(['&', '=', '\n', '\r', '\0'])
95}
96
97/// IICP client — discover → select → submit (ADR-016 §1).
98pub struct IicpClient {
99    config: ClientConfig,
100    http: HttpClient,
101}
102
103impl IicpClient {
104    /// Construct a client. Enforces SDK-04 (timeout_ms ≤ 120 000).
105    pub fn new(config: ClientConfig) -> Result<Self> {
106        if config.timeout_ms > MAX_TIMEOUT_MS {
107            return Err(IicpError::TimeoutTooLarge(config.timeout_ms));
108        }
109        let http = HttpClient::new(config.timeout_ms, config.node_token.clone())?;
110        Ok(Self { config, http })
111    }
112
113    /// Discover nodes for *intent* (SDK-01). Accepts an optional traceparent for propagation.
114    pub async fn discover(
115        &self,
116        intent: &str,
117        opts: Option<DiscoverOptions>,
118        traceparent: Option<&str>,
119    ) -> Result<NodeList> {
120        self.validate_intent(intent)?;
121        let opts = opts.unwrap_or_default();
122        let mut url = format!(
123            "{}/api/v1/discover?intent={}",
124            self.config.directory_url, intent
125        );
126        if let Some(region) = opts.region.as_ref().or(self.config.region.as_ref()) {
127            if is_safe_query_param(region) {
128                url.push_str(&format!("&region={region}"));
129            }
130        }
131        if let Some(model) = &opts.model {
132            if is_safe_query_param(model) {
133                url.push_str(&format!("&model={model}"));
134            }
135        }
136        if let Some(rep) = opts.min_reputation {
137            url.push_str(&format!("&min_reputation={rep}"));
138        }
139        url.push_str(&format!("&limit={}", opts.limit.unwrap_or(10)));
140        self.http.get_json(&url, traceparent).await
141    }
142
143    /// Discover → select best node → submit task (SDK-01/02).
144    /// Retries up to MAX_RETRIES on transient errors (SDK-05).
145    /// Generates one W3C traceparent shared across discover + POST (SDK-06).
146    pub async fn submit(&self, mut request: TaskRequest) -> Result<TaskResponse> {
147        self.validate_intent(&request.intent)?;
148        if request.task_id.is_empty() {
149            request.task_id = uuid::Uuid::new_v4().to_string();
150        }
151
152        let tp = make_traceparent(); // SDK-06: shared across discover + node POST
153        let nodes = self.discover(&request.intent, None, Some(&tp)).await?;
154        let node = nodes
155            .nodes
156            .into_iter()
157            .filter(|n| {
158                if !is_ssrf_safe(&n.endpoint) {
159                    eprintln!(
160                        "[iicp-client] SSRF guard: skipping node {} — endpoint {} is not publicly routable",
161                        &n.node_id[..n.node_id.len().min(8)],
162                        n.endpoint
163                    );
164                    false
165                } else {
166                    true
167                }
168            })
169            .find(|n| n.available)
170            .ok_or_else(|| IicpError::NoNodes {
171                intent: request.intent.clone(),
172            })?;
173
174        // IICP-CX S.16 §5: encrypt payload when use_confidentiality=true + node has cx_public_key
175        let body: serde_json::Value = if self.config.use_confidentiality {
176            if let Some(ref cx_key) = node.cx_public_key {
177                let iicp_conf =
178                    encrypt_payload(&request.payload, cx_key, &request.task_id, &request.intent)?;
179                let mut body = serde_json::to_value(&request)?;
180                if let Some(obj) = body.as_object_mut() {
181                    obj.remove("payload");
182                    obj.insert("iicp_conf".to_string(), serde_json::to_value(iicp_conf)?);
183                }
184                body
185            } else {
186                serde_json::to_value(&request)?
187            }
188        } else {
189            serde_json::to_value(&request)?
190        };
191
192        let mut last_err: Option<IicpError> = None;
193        for attempt in 0..MAX_RETRIES {
194            match self
195                .http
196                .post_json(
197                    &format!("{}/v1/task", node.endpoint),
198                    &body,
199                    None,
200                    Some(&tp),
201                )
202                .await
203            {
204                Ok(resp) => return Ok(resp),
205                Err(e) if e.is_transient() && attempt < MAX_RETRIES - 1 => {
206                    tokio::time::sleep(Duration::from_millis(200 * 2u64.pow(attempt))).await;
207                    last_err = Some(e);
208                }
209                Err(e) => return Err(e),
210            }
211        }
212        Err(last_err.unwrap())
213    }
214
215    /// Discover → select best LLM node → submit chat task (SDK-02).
216    pub async fn chat(
217        &self,
218        messages: Vec<ChatMessage>,
219        opts: Option<ChatOptions>,
220    ) -> Result<ChatResponse> {
221        let opts = opts.unwrap_or_default();
222        let mut payload = serde_json::json!({ "messages": messages });
223        if let Some(ref model) = opts.model {
224            payload["model"] = serde_json::Value::String(model.clone());
225        }
226        if let Some(temp) = opts.temperature {
227            payload["temperature"] = serde_json::json!(temp);
228        }
229        let request = TaskRequest {
230            task_id: uuid::Uuid::new_v4().to_string(),
231            intent: "urn:iicp:intent:llm:chat:v1".into(),
232            payload,
233            constraints: Some(TaskConstraints {
234                timeout_ms: opts.timeout_ms,
235                max_tokens: opts.max_tokens,
236                model: opts.model,
237            }),
238            auth: None,
239        };
240        let task_resp = self.submit(request).await?;
241        let node_id = task_resp.metrics.as_ref().and_then(|m| m.node_id.clone());
242        let task_id = task_resp.task_id.clone();
243        let result = task_resp.result.ok_or_else(|| IicpError::Protocol {
244            code: "no_result".into(),
245            message: "Node returned task without a result payload".into(),
246            status: 200,
247        })?;
248        let mut resp: ChatResponse = serde_json::from_value(result)?;
249        resp.task_id = task_id;
250        resp.node_id = node_id;
251        Ok(resp)
252    }
253
254    fn validate_intent(&self, intent: &str) -> Result<()> {
255        if !INTENT_RE.is_match(intent) {
256            return Err(IicpError::InvalidIntent(intent.into()));
257        }
258        Ok(())
259    }
260}