use std::collections::HashMap;
use std::{env, fs, thread};
use std::io::{Write};
use std::path::Path;
use json::{array, object, JsonValue};
use log::debug;
use crate::{parse_content_type_header_value, split_boundary, ContentType, Encoding, Protocol, Status, TransferEncoding};
use crate::stream::ClientStream;
pub struct ClientResponse {
protocol: Protocol,
debug: bool,
stream: ClientStream,
headers: HashMap<String, String>,
cookies: HashMap<String, String>,
status: Status,
body: Vec<u8>,
params: JsonValue,
boundary: String,
content_type: ContentType,
content_length: usize,
transfer_encoding: TransferEncoding,
content_encoding: Encoding,
location: String,
}
impl ClientResponse {
pub fn new(debug: bool, stream: ClientStream) -> Self {
Self { protocol: Protocol::Other("".to_string()), debug, stream, headers: Default::default(), cookies: Default::default(), status: Default::default(), body: vec![], params: object! {}, boundary: "".to_string(), content_type: ContentType::Other("".to_string()), content_length: 0, transfer_encoding: TransferEncoding::Other("".to_string()), content_encoding: Encoding::None, location: "".to_string() }
}
pub fn handle(mut self) -> Result<ClientResponse, String> {
let mut data = vec![];
loop {
let mut buf = [0; 1024];
match self.stream.read(&mut buf) {
Ok(e) => data.extend(buf[..e].to_vec()),
Err(e) => return Err(e.to_string())
}
if let Some(pos) = data.windows(4).position(|window| window == [13, 10, 13, 10]) {
self.handle_header(data.drain(..pos).collect::<Vec<u8>>())?;
data.drain(..4);
self.body = std::mem::take(&mut data);
break;
}
}
loop {
match self.transfer_encoding {
TransferEncoding::Chunked => {
self.body = self.handle_chunked(self.body.clone())?;
self.content_length = self.body.len();
self.handle_body(self.body.clone())?;
break;
}
TransferEncoding::Other(_) => {
if self.content_length == self.body.len() {
self.body = self.content_encoding.decompress(self.body.clone().as_slice())?;
self.content_length = self.body.len();
self.handle_body(self.body.clone())?;
break;
}
let mut buf = [0; 1024 * 1024];
match self.stream.read(&mut buf) {
Ok(e) => self.body.extend(buf[..e].to_vec()),
Err(e) => return Err(e.to_string())
}
}
}
}
Ok(self)
}
fn handle_chunked(&mut self, mut buffer: Vec<u8>) -> Result<Vec<u8>, String> {
let mut decoded = Vec::new();
loop {
if let Some(size_line_end) = buffer.windows(2).position(|w| w == b"\r\n") {
let size_line_bytes = buffer[..size_line_end].to_vec();
let size_line = String::from_utf8_lossy(&size_line_bytes);
let size_str = size_line.split(';').next().unwrap_or("").trim();
let size = usize::from_str_radix(size_str, 16).map_err(|_| format!("invalid chunk size: {}", size_str))?;
buffer.drain(..size_line_end + 2);
match size {
0 => {
if buffer.windows(2).any(|w| w == b"\r\n") {
break;
}
loop {
if buffer.windows(4).any(|w| w == b"\r\n\r\n") {
break;
}
self.stream.read_data(&mut buffer).map_err(|e| e.to_string())?;
}
break;
}
_ => {
while buffer.len() < size + 2 {
self.stream.read_data(&mut buffer).map_err(|e| e.to_string())?;
}
decoded.extend_from_slice(&buffer[..size]);
buffer.drain(..size + 2);
}
}
} else {
self.stream.read_data(&mut buffer).map_err(|e| e.to_string())?;
}
}
let decompressed = self.content_encoding.decompress(decoded.as_slice())?;
Ok(decompressed)
}
fn handle_header(&mut self, data: Vec<u8>) -> Result<(), String> {
let headers = String::from_utf8_lossy(data.as_slice());
if self.debug {
debug!("\r\n=================响应头 {:?}=================\r\n{headers}\r\n========================================",thread::current().id());
}
let header_line = headers.lines().next().unwrap();
let mut it = header_line.split_whitespace();
self.protocol = Protocol::from(it.next().unwrap_or(""));
self.status = Status::from(it.next().unwrap_or("").parse::<u16>().unwrap_or(0), it.next().unwrap_or(""));
match &self.protocol {
Protocol::HTTP1_0 | Protocol::HTTP1_1 => {
for item in headers.lines().skip(1) {
self.header_line_set(item)?;
}
}
Protocol::HTTP2 => {
return Err("HTTP2格式错误".to_string());
}
Protocol::HTTP3 => return Err("暂时未开放".to_string()),
Protocol::Other(name) => {
return Err(format!("未知协议格式: {}", name));
}
}
Ok(())
}
fn header_line_set(&mut self, line: &str) -> Result<(), String> {
match line.trim().find(":") {
None => return Err(format!("请求头[{line}]错误")),
Some(e) => {
let key = line[..e].trim().to_lowercase().clone();
let value = line[e + 1..].trim();
self.set_header(key.as_str(), value)?;
}
}
Ok(())
}
fn set_header(&mut self, key: &str, value: &str) -> Result<(), String> {
self.headers.insert(key.to_string(), value.to_string());
if value.len() > 8192 {
return Err("header longer than 8192 characters".to_string());
}
match key {
"content-type" => {
let (mime, params) = parse_content_type_header_value(value);
if !mime.is_empty() {
if mime == "multipart/form-data" {
if let Some(b) = params.get("boundary") {
self.boundary = b.to_string();
} else {
let lower = value.to_lowercase();
if let Some(pos) = lower.find("boundary=") {
let raw = &value[pos + "boundary=".len()..];
let raw = raw.split(';').next().unwrap_or(raw).trim();
let raw = raw.trim_matches('"');
self.boundary = raw.to_string();
}
}
self.content_type = ContentType::from("multipart/form-data");
} else {
self.content_type = ContentType::from(mime.as_str());
}
}
self.headers.insert(key.to_string(), self.content_type.str().to_string());
}
"content-length" => self.content_length = value.parse::<usize>().unwrap_or(0),
"set-cookie" => {
let _ = value.split(';').collect::<Vec<&str>>().iter().map(|&x| {
match x.find('=') {
None => {}
Some(index) => {
let key = x[..index].trim().to_string();
let val = x[index + 1..].trim().to_string();
let _ = self.cookies.insert(key.to_string(), val);
}
}
""
}).collect::<Vec<&str>>();
}
"transfer-encoding" => {
self.transfer_encoding = TransferEncoding::from(value);
}
"content-encoding" => {
self.content_encoding = Encoding::from(value);
}
"location" => self.location = value.to_string(),
_ => {}
}
Ok(())
}
fn handle_body(&mut self, data: Vec<u8>) -> Result<(), String> {
if self.debug {
debug!("\r\n=================响应体 {:?}=================\r\n长度: {}\r\n========================================",thread::current().id(),self.content_length);
}
if data.len() != self.content_length {
return Err(format!("Content-Length mismatch: header={}, actual={}", self.content_length, data.len()));
}
if self.content_length == 0 {
return Ok(());
}
match &self.content_type {
ContentType::FormData => {
let parts = split_boundary(data, &self.boundary)?;
let mut fields = object! {};
for part in parts {
let (header, body) = match part.windows(b"\r\n\r\n".len()).position(|window| window == b"\r\n\r\n") {
None => continue,
Some(e) => {
let header = part[..e].to_vec();
let body = part[e + 4..].to_vec();
let body = body[..body.len() - 2].to_vec();
(header, body)
}
};
let headers = String::from_utf8_lossy(header.as_slice());
let mut field_name = "";
let mut filename = "";
let mut content_type = ContentType::Other("".to_string());
for header in headers.lines() {
if header.to_lowercase().starts_with("content-disposition:") {
match header.find("filename=\"") {
None => {}
Some(filename_start) => {
let filename_len = filename_start + 10;
let filename_end = header[filename_len..].find('"').unwrap() + filename_len;
filename = &header[filename_len..filename_end];
}
}
match header.find("name=\"") {
None => {}
Some(name_start) => {
let name_start = name_start + 6;
let name_end = header[name_start..].find('"').unwrap() + name_start;
field_name = &header[name_start..name_end];
}
}
}
if header.to_lowercase().starts_with("content-type:") {
content_type = ContentType::from(header.to_lowercase().trim_start_matches("content-type:").trim());
}
}
if filename.is_empty() {
let text = String::from_utf8_lossy(body.as_slice());
fields[field_name.to_string()] = JsonValue::from(text.into_owned());
continue;
}
let extension = Path::new(filename).extension().and_then(|ext| ext.to_str()); let suffix = extension.unwrap_or("txt");
let filename = if extension.is_none() {
format!("{filename}.txt")
} else {
filename.to_string()
};
let mut temp_dir = env::temp_dir();
temp_dir.push(filename.clone());
let Ok(mut temp_file) = fs::File::create(&temp_dir) else { continue };
if temp_file.write(body.as_slice()).is_ok() {
if fields[field_name.to_string()].is_empty() {
fields[field_name.to_string()] = array![];
}
fields[field_name.to_string()].push(object! {
id:br_crypto::sha256::encrypt_hex(&body.clone()),
name:filename,
suffix:suffix,
size:body.len(),
type:content_type.str(),
file:temp_dir.to_str()
}).unwrap();
}
}
self.params = fields;
}
ContentType::FormUrlencoded => {
let input = String::from_utf8_lossy(&data);
let mut list = object! {};
for pair in input.split('&') {
if let Some((key, val)) = pair.split_once('=') {
let key = br_crypto::encoding::urlencoding_decode(key);
let val = br_crypto::encoding::urlencoding_decode(val);
let _ = list.insert(key.as_str(), val);
}
}
self.params = list;
}
ContentType::Json => {
let text = String::from_utf8_lossy(data.as_slice());
self.params = json::parse(text.into_owned().as_str()).unwrap_or(object! {});
}
ContentType::Xml | ContentType::Html | ContentType::Text | ContentType::Javascript => {
let text = String::from_utf8_lossy(data.as_slice());
self.params = text.into_owned().into();
}
ContentType::Other(_) => {}
ContentType::Stream => {}
}
Ok(())
}
pub fn status(&mut self) -> u16 {
self.status.code
}
pub fn location(&mut self) -> String {
self.location.clone()
}
pub fn status_text(&mut self) -> String {
self.status.reason.clone()
}
pub fn json(&mut self) -> JsonValue {
self.params.clone()
}
pub fn headers(&mut self) -> JsonValue {
self.headers.clone().into()
}
pub fn cookies(&mut self) -> JsonValue {
self.cookies.clone().into()
}
pub fn text(&mut self) -> String {
String::from_utf8(self.body.to_vec()).unwrap_or_default()
}
pub fn stream(&mut self) -> Vec<u8> {
self.body.clone()
}
}