1use 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
12static 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
19fn 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 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 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; }
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
92fn is_safe_query_param(s: &str) -> bool {
94 !s.contains(['&', '=', '\n', '\r', '\0'])
95}
96
97pub struct IicpClient {
99 config: ClientConfig,
100 http: HttpClient,
101}
102
103impl IicpClient {
104 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 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!("®ion={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 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(); 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 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 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}