agent_core/core/auth/
aws_bedrock.rs

1//! Provider-local IAM Identity Center and Bedrock broker boundary.
2//! The transport trait is intentionally typed: it cannot sign arbitrary URLs or vend credentials.
3use super::cloud::{AwsBedrockConfig, InvokeRequest};
4use async_trait::async_trait;
5use futures::Stream;
6use serde::Serialize;
7use sha2::{Digest, Sha256};
8use std::{collections::BTreeMap, fmt, pin::Pin};
9
10#[derive(Clone)]
11pub struct RegisteredClient {
12    id: String,
13    secret: String,
14    pub expires_at: u64,
15}
16impl RegisteredClient {
17    pub fn new(id: impl Into<String>, secret: impl Into<String>, expires_at: u64) -> Self {
18        Self {
19            id: id.into(),
20            secret: secret.into(),
21            expires_at,
22        }
23    }
24    pub fn id(&self) -> &str {
25        &self.id
26    }
27    pub fn secret(&self) -> &str {
28        &self.secret
29    }
30}
31impl fmt::Debug for RegisteredClient {
32    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
33        f.debug_struct("RegisteredClient")
34            .field("id", &"[REDACTED]")
35            .field("secret", &"[REDACTED]")
36            .field("expires_at", &self.expires_at)
37            .finish()
38    }
39}
40#[derive(Clone)]
41pub struct DeviceAuthorization {
42    device_code: String,
43    pub user_code: String,
44    pub verification_uri: String,
45    pub interval: u64,
46    pub expires_in: u64,
47}
48impl DeviceAuthorization {
49    pub fn new(
50        d: impl Into<String>,
51        u: impl Into<String>,
52        v: impl Into<String>,
53        interval: u64,
54        expires_in: u64,
55    ) -> Self {
56        Self {
57            device_code: d.into(),
58            user_code: u.into(),
59            verification_uri: v.into(),
60            interval,
61            expires_in,
62        }
63    }
64    pub fn device_code(&self) -> &str {
65        &self.device_code
66    }
67}
68impl fmt::Debug for DeviceAuthorization {
69    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
70        f.write_str("DeviceAuthorization([REDACTED])")
71    }
72}
73#[derive(Clone)]
74pub struct SsoToken {
75    access: String,
76    refresh: Option<String>,
77    pub expires_in: u64,
78}
79impl SsoToken {
80    pub fn new(a: impl Into<String>, r: Option<&str>, e: u64) -> Self {
81        Self {
82            access: a.into(),
83            refresh: r.map(str::to_owned),
84            expires_in: e,
85        }
86    }
87    pub fn access(&self) -> &str {
88        &self.access
89    }
90    pub fn refresh(&self) -> Option<&str> {
91        self.refresh.as_deref()
92    }
93}
94impl fmt::Debug for SsoToken {
95    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
96        f.write_str("SsoToken([REDACTED])")
97    }
98}
99#[derive(Clone)]
100pub struct RoleCredentials {
101    access_key: String,
102    secret_key: String,
103    session_token: String,
104    pub expires_at: u64,
105}
106impl RoleCredentials {
107    pub fn new(a: impl Into<String>, s: impl Into<String>, t: impl Into<String>, e: u64) -> Self {
108        Self {
109            access_key: a.into(),
110            secret_key: s.into(),
111            session_token: t.into(),
112            expires_at: e,
113        }
114    }
115    pub fn access_key(&self) -> &str {
116        &self.access_key
117    }
118    pub fn secret_key(&self) -> &str {
119        &self.secret_key
120    }
121    pub fn session_token(&self) -> &str {
122        &self.session_token
123    }
124}
125impl fmt::Debug for RoleCredentials {
126    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
127        f.write_str("RoleCredentials([REDACTED])")
128    }
129}
130#[derive(Debug, Clone)]
131pub struct Account {
132    pub id: String,
133    pub name: String,
134}
135#[derive(Debug, Clone)]
136pub struct Role {
137    pub name: String,
138}
139#[derive(Debug, Clone)]
140pub struct FoundationModel {
141    pub id: String,
142    pub name: String,
143    pub supports_text: bool,
144    pub supports_converse: bool,
145}
146impl FoundationModel {
147    pub fn new(i: impl Into<String>, n: impl Into<String>, t: bool, c: bool) -> Self {
148        Self {
149            id: i.into(),
150            name: n.into(),
151            supports_text: t,
152            supports_converse: c,
153        }
154    }
155}
156#[derive(Debug, Clone, Serialize)]
157pub struct CatalogEntry {
158    pub id: String,
159    pub display_name: String,
160    pub source: &'static str,
161    pub region: String,
162}
163#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
164pub struct Usage {
165    pub input_tokens: u64,
166    pub output_tokens: u64,
167}
168#[derive(Debug, Clone, PartialEq, Eq)]
169pub struct ToolCall {
170    pub id: String,
171    pub name: String,
172    pub arguments: String,
173}
174#[derive(Debug, Clone, PartialEq, Eq)]
175pub struct ConverseOutput {
176    pub text: String,
177    pub tool_calls: Vec<ToolCall>,
178    pub usage: Usage,
179}
180#[derive(Debug, Clone, PartialEq, Eq)]
181pub enum ConverseEvent {
182    TextDelta(String),
183    ToolArguments { id: String, delta: String },
184    Usage(Usage),
185    Done,
186}
187#[derive(Debug, Clone, Copy)]
188pub enum TokenGrant<'a> {
189    DeviceCode(&'a str),
190    RefreshToken(&'a str),
191}
192#[derive(Debug, Clone, Copy)]
193pub enum Selection {
194    Explicit,
195}
196#[derive(Debug, thiserror::Error)]
197pub enum AwsError {
198    #[error("explicit selection required: {0}")]
199    SelectionRequired(&'static str),
200    #[error("authorization pending")]
201    AuthorizationPending,
202    #[error("authorization polling slowed down")]
203    SlowDown,
204    #[error("AWS authorization denied or expired")]
205    AuthorizationFailed,
206    #[error("invalid model id")]
207    InvalidModel,
208    #[error("AWS operation failed (details redacted)")]
209    Upstream,
210}
211
212#[derive(Clone)]
213pub struct SignedRequest {
214    pub method: &'static str,
215    pub host: String,
216    pub path: String,
217    pub body: Vec<u8>,
218    headers: BTreeMap<String, String>,
219}
220impl SignedRequest {
221    pub fn header(&self, n: &str) -> Option<&str> {
222        self.headers.get(n).map(String::as_str)
223    }
224    pub fn has_header(&self, n: &str) -> bool {
225        self.headers.contains_key(n)
226    }
227}
228impl fmt::Debug for SignedRequest {
229    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
230        f.debug_struct("SignedRequest")
231            .field("method", &self.method)
232            .field("host", &self.host)
233            .field("path", &self.path)
234            .field("headers", &"[REDACTED]")
235            .finish()
236    }
237}
238
239pub type ConverseEventStream =
240    Pin<Box<dyn Stream<Item = Result<ConverseEvent, AwsError>> + Send + 'static>>;
241
242#[async_trait]
243pub trait AwsApi: Send + Sync + Clone + 'static {
244    async fn register_client(&self, region: &str) -> Result<RegisteredClient, AwsError>;
245    async fn start_device_authorization(
246        &self,
247        client: &RegisteredClient,
248        start_url: &str,
249    ) -> Result<DeviceAuthorization, AwsError>;
250    async fn create_token(
251        &self,
252        client: &RegisteredClient,
253        region: &str,
254        grant: TokenGrant<'_>,
255    ) -> Result<SsoToken, AwsError>;
256    async fn list_accounts(
257        &self,
258        region: &str,
259        access_token: &str,
260    ) -> Result<Vec<Account>, AwsError>;
261    async fn list_account_roles(
262        &self,
263        region: &str,
264        access_token: &str,
265        account: &str,
266    ) -> Result<Vec<Role>, AwsError>;
267    async fn get_role_credentials(
268        &self,
269        region: &str,
270        access_token: &str,
271        account: &str,
272        role: &str,
273    ) -> Result<RoleCredentials, AwsError>;
274    async fn list_foundation_models(
275        &self,
276        request: SignedRequest,
277    ) -> Result<Vec<FoundationModel>, AwsError>;
278    async fn converse(&self, request: SignedRequest) -> Result<ConverseOutput, AwsError>;
279    async fn converse_stream(
280        &self,
281        request: SignedRequest,
282    ) -> Result<ConverseEventStream, AwsError>;
283}
284#[derive(Debug, Clone, Serialize)]
285pub struct PublicContext {
286    pub sso_region: String,
287    pub account_id: String,
288    pub role_name: String,
289    pub bedrock_region: String,
290}
291pub struct AwsBedrockBroker<A: AwsApi> {
292    api: A,
293    config: AwsBedrockConfig,
294    credentials: RoleCredentials,
295}
296impl<A: AwsApi> AwsBedrockBroker<A> {
297    /// Restore broker-owned temporary credentials. This is intentionally only a
298    /// typed Bedrock adapter; credentials can never be extracted or used to sign
299    /// caller-selected requests.
300    pub fn from_credentials(
301        api: A,
302        config: AwsBedrockConfig,
303        credentials: RoleCredentials,
304    ) -> Self {
305        Self {
306            api,
307            config,
308            credentials,
309        }
310    }
311
312    pub async fn login(api: A, config: AwsBedrockConfig, _: Selection) -> Result<Self, AwsError> {
313        let c = api.register_client(&config.sso_region).await?;
314        let d = api
315            .start_device_authorization(&c, &config.sso_start_url)
316            .await?;
317        let token = api
318            .create_token(
319                &c,
320                &config.sso_region,
321                TokenGrant::DeviceCode(&d.device_code),
322            )
323            .await?;
324        let accounts = api.list_accounts(&config.sso_region, &token.access).await?;
325        if accounts.len() != 1 {
326            return Err(AwsError::SelectionRequired("account"));
327        }
328        if accounts[0].id != config.account_id {
329            return Err(AwsError::SelectionRequired("configured account"));
330        }
331        let roles = api
332            .list_account_roles(&config.sso_region, &token.access, &config.account_id)
333            .await?;
334        if roles.len() != 1 || roles[0].name != config.role_name {
335            return Err(AwsError::SelectionRequired("role"));
336        }
337        let credentials = api
338            .get_role_credentials(
339                &config.sso_region,
340                &token.access,
341                &config.account_id,
342                &config.role_name,
343            )
344            .await?;
345        Ok(Self {
346            api,
347            config,
348            credentials,
349        })
350    }
351    pub fn public_context(&self) -> PublicContext {
352        PublicContext {
353            sso_region: self.config.sso_region.clone(),
354            account_id: self.config.account_id.clone(),
355            role_name: self.config.role_name.clone(),
356            bedrock_region: self.config.bedrock_region.clone(),
357        }
358    }
359    pub async fn catalog(&self) -> Result<Vec<CatalogEntry>, AwsError> {
360        let r = sign_bedrock_request(
361            &self.config.bedrock_region,
362            "GET",
363            "/foundation-models",
364            b"",
365            &self.credentials,
366            chrono::Utc::now().timestamp().max(0) as u64,
367        )?;
368        Ok(self
369            .api
370            .list_foundation_models(r)
371            .await?
372            .into_iter()
373            .filter(|m| m.supports_text && m.supports_converse)
374            .map(|m| CatalogEntry {
375                id: format!("aws-bedrock/{}", m.id),
376                display_name: m.name,
377                source: "dynamic",
378                region: self.config.bedrock_region.clone(),
379            })
380            .collect())
381    }
382    pub async fn converse(
383        &self,
384        model: &str,
385        request: InvokeRequest,
386    ) -> Result<ConverseOutput, AwsError> {
387        let id = valid_model(model)?;
388        let body = serde_json::to_vec(&request).map_err(|_| AwsError::Upstream)?;
389        let path = format!("/model/{id}/converse");
390        let r = sign_bedrock_request(
391            &self.config.bedrock_region,
392            "POST",
393            &path,
394            &body,
395            &self.credentials,
396            chrono::Utc::now().timestamp().max(0) as u64,
397        )?;
398        self.api.converse(r).await
399    }
400    pub async fn converse_stream(
401        &self,
402        model: &str,
403        request: InvokeRequest,
404    ) -> Result<ConverseEventStream, AwsError> {
405        let id = valid_model(model)?;
406        let body = serde_json::to_vec(&request).map_err(|_| AwsError::Upstream)?;
407        let path = format!("/model/{id}/converse-stream");
408        let r = sign_bedrock_request(
409            &self.config.bedrock_region,
410            "POST",
411            &path,
412            &body,
413            &self.credentials,
414            chrono::Utc::now().timestamp().max(0) as u64,
415        )?;
416        self.api.converse_stream(r).await
417    }
418}
419
420#[derive(Clone)]
421pub struct AwsHttpApi {
422    http: reqwest::Client,
423    sso_region: String,
424}
425impl AwsHttpApi {
426    pub fn new(http: reqwest::Client, sso_region: impl Into<String>) -> Self {
427        Self {
428            http,
429            sso_region: sso_region.into(),
430        }
431    }
432}
433fn aws_json_error(_status: reqwest::StatusCode) -> AwsError {
434    AwsError::Upstream
435}
436#[async_trait]
437impl AwsApi for AwsHttpApi {
438    async fn register_client(&self, region: &str) -> Result<RegisteredClient, AwsError> {
439        let u = format!("https://oidc.{region}.amazonaws.com/client/register");
440        let v:serde_json::Value=self.http.post(u).json(&serde_json::json!({"clientName":"synaps-cli","clientType":"public","scopes":["sso:account:access"]})).send().await.map_err(|_|AwsError::Upstream)?.error_for_status().map_err(|e|aws_json_error(e.status().unwrap_or(reqwest::StatusCode::BAD_GATEWAY)))?.json().await.map_err(|_|AwsError::Upstream)?;
441        Ok(RegisteredClient::new(
442            v["clientId"].as_str().ok_or(AwsError::Upstream)?,
443            v["clientSecret"].as_str().ok_or(AwsError::Upstream)?,
444            v["clientSecretExpiresAt"]
445                .as_u64()
446                .ok_or(AwsError::Upstream)?,
447        ))
448    }
449    async fn start_device_authorization(
450        &self,
451        c: &RegisteredClient,
452        start: &str,
453    ) -> Result<DeviceAuthorization, AwsError> {
454        let u = format!(
455            "https://oidc.{}.amazonaws.com/device_authorization",
456            self.sso_region
457        );
458        let v: serde_json::Value = self
459            .http
460            .post(u)
461            .json(&serde_json::json!({"clientId":c.id,"clientSecret":c.secret,"startUrl":start}))
462            .send()
463            .await
464            .map_err(|_| AwsError::Upstream)?
465            .error_for_status()
466            .map_err(|_| AwsError::Upstream)?
467            .json()
468            .await
469            .map_err(|_| AwsError::Upstream)?;
470        Ok(DeviceAuthorization::new(
471            v["deviceCode"].as_str().ok_or(AwsError::Upstream)?,
472            v["userCode"].as_str().ok_or(AwsError::Upstream)?,
473            v.get("verificationUriComplete")
474                .or_else(|| v.get("verificationUri"))
475                .and_then(|x| x.as_str())
476                .ok_or(AwsError::Upstream)?,
477            v["interval"].as_u64().unwrap_or(5),
478            v["expiresIn"].as_u64().ok_or(AwsError::Upstream)?,
479        ))
480    }
481    async fn create_token(
482        &self,
483        c: &RegisteredClient,
484        region: &str,
485        g: TokenGrant<'_>,
486    ) -> Result<SsoToken, AwsError> {
487        let (grant, key) = match g {
488            TokenGrant::DeviceCode(x) => (
489                "urn:ietf:params:oauth:grant-type:device_code",
490                ("deviceCode", x),
491            ),
492            TokenGrant::RefreshToken(x) => ("refresh_token", ("refreshToken", x)),
493        };
494        let mut body =
495            serde_json::json!({"clientId":c.id,"clientSecret":c.secret,"grantType":grant});
496        body[key.0] = key.1.into();
497        let r = self
498            .http
499            .post(format!("https://oidc.{region}.amazonaws.com/token"))
500            .json(&body)
501            .send()
502            .await
503            .map_err(|_| AwsError::Upstream)?;
504        if !r.status().is_success() {
505            let t = r.text().await.unwrap_or_default();
506            return Err(if t.contains("AuthorizationPending") {
507                AwsError::AuthorizationPending
508            } else if t.contains("SlowDown") {
509                AwsError::SlowDown
510            } else {
511                AwsError::AuthorizationFailed
512            });
513        }
514        let v: serde_json::Value = r.json().await.map_err(|_| AwsError::Upstream)?;
515        Ok(SsoToken::new(
516            v["accessToken"].as_str().ok_or(AwsError::Upstream)?,
517            v["refreshToken"].as_str(),
518            v["expiresIn"].as_u64().ok_or(AwsError::Upstream)?,
519        ))
520    }
521    async fn list_accounts(&self, r: &str, t: &str) -> Result<Vec<Account>, AwsError> {
522        let v: serde_json::Value = self
523            .http
524            .get(format!(
525                "https://portal.sso.{r}.amazonaws.com/assignment/accounts"
526            ))
527            .header("x-amz-sso_bearer_token", t)
528            .send()
529            .await
530            .map_err(|_| AwsError::Upstream)?
531            .error_for_status()
532            .map_err(|_| AwsError::Upstream)?
533            .json()
534            .await
535            .map_err(|_| AwsError::Upstream)?;
536        Ok(v["accountList"]
537            .as_array()
538            .ok_or(AwsError::Upstream)?
539            .iter()
540            .filter_map(|x| {
541                Some(Account {
542                    id: x["accountId"].as_str()?.into(),
543                    name: x["accountName"].as_str().unwrap_or("").into(),
544                })
545            })
546            .collect())
547    }
548    async fn list_account_roles(&self, r: &str, t: &str, a: &str) -> Result<Vec<Role>, AwsError> {
549        let v: serde_json::Value = self
550            .http
551            .get(format!(
552                "https://portal.sso.{r}.amazonaws.com/assignment/roles?account_id={a}"
553            ))
554            .header("x-amz-sso_bearer_token", t)
555            .send()
556            .await
557            .map_err(|_| AwsError::Upstream)?
558            .error_for_status()
559            .map_err(|_| AwsError::Upstream)?
560            .json()
561            .await
562            .map_err(|_| AwsError::Upstream)?;
563        Ok(v["roleList"]
564            .as_array()
565            .ok_or(AwsError::Upstream)?
566            .iter()
567            .filter_map(|x| {
568                Some(Role {
569                    name: x["roleName"].as_str()?.into(),
570                })
571            })
572            .collect())
573    }
574    async fn get_role_credentials(
575        &self,
576        r: &str,
577        t: &str,
578        a: &str,
579        role: &str,
580    ) -> Result<RoleCredentials, AwsError> {
581        let v:serde_json::Value=self.http.get(format!("https://portal.sso.{r}.amazonaws.com/federation/credentials?account_id={a}&role_name={role}")).header("x-amz-sso_bearer_token",t).send().await.map_err(|_|AwsError::Upstream)?.error_for_status().map_err(|_|AwsError::Upstream)?.json().await.map_err(|_|AwsError::Upstream)?;
582        let c = &v["roleCredentials"];
583        Ok(RoleCredentials::new(
584            c["accessKeyId"].as_str().ok_or(AwsError::Upstream)?,
585            c["secretAccessKey"].as_str().ok_or(AwsError::Upstream)?,
586            c["sessionToken"].as_str().ok_or(AwsError::Upstream)?,
587            c["expiration"].as_u64().ok_or(AwsError::Upstream)?,
588        ))
589    }
590    async fn list_foundation_models(
591        &self,
592        q: SignedRequest,
593    ) -> Result<Vec<FoundationModel>, AwsError> {
594        let v = self.send_signed(q).await?;
595        Ok(v["modelSummaries"]
596            .as_array()
597            .ok_or(AwsError::Upstream)?
598            .iter()
599            .filter_map(|x| {
600                Some(FoundationModel::new(
601                    x["modelId"].as_str()?,
602                    x["modelName"].as_str().unwrap_or(""),
603                    x["inputModalities"].as_array()?.iter().any(|v| v == "TEXT"),
604                    x["inferenceTypesSupported"].as_array().is_some(),
605                ))
606            })
607            .collect())
608    }
609    async fn converse(&self, q: SignedRequest) -> Result<ConverseOutput, AwsError> {
610        let v = self.send_signed(q).await?;
611        Ok(ConverseOutput {
612            text: v
613                .pointer("/output/message/content/0/text")
614                .and_then(|x| x.as_str())
615                .unwrap_or("")
616                .into(),
617            tool_calls: vec![],
618            usage: Usage {
619                input_tokens: v["usage"]["inputTokens"].as_u64().unwrap_or(0),
620                output_tokens: v["usage"]["outputTokens"].as_u64().unwrap_or(0),
621            },
622        })
623    }
624    async fn converse_stream(&self, q: SignedRequest) -> Result<ConverseEventStream, AwsError> {
625        use futures::StreamExt;
626        let url = format!("https://{}{}", q.host, q.path);
627        let mut request = self.http.post(url).body(q.body);
628        for (name, value) in q.headers {
629            request = request.header(name, value);
630        }
631        let response = request
632            .send()
633            .await
634            .map_err(|_| AwsError::Upstream)?
635            .error_for_status()
636            .map_err(|_| AwsError::Upstream)?;
637        let chunks = response.bytes_stream();
638        // `unfold` owns the response and parser buffer. Dropping the returned
639        // stream drops the response immediately (cancellation), while yielding
640        // one frame at a time and retaining at most one bounded frame.
641        const MAX_FRAME: usize = 1024 * 1024;
642        let stream = futures::stream::unfold(
643            (Box::pin(chunks), Vec::<u8>::new(), false),
644            |(mut chunks, mut buffer, done)| async move {
645                if done {
646                    return None;
647                }
648                loop {
649                    if buffer.len() >= 4 {
650                        let total = u32::from_be_bytes(buffer[..4].try_into().ok()?) as usize;
651                        if !(16..=MAX_FRAME).contains(&total) {
652                            return Some((Err(AwsError::Upstream), (chunks, buffer, true)));
653                        }
654                        if buffer.len() >= total {
655                            let frame: Vec<u8> = buffer.drain(..total).collect();
656                            let event = decode_event_frame(&frame);
657                            let terminal = matches!(event, Ok(Some(ConverseEvent::Done)) | Err(_));
658                            match event {
659                                Ok(Some(event)) => {
660                                    return Some((Ok(event), (chunks, buffer, terminal)))
661                                }
662                                Ok(None) => continue,
663                                Err(error) => return Some((Err(error), (chunks, buffer, true))),
664                            }
665                        }
666                    }
667                    match chunks.next().await {
668                        Some(Ok(chunk)) if buffer.len() + chunk.len() <= MAX_FRAME => {
669                            buffer.extend_from_slice(&chunk)
670                        }
671                        Some(_) => return Some((Err(AwsError::Upstream), (chunks, buffer, true))),
672                        None => return Some((Err(AwsError::Upstream), (chunks, buffer, true))),
673                    }
674                }
675            },
676        );
677        Ok(Box::pin(stream))
678    }
679}
680
681fn decode_event_frame(bytes: &[u8]) -> Result<Option<ConverseEvent>, AwsError> {
682    if bytes.len() < 16 {
683        return Err(AwsError::Upstream);
684    }
685    let total = u32::from_be_bytes(bytes[0..4].try_into().unwrap()) as usize;
686    let headers = u32::from_be_bytes(bytes[4..8].try_into().unwrap()) as usize;
687    if total != bytes.len()
688        || headers > total - 16
689        || crc32(&bytes[..8]) != u32::from_be_bytes(bytes[8..12].try_into().unwrap())
690        || crc32(&bytes[..total - 4]) != u32::from_be_bytes(bytes[total - 4..].try_into().unwrap())
691    {
692        return Err(AwsError::Upstream);
693    }
694    let header_bytes = &bytes[12..12 + headers];
695    let payload = &bytes[12 + headers..total - 4];
696    let value: serde_json::Value =
697        serde_json::from_slice(payload).map_err(|_| AwsError::Upstream)?;
698    validate_event_headers(header_bytes, &value)?;
699    if let Some(delta) = value
700        .pointer("/contentBlockDelta/delta/text")
701        .and_then(|v| v.as_str())
702    {
703        Ok(Some(ConverseEvent::TextDelta(delta.to_owned())))
704    } else if let Some(delta) = value
705        .pointer("/contentBlockDelta/delta/toolUse/input")
706        .and_then(|v| v.as_str())
707    {
708        Ok(Some(ConverseEvent::ToolArguments {
709            id: value
710                .pointer("/contentBlockDelta/delta/toolUse/toolUseId")
711                .and_then(|v| v.as_str())
712                .unwrap_or_default()
713                .to_owned(),
714            delta: delta.to_owned(),
715        }))
716    } else if let Some(usage) = value.pointer("/metadata/usage") {
717        Ok(Some(ConverseEvent::Usage(Usage {
718            input_tokens: usage["inputTokens"].as_u64().unwrap_or(0),
719            output_tokens: usage["outputTokens"].as_u64().unwrap_or(0),
720        })))
721    } else if value.get("messageStop").is_some() {
722        Ok(Some(ConverseEvent::Done))
723    } else if value.get("internalServerException").is_some()
724        || value.get("modelStreamErrorException").is_some()
725    {
726        Err(AwsError::Upstream)
727    } else {
728        Ok(None)
729    }
730}
731
732fn validate_event_headers(bytes: &[u8], value: &serde_json::Value) -> Result<(), AwsError> {
733    let mut cursor = 0;
734    let mut values = std::collections::HashMap::new();
735    while cursor < bytes.len() {
736        let name_len = *bytes.get(cursor).ok_or(AwsError::Upstream)? as usize;
737        cursor += 1;
738        let name = std::str::from_utf8(
739            bytes
740                .get(cursor..cursor + name_len)
741                .ok_or(AwsError::Upstream)?,
742        )
743        .map_err(|_| AwsError::Upstream)?;
744        cursor += name_len;
745        if bytes.get(cursor) != Some(&7) {
746            return Err(AwsError::Upstream);
747        }
748        cursor += 1;
749        let len = u16::from_be_bytes(
750            bytes
751                .get(cursor..cursor + 2)
752                .ok_or(AwsError::Upstream)?
753                .try_into()
754                .unwrap(),
755        ) as usize;
756        cursor += 2;
757        let text = std::str::from_utf8(bytes.get(cursor..cursor + len).ok_or(AwsError::Upstream)?)
758            .map_err(|_| AwsError::Upstream)?;
759        cursor += len;
760        if values.insert(name, text).is_some() {
761            return Err(AwsError::Upstream);
762        }
763    }
764    if values.get(":content-type") != Some(&"application/json") {
765        return Err(AwsError::Upstream);
766    }
767    let message = *values.get(":message-type").ok_or(AwsError::Upstream)?;
768    let kind = match message {
769        "event" => *values.get(":event-type").ok_or(AwsError::Upstream)?,
770        "exception" => *values.get(":exception-type").ok_or(AwsError::Upstream)?,
771        _ => return Err(AwsError::Upstream),
772    };
773    if values.contains_key(if message == "event" {
774        ":exception-type"
775    } else {
776        ":event-type"
777    }) || value.get(kind).is_none()
778    {
779        return Err(AwsError::Upstream);
780    }
781    Ok(())
782}
783
784/// Decode AWS EventStream frames returned by Bedrock ConverseStream. CRCs are
785/// validated before JSON is inspected, preventing truncated/corrupt frames from
786/// being interpreted as model output.
787pub fn decode_converse_stream(mut bytes: &[u8]) -> Result<Vec<ConverseEvent>, AwsError> {
788    let mut out = Vec::new();
789    while !bytes.is_empty() {
790        if bytes.len() < 16 {
791            return Err(AwsError::Upstream);
792        }
793        let total = u32::from_be_bytes(bytes[0..4].try_into().unwrap()) as usize;
794        let headers = u32::from_be_bytes(bytes[4..8].try_into().unwrap()) as usize;
795        if total < 16 || total > bytes.len() || headers > total - 16 {
796            return Err(AwsError::Upstream);
797        }
798        if crc32(&bytes[..8]) != u32::from_be_bytes(bytes[8..12].try_into().unwrap())
799            || crc32(&bytes[..total - 4])
800                != u32::from_be_bytes(bytes[total - 4..total].try_into().unwrap())
801        {
802            return Err(AwsError::Upstream);
803        }
804        match decode_event_frame(&bytes[..total])? {
805            Some(ConverseEvent::Done) if out.iter().any(|e| matches!(e, ConverseEvent::Done)) => {
806                return Err(AwsError::Upstream);
807            }
808            Some(event) => out.push(event),
809            None => {}
810        }
811        bytes = &bytes[total..];
812    }
813    if !matches!(out.last(), Some(ConverseEvent::Done)) {
814        // EOF without messageStop is truncation, never successful completion.
815        return Err(AwsError::Upstream);
816    }
817    if out
818        .iter()
819        .filter(|e| matches!(e, ConverseEvent::Done))
820        .count()
821        != 1
822    {
823        return Err(AwsError::Upstream);
824    }
825    Ok(out)
826}
827
828fn crc32(bytes: &[u8]) -> u32 {
829    let mut crc = !0u32;
830    for byte in bytes {
831        crc ^= *byte as u32;
832        for _ in 0..8 {
833            crc = if crc & 1 != 0 {
834                (crc >> 1) ^ 0xedb8_8320
835            } else {
836                crc >> 1
837            };
838        }
839    }
840    !crc
841}
842impl AwsHttpApi {
843    async fn send_signed(&self, q: SignedRequest) -> Result<serde_json::Value, AwsError> {
844        let url = format!("https://{}{}", q.host, q.path);
845        let mut b = match q.method {
846            "GET" => self.http.get(url),
847            _ => self.http.post(url).body(q.body),
848        };
849        for (k, v) in q.headers {
850            b = b.header(k, v)
851        }
852        b.send()
853            .await
854            .map_err(|_| AwsError::Upstream)?
855            .error_for_status()
856            .map_err(|_| AwsError::Upstream)?
857            .json()
858            .await
859            .map_err(|_| AwsError::Upstream)
860    }
861}
862
863fn valid_model(v: &str) -> Result<&str, AwsError> {
864    let v = v
865        .strip_prefix("aws-bedrock/")
866        .ok_or(AwsError::InvalidModel)?;
867    if v.is_empty()
868        || v.contains('/')
869        || !v
870            .bytes()
871            .all(|b| b.is_ascii_alphanumeric() || matches!(b, b'.' | b'-' | b'_' | b':'))
872    {
873        Err(AwsError::InvalidModel)
874    } else {
875        Ok(v)
876    }
877}
878fn hmac(key: &[u8], data: &[u8]) -> Vec<u8> {
879    let mut k = [0u8; 64];
880    if key.len() > 64 {
881        k[..32].copy_from_slice(&Sha256::digest(key))
882    } else {
883        k[..key.len()].copy_from_slice(key)
884    }
885    let mut i = [0x36; 64];
886    let mut o = [0x5c; 64];
887    for x in 0..64 {
888        i[x] ^= k[x];
889        o[x] ^= k[x]
890    }
891    let inner = Sha256::new().chain_update(i).chain_update(data).finalize();
892    Sha256::new()
893        .chain_update(o)
894        .chain_update(inner)
895        .finalize()
896        .to_vec()
897}
898fn hex(v: impl AsRef<[u8]>) -> String {
899    v.as_ref().iter().map(|b| format!("{b:02x}")).collect()
900}
901pub fn sign_bedrock_request(
902    region: &str,
903    method: &str,
904    path: &str,
905    body: &[u8],
906    c: &RoleCredentials,
907    epoch: u64,
908) -> Result<SignedRequest, AwsError> {
909    if method != "GET" && method != "POST" {
910        return Err(AwsError::Upstream);
911    }
912    let dt = chrono::DateTime::from_timestamp(epoch as i64, 0).ok_or(AwsError::Upstream)?;
913    let date = dt.format("%Y%m%d").to_string();
914    let amz = dt.format("%Y%m%dT%H%M%SZ").to_string();
915    let runtime = path.starts_with("/model/");
916    let host = if runtime {
917        format!("bedrock-runtime.{region}.amazonaws.com")
918    } else {
919        format!("bedrock.{region}.amazonaws.com")
920    };
921    let payload = hex(Sha256::digest(body));
922    let content_type = "application/json";
923    let accepts_eventstream = runtime && path.ends_with("converse-stream");
924    let canonical=format!("{method}\n{path}\n\ncontent-type:{content_type}\nhost:{host}\nx-amz-content-sha256:{payload}\nx-amz-date:{amz}\nx-amz-security-token:{}\n\ncontent-type;host;x-amz-content-sha256;x-amz-date;x-amz-security-token\n{payload}",c.session_token);
925    let scope = format!("{date}/{region}/bedrock/aws4_request");
926    let sts = format!(
927        "AWS4-HMAC-SHA256\n{amz}\n{scope}\n{}",
928        hex(Sha256::digest(canonical))
929    );
930    let kd = hmac(format!("AWS4{}", c.secret_key).as_bytes(), date.as_bytes());
931    let kr = hmac(&kd, region.as_bytes());
932    let ks = hmac(&kr, b"bedrock");
933    let key = hmac(&ks, b"aws4_request");
934    let sig = hex(hmac(&key, sts.as_bytes()));
935    let mut headers = BTreeMap::new();
936    headers.insert("content-type".into(), content_type.into());
937    if accepts_eventstream {
938        headers.insert("accept".into(), "application/vnd.amazon.eventstream".into());
939    }
940    headers.insert("host".into(), host.clone());
941    headers.insert("x-amz-content-sha256".into(), payload);
942    headers.insert("x-amz-date".into(), amz);
943    headers.insert("x-amz-security-token".into(), c.session_token.clone());
944    headers.insert("authorization".into(),format!("AWS4-HMAC-SHA256 Credential={}/{scope}, SignedHeaders=content-type;host;x-amz-content-sha256;x-amz-date;x-amz-security-token, Signature={sig}",c.access_key));
945    Ok(SignedRequest {
946        method: if method == "GET" { "GET" } else { "POST" },
947        host,
948        path: path.into(),
949        body: body.into(),
950        headers,
951    })
952}
agent_core/core/auth/aws_bedrock.rs

agent_core/core/auth/
aws_bedrock.rs
