sigv4 wip

Author: Rutik7066

Cargo.lock

@@ -28,6 +28,7 @@ base64 = { workspace = true }
 hmac = "0.12"
 sha2 = "0.10"
 time = { version = "0.3", features = ["formatting"] }
+percent-encoding = "2.3"
 
 [package.metadata.component]
 package = "golem:llm-bedrock"

llm/bedrock/Cargo.toml

@@ -66,7 +66,6 @@ impl BedrockClient {
         })?;
 
         let mut request_builder = self.client.request(Method::POST, &url);
-        request_builder = request_builder.header("content-type", "application/json");
         for (key, value) in headers {
             request_builder = request_builder.header(key, value);
         }
@@ -116,7 +115,6 @@ impl BedrockClient {
         })?;
 
         let mut request_builder = self.client.request(Method::POST, &url);
-        request_builder = request_builder.header("content-type", "application/json");
         for (key, value) in headers {
             request_builder = request_builder.header(key, value);
         }
@@ -134,6 +132,7 @@ impl BedrockClient {
     }
 }
 
+/// FIXED: AWS SigV4 headers generation with proper content-type handling
 pub fn generate_sigv4_headers(
     access_key: &str,
     secret_key: &str,
@@ -167,16 +166,18 @@ pub fn generate_sigv4_headers(
     let path = if uri.starts_with('/') { uri } else { "/" };
     let query = "";
 
-    // Create canonical headers
+    // FIXED: Create canonical headers with content-type included and proper sorting
     let mut headers: Vec<(String, String)> = vec![
         ("host".to_string(), host.to_string()),
         ("x-amz-date".to_string(), datetime_str.clone()),
+        ("content-type".to_string(), "application/x-amz-json-1.0".to_string()),
     ];
     headers.sort_by(|a, b| a.0.cmp(&b.0));
 
+    // FIXED: Proper trimming of header keys and values
     let canonical_headers = headers
         .iter()
-        .map(|(k, v)| format!("{}:{}", k, v))
+        .map(|(k, v)| format!("{}:{}", k.trim(), v.trim()))
         .collect::<Vec<_>>()
         .join("\n")
         + "\n";
@@ -204,7 +205,7 @@ pub fn generate_sigv4_headers(
         Sha256::digest(canonical_request.as_bytes())
     );
 
-    // Calculate signature
+    // Calculate signature using AWS SigV4 key derivation
     type HmacSha256 = Hmac<Sha256>;
 
     let mut mac = HmacSha256::new_from_slice(format!("AWS4{}", secret_key).as_bytes())?;
@@ -233,9 +234,11 @@ pub fn generate_sigv4_headers(
         access_key, credential_scope, signed_headers, signature
     );
 
-    let mut result_headers = vec![
+    // FIXED: Return all required headers including content-type
+    let result_headers = vec![
         ("authorization".to_string(), auth_header),
         ("x-amz-date".to_string(), datetime_str),
+        ("content-type".to_string(), "application/x-amz-json-1.0".to_string()),
     ];
 
     Ok(result_headers)
@@ -373,20 +376,17 @@ pub struct ToolConfig {
 }
 
 #[derive(Debug, Clone, Serialize, Deserialize)]
-#[serde(tag = "type")]
-pub enum Tool {
+pub struct Tool {
     #[serde(rename = "toolSpec")]
-    ToolSpec {
-        name: String,
-        description: String,
-        #[serde(rename = "inputSchema")]
-        input_schema: ToolInputSchema,
-    },
+    pub tool_spec: ToolSpec,
 }
 
 #[derive(Debug, Clone, Serialize, Deserialize)]
-pub struct ToolInputSchema {
-    pub json: Value,
+pub struct ToolSpec {
+    pub name: String,
+    pub description: String,
+    #[serde(rename = "inputSchema")]
+    pub input_schema: Value,
 }
 
 #[derive(Debug, Clone, Serialize, Deserialize)]
@@ -406,16 +406,7 @@ pub struct GuardrailConfig {
     pub guardrail_identifier: String,
     #[serde(rename = "guardrailVersion")]
     pub guardrail_version: String,
-    #[serde(skip_serializing_if = "Option::is_none")]
-    pub trace: Option<GuardrailTrace>,
-}
-
-#[derive(Debug, Clone, Serialize, Deserialize)]
-pub enum GuardrailTrace {
-    #[serde(rename = "enabled")]
-    Enabled,
-    #[serde(rename = "disabled")]
-    Disabled,
+    pub trace: Option<String>,
 }
 
 #[derive(Debug, Clone, Serialize, Deserialize)]
@@ -503,4 +494,135 @@ fn parse_response<T: DeserializeOwned + Debug>(response: Response) -> Result<T,
     }
 }
 
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_generate_sigv4_headers_basic() {
+        let access_key = "AKIAIOSFODNN7EXAMPLE";
+        let secret_key = "wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY";
+        let region = "us-east-1";
+        let service = "bedrock";
+        let method = "POST";
+        let uri = "/model/anthropic.claude-3-sonnet-20240229-v1:0/converse";
+        let host = "bedrock-runtime.us-east-1.amazonaws.com";
+        let body = r#"{"messages":[{"role":"user","content":[{"type":"text","text":"Hello"}]}]}"#;
+
+        let result = generate_sigv4_headers(
+            access_key,
+            secret_key,
+            region,
+            service,
+            method,
+            uri,
+            host,
+            body,
+        );
+
+        assert!(result.is_ok());
+        let headers = result.unwrap();
+        
+        // Check that required headers are present
+        let header_map: std::collections::HashMap<String, String> = headers.into_iter().collect();
+        
+        assert!(header_map.contains_key("authorization"));
+        assert!(header_map.contains_key("x-amz-date"));
+        assert!(header_map.contains_key("content-type"));
+        
+        // Check authorization header format
+        let auth_header = &header_map["authorization"];
+        assert!(auth_header.starts_with("AWS4-HMAC-SHA256 Credential="));
+        assert!(auth_header.contains("SignedHeaders="));
+        assert!(auth_header.contains("Signature="));
+        
+        // Check content-type
+        assert_eq!(header_map["content-type"], "application/x-amz-json-1.0");
+        
+        // Check x-amz-date format (should be ISO8601)
+        let date_header = &header_map["x-amz-date"];
+        assert_eq!(date_header.len(), 16); // YYYYMMDDTHHMMSSZ
+        assert!(date_header.ends_with('Z'));
+        assert!(date_header.contains('T'));
+    }
+
+    #[test]
+    fn test_canonical_headers_ordering() {
+        let access_key = "AKIAIOSFODNN7EXAMPLE";
+        let secret_key = "wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY";
+        let region = "us-east-1";
+        let service = "bedrock";
+        let method = "POST";
+        let uri = "/model/test/converse";
+        let host = "bedrock-runtime.us-east-1.amazonaws.com";
+        let body = "{}";
+
+        let result = generate_sigv4_headers(
+            access_key,
+            secret_key,
+            region,
+            service,
+            method,
+            uri,
+            host,
+            body,
+        );
+
+        assert!(result.is_ok());
+        let headers = result.unwrap();
+        let header_map: std::collections::HashMap<String, String> = headers.into_iter().collect();
+        
+        let auth_header = &header_map["authorization"];
+        
+        // SignedHeaders should be in alphabetical order: content-type;host;x-amz-date
+        assert!(auth_header.contains("SignedHeaders=content-type;host;x-amz-date"));
+    }
+
+    #[test]
+    fn test_bedrock_client_integration() {
+        let client = BedrockClient::new(
+            "AKIAIOSFODNN7EXAMPLE".to_string(),
+            "wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY".to_string(),
+            "us-east-1".to_string(),
+        );
+
+        let request = ConverseRequest {
+            model_id: "anthropic.claude-3-sonnet-20240229-v1:0".to_string(),
+            messages: vec![Message {
+                role: Role::User,
+                content: vec![ContentBlock::Text {
+                    text: "Hello, how are you?".to_string(),
+                }],
+            }],
+            system: None,
+            inference_config: None,
+            tool_config: None,
+            guardrail_config: None,
+            additional_model_request_fields: None,
+        };
+
+        // This test verifies that serialization and signing work together
+        let body = serde_json::to_string(&request).expect("Failed to serialize request");
+        let host = format!("bedrock-runtime.{}.amazonaws.com", client.region);
 
+        let headers_result = generate_sigv4_headers(
+            &client.access_key_id,
+            &client.secret_access_key,
+            &client.region,
+            "bedrock",
+            "POST",
+            "/model/anthropic.claude-3-sonnet-20240229-v1:0/converse",
+            &host,
+            &body,
+        );
+
+        assert!(headers_result.is_ok());
+        let headers = headers_result.unwrap();
+        
+        // Verify all required headers are present for AWS API call
+        let header_names: Vec<&str> = headers.iter().map(|(k, _)| k.as_str()).collect();
+        assert!(header_names.contains(&"authorization"));
+        assert!(header_names.contains(&"x-amz-date"));
+        assert!(header_names.contains(&"content-type"));
+    }
+}

llm/bedrock/src/client.rs

@@ -1,7 +1,7 @@
 use crate::client::{
     ContentBlock, ConverseRequest, ConverseResponse, ImageFormat, ImageSource as ClientImageSource,
     InferenceConfig, Message as ClientMessage, Role as ClientRole, StopReason, SystemContentBlock,
-    Tool, ToolChoice, ToolConfig, ToolInputSchema, ToolResultContentBlock, ToolResultStatus,
+    Tool, ToolChoice, ToolConfig, ToolSpec, ToolResultContentBlock, ToolResultStatus,
 };
 use base64::{engine::general_purpose, Engine as _};
 use golem_llm::golem::llm::llm::{
@@ -287,10 +287,12 @@ fn message_to_system_content(message: &Message) -> Vec<SystemContentBlock> {
 
 fn tool_definition_to_tool(tool: &ToolDefinition) -> Result<Tool, Error> {
     match serde_json::from_str(&tool.parameters_schema) {
-        Ok(json_schema) => Ok(Tool::ToolSpec {
-            name: tool.name.clone(),
-            description: tool.description.clone().unwrap_or_default(),
-            input_schema: ToolInputSchema { json: json_schema },
+        Ok(json_schema) => Ok(Tool {
+            tool_spec: ToolSpec {
+                name: tool.name.clone(),
+                description: tool.description.clone().unwrap_or_default(),
+                input_schema: json_schema,
+            },
         }),
         Err(error) => Err(Error {
             code: ErrorCode::InternalError,

llm/bedrock/src/conversions.rs

@@ -38,8 +38,8 @@ path = "wit-generated"
 
 [package.metadata.component.target.dependencies]
 "golem:llm" = { path = "wit-generated/deps/golem-llm" }
-"wasi:io" = { path = "wit-generated/deps/io" }
 "wasi:clocks" = { path = "wit-generated/deps/clocks" }
+"wasi:io" = { path = "wit-generated/deps/io" }
 "golem:rpc" = { path = "wit-generated/deps/golem-rpc" }
 "test:helper-client" = { path = "wit-generated/deps/test_helper-client" }
 "test:llm-exports" = { path = "wit-generated/deps/test_llm-exports" }