@
montaro2017 我刚好昨天也遇到了类似的问题,我的做法是在一个长连接里不断检测 header ,然后根据内容再做转发
```rust
loop {
// 读取到完整头部
let header_end = loop {
if let Some(pos) = headers_end_pos(&buf) {
break Some(pos);
}
let n = client_stream.read(&mut tmp).await?;
if n == 0 {
// 客户端关闭或无更多数据
if buf.is_empty() {
return Ok(());
} else {
return Ok(());
}
}
buf.extend_from_slice(&tmp[..n]);
if buf.len() > 128 * 1024 {
error!("请求头过大,终止连接");
return Ok(());
}
};
let header_end = header_end.unwrap();
let headers_vec: Vec<u8> = buf[..header_end].to_vec();
let virtual_env = parse_virtual_env_from_headers(&headers_vec);
let content_length = parse_content_length(&headers_vec);
let chunked = is_chunked(&headers_vec);
// 读取完整正文
let body_len = if let Some(cl) = content_length {
while buf.len() < header_end + cl {
let n = client_stream.read(&mut tmp).await?;
if n == 0 {
error!("Content-Length 指定但连接提前关闭");
return Ok(());
}
buf.extend_from_slice(&tmp[..n]);
}
cl
} else if chunked {
loop {
if let Some(end) = chunked_body_end_pos(&buf[header_end..]) {
break end;
}
let n = client_stream.read(&mut tmp).await?;
if n == 0 {
error!("chunked 正文未完整但连接已关闭");
return Ok(());
}
buf.extend_from_slice(&tmp[..n]);
}
} else {
0
};
let request_end = header_end + body_len;
let body = &buf[header_end..request_end];
let mut new_headers = rewrite_connection_close(&headers_vec);
new_headers.extend_from_slice(body);
// 路由:每个请求一个目标连接;响应仅回写到客户端
if let Some(env) = virtual_env {
let ctrl_opt = {
let envs = (*VIRTUAL_ENVS).lock().unwrap();
envs.get(&env).cloned()
};
if let Some(mut ctrl) = ctrl_opt {
match ctrl.open_stream().await {
Ok(mut sub) => {
sub.write_all(&new_headers).await?;
let _ = tokio::io::copy(&mut sub, &mut client_stream).await;
}
Err(e) => {
error!("打开虚拟环境 {} 的子流失败: {:?}", env, e);
let mut upstream_stream = TcpStream::connect(upstream_addr).await?;
upstream_stream.write_all(&new_headers).await?;
let _ = tokio::io::copy(&mut upstream_stream, &mut client_stream).await;
}
}
} else {
let mut upstream_stream = TcpStream::connect(upstream_addr).await?;
upstream_stream.write_all(&new_headers).await?;
let _ = tokio::io::copy(&mut upstream_stream, &mut client_stream).await;
}
} else {
let mut upstream_stream = TcpStream::connect(upstream_addr).await?;
upstream_stream.write_all(&new_headers).await?;
let _ = tokio::io::copy(&mut upstream_stream, &mut client_stream).await;
}
// 删除已消费的请求字节,保留后续请求(若已到达)
if request_end < buf.len() {
let remaining = buf.split_off(request_end);
buf = remaining;
} else {
buf.clear();
}
}
```