mirror of
https://github.com/mountain-loop/yaak.git
synced 2026-07-13 08:13:00 +02:00
Fix multiple Set-Cookie headers not being preserved
Changed HttpResponse.headers from HashMap<String, String> to Vec<(String, String)> to preserve duplicate headers. This fixes the bug where only the last Set-Cookie header was stored when servers sent multiple cookies. Added test case for multiple Set-Cookie headers to prevent regression. Generated with [Claude Code](https://claude.ai/code) via [Happy](https://happy.engineering) Co-Authored-By: Claude <noreply@anthropic.com> Co-Authored-By: Happy <yesreply@happy.engineering>
This commit is contained in:
@@ -4,7 +4,6 @@ use crate::types::{SendableBody, SendableHttpRequest};
|
|||||||
use async_trait::async_trait;
|
use async_trait::async_trait;
|
||||||
use futures_util::StreamExt;
|
use futures_util::StreamExt;
|
||||||
use reqwest::{Client, Method, Version};
|
use reqwest::{Client, Method, Version};
|
||||||
use std::collections::HashMap;
|
|
||||||
use std::fmt::Display;
|
use std::fmt::Display;
|
||||||
use std::pin::Pin;
|
use std::pin::Pin;
|
||||||
use std::task::{Context, Poll};
|
use std::task::{Context, Poll};
|
||||||
@@ -153,10 +152,10 @@ pub struct HttpResponse {
|
|||||||
pub status: u16,
|
pub status: u16,
|
||||||
/// HTTP status reason phrase (e.g., "OK", "Not Found")
|
/// HTTP status reason phrase (e.g., "OK", "Not Found")
|
||||||
pub status_reason: Option<String>,
|
pub status_reason: Option<String>,
|
||||||
/// Response headers
|
/// Response headers (Vec to support multiple headers with same name, e.g., Set-Cookie)
|
||||||
pub headers: HashMap<String, String>,
|
pub headers: Vec<(String, String)>,
|
||||||
/// Request headers
|
/// Request headers (Vec to support multiple headers with same name)
|
||||||
pub request_headers: HashMap<String, String>,
|
pub request_headers: Vec<(String, String)>,
|
||||||
/// Content-Length from headers (may differ from actual body size)
|
/// Content-Length from headers (may differ from actual body size)
|
||||||
pub content_length: Option<u64>,
|
pub content_length: Option<u64>,
|
||||||
/// Final URL (after redirects)
|
/// Final URL (after redirects)
|
||||||
@@ -194,8 +193,8 @@ impl HttpResponse {
|
|||||||
pub fn new(
|
pub fn new(
|
||||||
status: u16,
|
status: u16,
|
||||||
status_reason: Option<String>,
|
status_reason: Option<String>,
|
||||||
headers: HashMap<String, String>,
|
headers: Vec<(String, String)>,
|
||||||
request_headers: HashMap<String, String>,
|
request_headers: Vec<(String, String)>,
|
||||||
content_length: Option<u64>,
|
content_length: Option<u64>,
|
||||||
url: String,
|
url: String,
|
||||||
remote_addr: Option<String>,
|
remote_addr: Option<String>,
|
||||||
@@ -395,10 +394,10 @@ impl HttpSender for ReqwestSender {
|
|||||||
method: sendable_req.method().to_string(),
|
method: sendable_req.method().to_string(),
|
||||||
});
|
});
|
||||||
|
|
||||||
let mut request_headers = HashMap::new();
|
let mut request_headers = Vec::new();
|
||||||
for (name, value) in sendable_req.headers() {
|
for (name, value) in sendable_req.headers() {
|
||||||
let v = value.to_str().unwrap_or_default().to_string();
|
let v = value.to_str().unwrap_or_default().to_string();
|
||||||
request_headers.insert(name.to_string(), v.clone());
|
request_headers.push((name.to_string(), v.clone()));
|
||||||
send_event(HttpResponseEvent::HeaderUp(name.to_string(), v));
|
send_event(HttpResponseEvent::HeaderUp(name.to_string(), v));
|
||||||
}
|
}
|
||||||
send_event(HttpResponseEvent::Info("Sending request to server".to_string()));
|
send_event(HttpResponseEvent::Info("Sending request to server".to_string()));
|
||||||
@@ -426,12 +425,12 @@ impl HttpSender for ReqwestSender {
|
|||||||
status: response.status().to_string(),
|
status: response.status().to_string(),
|
||||||
});
|
});
|
||||||
|
|
||||||
// Extract headers
|
// Extract headers (use Vec to preserve duplicates like Set-Cookie)
|
||||||
let mut headers = HashMap::new();
|
let mut headers = Vec::new();
|
||||||
for (key, value) in response.headers() {
|
for (key, value) in response.headers() {
|
||||||
if let Ok(v) = value.to_str() {
|
if let Ok(v) = value.to_str() {
|
||||||
send_event(HttpResponseEvent::HeaderDown(key.to_string(), v.to_string()));
|
send_event(HttpResponseEvent::HeaderDown(key.to_string(), v.to_string()));
|
||||||
headers.insert(key.to_string(), v.to_string());
|
headers.push((key.to_string(), v.to_string()));
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@@ -264,7 +264,6 @@ mod tests {
|
|||||||
use crate::decompress::ContentEncoding;
|
use crate::decompress::ContentEncoding;
|
||||||
use crate::sender::{HttpResponseEvent, HttpSender};
|
use crate::sender::{HttpResponseEvent, HttpSender};
|
||||||
use async_trait::async_trait;
|
use async_trait::async_trait;
|
||||||
use std::collections::HashMap;
|
|
||||||
use std::pin::Pin;
|
use std::pin::Pin;
|
||||||
use std::sync::Arc;
|
use std::sync::Arc;
|
||||||
use tokio::io::AsyncRead;
|
use tokio::io::AsyncRead;
|
||||||
@@ -277,7 +276,7 @@ mod tests {
|
|||||||
|
|
||||||
struct MockResponse {
|
struct MockResponse {
|
||||||
status: u16,
|
status: u16,
|
||||||
headers: HashMap<String, String>,
|
headers: Vec<(String, String)>,
|
||||||
body: Vec<u8>,
|
body: Vec<u8>,
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -306,7 +305,7 @@ mod tests {
|
|||||||
mock.status,
|
mock.status,
|
||||||
None, // status_reason
|
None, // status_reason
|
||||||
mock.headers,
|
mock.headers,
|
||||||
HashMap::new(),
|
Vec::new(),
|
||||||
None, // content_length
|
None, // content_length
|
||||||
"https://example.com".to_string(), // url
|
"https://example.com".to_string(), // url
|
||||||
None, // remote_addr
|
None, // remote_addr
|
||||||
@@ -320,7 +319,7 @@ mod tests {
|
|||||||
|
|
||||||
#[tokio::test]
|
#[tokio::test]
|
||||||
async fn test_transaction_no_redirect() {
|
async fn test_transaction_no_redirect() {
|
||||||
let response = MockResponse { status: 200, headers: HashMap::new(), body: b"OK".to_vec() };
|
let response = MockResponse { status: 200, headers: Vec::new(), body: b"OK".to_vec() };
|
||||||
let sender = MockSender::new(vec![response]);
|
let sender = MockSender::new(vec![response]);
|
||||||
let transaction = HttpTransaction::new(sender);
|
let transaction = HttpTransaction::new(sender);
|
||||||
|
|
||||||
@@ -343,12 +342,11 @@ mod tests {
|
|||||||
|
|
||||||
#[tokio::test]
|
#[tokio::test]
|
||||||
async fn test_transaction_single_redirect() {
|
async fn test_transaction_single_redirect() {
|
||||||
let mut redirect_headers = HashMap::new();
|
let redirect_headers = vec![("Location".to_string(), "https://example.com/new".to_string())];
|
||||||
redirect_headers.insert("Location".to_string(), "https://example.com/new".to_string());
|
|
||||||
|
|
||||||
let responses = vec![
|
let responses = vec![
|
||||||
MockResponse { status: 302, headers: redirect_headers, body: vec![] },
|
MockResponse { status: 302, headers: redirect_headers, body: vec![] },
|
||||||
MockResponse { status: 200, headers: HashMap::new(), body: b"Final".to_vec() },
|
MockResponse { status: 200, headers: Vec::new(), body: b"Final".to_vec() },
|
||||||
];
|
];
|
||||||
|
|
||||||
let sender = MockSender::new(responses);
|
let sender = MockSender::new(responses);
|
||||||
@@ -375,8 +373,7 @@ mod tests {
|
|||||||
|
|
||||||
#[tokio::test]
|
#[tokio::test]
|
||||||
async fn test_transaction_max_redirects_exceeded() {
|
async fn test_transaction_max_redirects_exceeded() {
|
||||||
let mut redirect_headers = HashMap::new();
|
let redirect_headers = vec![("Location".to_string(), "https://example.com/loop".to_string())];
|
||||||
redirect_headers.insert("Location".to_string(), "https://example.com/loop".to_string());
|
|
||||||
|
|
||||||
// Create more redirects than allowed
|
// Create more redirects than allowed
|
||||||
let responses: Vec<MockResponse> = (0..12)
|
let responses: Vec<MockResponse> = (0..12)
|
||||||
@@ -474,8 +471,8 @@ mod tests {
|
|||||||
Ok(HttpResponse::new(
|
Ok(HttpResponse::new(
|
||||||
200,
|
200,
|
||||||
None,
|
None,
|
||||||
HashMap::new(),
|
Vec::new(),
|
||||||
HashMap::new(),
|
Vec::new(),
|
||||||
None,
|
None,
|
||||||
"https://example.com".to_string(),
|
"https://example.com".to_string(),
|
||||||
None,
|
None,
|
||||||
@@ -528,8 +525,7 @@ mod tests {
|
|||||||
_request: SendableHttpRequest,
|
_request: SendableHttpRequest,
|
||||||
_event_tx: mpsc::Sender<HttpResponseEvent>,
|
_event_tx: mpsc::Sender<HttpResponseEvent>,
|
||||||
) -> Result<HttpResponse> {
|
) -> Result<HttpResponse> {
|
||||||
let mut headers = HashMap::new();
|
let headers = vec![("set-cookie".to_string(), "session=xyz789; Path=/".to_string())];
|
||||||
headers.insert("set-cookie".to_string(), "session=xyz789; Path=/".to_string());
|
|
||||||
|
|
||||||
let body_stream: Pin<Box<dyn AsyncRead + Send>> =
|
let body_stream: Pin<Box<dyn AsyncRead + Send>> =
|
||||||
Box::pin(std::io::Cursor::new(vec![]));
|
Box::pin(std::io::Cursor::new(vec![]));
|
||||||
@@ -537,7 +533,7 @@ mod tests {
|
|||||||
200,
|
200,
|
||||||
None,
|
None,
|
||||||
headers,
|
headers,
|
||||||
HashMap::new(),
|
Vec::new(),
|
||||||
None,
|
None,
|
||||||
"https://example.com".to_string(),
|
"https://example.com".to_string(),
|
||||||
None,
|
None,
|
||||||
@@ -569,6 +565,79 @@ mod tests {
|
|||||||
assert!(cookies[0].raw_cookie.contains("session=xyz789"));
|
assert!(cookies[0].raw_cookie.contains("session=xyz789"));
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn test_multiple_set_cookie_headers() {
|
||||||
|
// Create a cookie store
|
||||||
|
let cookie_store = CookieStore::new();
|
||||||
|
|
||||||
|
// Mock sender that returns multiple Set-Cookie headers
|
||||||
|
struct MultiSetCookieSender;
|
||||||
|
|
||||||
|
#[async_trait]
|
||||||
|
impl HttpSender for MultiSetCookieSender {
|
||||||
|
async fn send(
|
||||||
|
&self,
|
||||||
|
_request: SendableHttpRequest,
|
||||||
|
_event_tx: mpsc::Sender<HttpResponseEvent>,
|
||||||
|
) -> Result<HttpResponse> {
|
||||||
|
// Multiple Set-Cookie headers (this is standard HTTP behavior)
|
||||||
|
let headers = vec![
|
||||||
|
("set-cookie".to_string(), "session=abc123; Path=/".to_string()),
|
||||||
|
("set-cookie".to_string(), "user_id=42; Path=/".to_string()),
|
||||||
|
("set-cookie".to_string(), "preferences=dark; Path=/; Max-Age=86400".to_string()),
|
||||||
|
];
|
||||||
|
|
||||||
|
let body_stream: Pin<Box<dyn AsyncRead + Send>> =
|
||||||
|
Box::pin(std::io::Cursor::new(vec![]));
|
||||||
|
Ok(HttpResponse::new(
|
||||||
|
200,
|
||||||
|
None,
|
||||||
|
headers,
|
||||||
|
Vec::new(),
|
||||||
|
None,
|
||||||
|
"https://example.com".to_string(),
|
||||||
|
None,
|
||||||
|
Some("HTTP/1.1".to_string()),
|
||||||
|
body_stream,
|
||||||
|
ContentEncoding::Identity,
|
||||||
|
))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
let sender = MultiSetCookieSender;
|
||||||
|
let transaction = HttpTransaction::with_cookie_store(sender, cookie_store.clone());
|
||||||
|
|
||||||
|
let request = SendableHttpRequest {
|
||||||
|
url: "https://example.com/login".to_string(),
|
||||||
|
method: "POST".to_string(),
|
||||||
|
headers: vec![],
|
||||||
|
..Default::default()
|
||||||
|
};
|
||||||
|
|
||||||
|
let (_tx, rx) = tokio::sync::watch::channel(false);
|
||||||
|
let (event_tx, _event_rx) = mpsc::channel(100);
|
||||||
|
let result = transaction.execute_with_cancellation(request, rx, event_tx).await;
|
||||||
|
assert!(result.is_ok());
|
||||||
|
|
||||||
|
// Verify all three cookies were stored
|
||||||
|
let cookies = cookie_store.get_all_cookies();
|
||||||
|
assert_eq!(cookies.len(), 3, "All three Set-Cookie headers should be parsed and stored");
|
||||||
|
|
||||||
|
let cookie_values: Vec<&str> = cookies.iter().map(|c| c.raw_cookie.as_str()).collect();
|
||||||
|
assert!(
|
||||||
|
cookie_values.iter().any(|c| c.contains("session=abc123")),
|
||||||
|
"session cookie should be stored"
|
||||||
|
);
|
||||||
|
assert!(
|
||||||
|
cookie_values.iter().any(|c| c.contains("user_id=42")),
|
||||||
|
"user_id cookie should be stored"
|
||||||
|
);
|
||||||
|
assert!(
|
||||||
|
cookie_values.iter().any(|c| c.contains("preferences=dark")),
|
||||||
|
"preferences cookie should be stored"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
#[tokio::test]
|
#[tokio::test]
|
||||||
async fn test_cookies_across_redirects() {
|
async fn test_cookies_across_redirects() {
|
||||||
use std::sync::atomic::{AtomicUsize, Ordering};
|
use std::sync::atomic::{AtomicUsize, Ordering};
|
||||||
@@ -595,9 +664,10 @@ mod tests {
|
|||||||
|
|
||||||
let (status, headers) = if count == 0 {
|
let (status, headers) = if count == 0 {
|
||||||
// First request: return redirect with Set-Cookie
|
// First request: return redirect with Set-Cookie
|
||||||
let mut h = HashMap::new();
|
let h = vec![
|
||||||
h.insert("location".to_string(), "https://example.com/final".to_string());
|
("location".to_string(), "https://example.com/final".to_string()),
|
||||||
h.insert("set-cookie".to_string(), "redirect_cookie=value1".to_string());
|
("set-cookie".to_string(), "redirect_cookie=value1".to_string()),
|
||||||
|
];
|
||||||
(302, h)
|
(302, h)
|
||||||
} else {
|
} else {
|
||||||
// Second request: verify cookie was sent
|
// Second request: verify cookie was sent
|
||||||
@@ -610,7 +680,7 @@ mod tests {
|
|||||||
"Redirect cookie should be included"
|
"Redirect cookie should be included"
|
||||||
);
|
);
|
||||||
|
|
||||||
(200, HashMap::new())
|
(200, Vec::new())
|
||||||
};
|
};
|
||||||
|
|
||||||
let body_stream: Pin<Box<dyn AsyncRead + Send>> =
|
let body_stream: Pin<Box<dyn AsyncRead + Send>> =
|
||||||
@@ -619,7 +689,7 @@ mod tests {
|
|||||||
status,
|
status,
|
||||||
None,
|
None,
|
||||||
headers,
|
headers,
|
||||||
HashMap::new(),
|
Vec::new(),
|
||||||
None,
|
None,
|
||||||
"https://example.com".to_string(),
|
"https://example.com".to_string(),
|
||||||
None,
|
None,
|
||||||
|
|||||||
Reference in New Issue
Block a user