Skip to content

Commit df4d388

Browse files
refactor(pyth-lazer-agent): improve proxy CONNECT implementation
- Replace manual string construction with proper HTTP request building - Add robust HTTP response parsing with status code extraction - Read response headers until complete (detect \r\n\r\n terminator) - Provide detailed error messages with status codes and descriptions - Add Proxy-Connection: Keep-Alive header for better compatibility - Fix clippy warnings for inline format args and safe slice access - Improve error handling for edge cases (empty response, invalid status) Co-Authored-By: Mike Rolish <[email protected]>
1 parent 00a4442 commit df4d388

File tree

1 file changed

+58
-14
lines changed

1 file changed

+58
-14
lines changed

apps/pyth-lazer-agent/src/relayer_session.rs

Lines changed: 58 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -60,19 +60,23 @@ async fn connect_through_proxy(
6060
80
6161
});
6262

63-
let mut connect_request = format!(
64-
"CONNECT {target_host}:{target_port} HTTP/1.1\r\nHost: {target_host}:{target_port}\r\n"
65-
);
63+
let target_authority = format!("{target_host}:{target_port}");
64+
let mut request_parts = vec![format!("CONNECT {target_authority} HTTP/1.1")];
65+
request_parts.push(format!("Host: {target_authority}"));
6666

6767
let username = proxy_url.username();
6868
if !username.is_empty() {
6969
let password = proxy_url.password().unwrap_or("");
7070
let credentials = format!("{username}:{password}");
7171
let encoded = base64::engine::general_purpose::STANDARD.encode(credentials.as_bytes());
72-
connect_request = format!("{connect_request}Proxy-Authorization: Basic {encoded}\r\n");
72+
request_parts.push(format!("Proxy-Authorization: Basic {encoded}"));
7373
}
7474

75-
connect_request = format!("{connect_request}\r\n");
75+
request_parts.push("Proxy-Connection: Keep-Alive".to_string());
76+
request_parts.push(String::new()); // Empty line to end headers
77+
request_parts.push(String::new()); // CRLF to end request
78+
79+
let connect_request = request_parts.join("\r\n");
7680

7781
stream
7882
.write_all(connect_request.as_bytes())
@@ -81,18 +85,58 @@ async fn connect_through_proxy(
8185
"Failed to send CONNECT request to proxy at {proxy_url}"
8286
))?;
8387

84-
let mut response = vec![0u8; 1024];
85-
let n = stream.read(&mut response).await.context(format!(
86-
"Failed to read CONNECT response from proxy at {proxy_url}"
87-
))?;
88+
let mut response_buffer = Vec::new();
89+
let mut temp_buf = [0u8; 1024];
90+
let mut headers_complete = false;
91+
92+
while !headers_complete {
93+
let n = stream.read(&mut temp_buf).await.context(format!(
94+
"Failed to read CONNECT response from proxy at {proxy_url}"
95+
))?;
96+
97+
if n == 0 {
98+
bail!("Proxy closed connection before sending complete response");
99+
}
100+
101+
response_buffer.extend_from_slice(temp_buf.get(..n).context("Invalid buffer slice")?);
88102

89-
let response_str =
90-
String::from_utf8_lossy(response.get(..n).context("Invalid response slice range")?);
103+
if response_buffer.windows(4).any(|w| w == b"\r\n\r\n") {
104+
headers_complete = true;
105+
}
106+
}
107+
108+
let response_str = String::from_utf8_lossy(&response_buffer);
109+
110+
let status_line = response_str
111+
.lines()
112+
.next()
113+
.context("Empty response from proxy")?;
114+
115+
let parts: Vec<&str> = status_line.split_whitespace().collect();
116+
if parts.len() < 2 {
117+
bail!(
118+
"Invalid HTTP response from proxy at {}: {}",
119+
proxy_url,
120+
status_line
121+
);
122+
}
91123

92-
if !response_str.starts_with("HTTP/1.1 200") && !response_str.starts_with("HTTP/1.0 200") {
124+
let status_code = parts
125+
.get(1)
126+
.context("Missing status code in proxy response")?
127+
.parse::<u16>()
128+
.context("Invalid status code in proxy response")?;
129+
130+
if status_code != 200 {
131+
let status_text = parts
132+
.get(2..)
133+
.map(|s| s.join(" "))
134+
.unwrap_or_else(|| "Unknown".to_string());
93135
bail!(
94-
"Proxy CONNECT failed: {}",
95-
response_str.lines().next().unwrap_or("Unknown error")
136+
"Proxy CONNECT failed with status {} {}: {}",
137+
status_code,
138+
status_text,
139+
status_line
96140
);
97141
}
98142

0 commit comments

Comments
 (0)