br-web-server 0.5.15

use crate::request::{Request};
use crate::websocket::{CloseCode, ErrorCode, Message, MessageMode, MessageType};
use std::io::{ErrorKind, Read, Write};
use std::net::TcpStream;
use std::{thread};
use std::sync::{Arc, Mutex};
use std::time::Duration;
use hpack::{Decoder};
use log::{info, warn};
use rustls::{ClientConnection, ServerConnection, StreamOwned};
use crate::{HttpError, Method, Uri};

#[derive(Debug, Clone)]
pub enum Scheme {
    Http(Arc<Mutex<TcpStream>>),
    Https(Arc<Mutex<StreamOwned<ServerConnection,TcpStream>>>),
}

impl Scheme {
    /// 读取
    pub fn read(&mut self, data: &mut Vec<u8>) -> Result<(), HttpError> {
        let mut buf = [0u8; 1024 * 1024];

        let mut index = 2;
        loop {
            let result = match self {
                Scheme::Http(stream) => stream.lock().unwrap().read(&mut buf),
                Scheme::Https(stream) => stream.lock().unwrap().read(&mut buf)
            };
            return match result {
                Ok(0) => Err(HttpError::new(500, "read: 客户端主动关闭")),
                Ok(n) => {
                    data.extend(&buf[..n]);
                    return Ok(());
                }
                Err(ref e) if e.kind() == ErrorKind::Interrupted => {
                    if !data.is_empty() {
                        return Ok(());
                    }
                    if index > 0 {
                        index -= 1;
                        continue;
                    }
                    Err(HttpError::new(500, format!("read现在没数据可读: {}", e.to_string().as_str()).as_str()))
                }
                Err(e) => Err(HttpError::new(500, format!("read: {}", e.to_string().as_str()).as_str())),
            };
        }
    }
    /// 指定长度
    fn read_data(&self, init_data: &mut Vec<u8>, length: usize) -> Result<(), HttpError> {
        loop {
            if init_data.len() >= length {
                return Ok(());
            }
            let mut buf = [0u8; 1024 * 1024];
            let result = match self {
                Scheme::Http(stream) => stream.lock().unwrap().read(&mut buf),
                Scheme::Https(stream) => stream.lock().unwrap().read(&mut buf)
            };
            return match result {
                Ok(0) => Err(HttpError::new(500, "read_data: 客户端主动关闭")),
                Ok(n) => {
                    init_data.extend(&buf[..n]);
                    Ok(())
                }
                Err(ref e) if e.kind() == ErrorKind::WouldBlock => {
                    // ⚠️ 当前无数据可读，稍后再试
                    thread::sleep(Duration::from_millis(100));
                    continue;
                }
                Err(e) => Err(HttpError::new(500, format!("read_data: {}", e.to_string().as_str()).as_str())),
            };
        }
    }

    /// 写入
    pub fn write(&mut self, data: &[u8]) -> Result<(), HttpError> {
        let mut off = 0;
        loop {
            let result = match self {
                Scheme::Http(stream) => stream.lock().unwrap().write(&data[off..]),
                Scheme::Https(stream) => stream.lock().unwrap().get_mut().write(&data[off..])
            };
            match result {
                Ok(0) => return Err(HttpError::new(500, "write: 客户端主动关闭")),
                Ok(e) => {
                    if e != data.len() {
                        off = e;
                        continue;
                    }
                    self.flush()?;
                    return Ok(());
                }
                Err(ref e) if e.kind() == ErrorKind::WouldBlock || e.kind() == ErrorKind::Interrupted => {}
                Err(e) => return Err(HttpError::new(500, format!("write: {}", e.to_string().as_str()).as_str())),
            };
        }
    }
    pub fn write_all(&mut self, data: &[u8]) -> Result<(), HttpError> {
        let result = match self {
            Scheme::Http(stream) => stream.lock().unwrap().write_all(data),
            Scheme::Https(stream) => stream.lock().unwrap().write_all(data)
        };
        match result {
            Ok(()) => {
                self.flush()?;
                Ok(())
            }
            Err(e) => Err(HttpError::new(500, format!("write: {}", e.to_string().as_str()).as_str())),
        }
    }

    pub fn flush(&mut self) -> Result<(), HttpError> {
        let result = match self {
            Scheme::Http(stream) => stream.lock().unwrap().flush(),
            Scheme::Https(stream) => stream.lock().unwrap().flush()
        };
        match result {
            Ok(()) => Ok(()),
            Err(e) => Err(HttpError::new(500, format!("flush: {}", e.to_string().as_str()).as_str())),
        }
    }

    pub fn read_ws_data(&mut self) -> Result<Message, HttpError> {
        let mut response = vec![];
        loop {
            let mut buffer = [0u8; 1024 * 1024];
            let res = match self {
                Scheme::Http(stream) => stream.lock().unwrap().read(&mut buffer),
                Scheme::Https(ref mut stream) => stream.lock().unwrap().read(&mut buffer)
            };
            return match res {
                Ok(0) => Ok(Message {
                    mode: MessageMode::Client,
                    message_type: MessageType::Close,
                    payload: vec![],
                    text: CloseCode::ClientClose.str(),
                    close: CloseCode::ClientClose,
                    error: ErrorCode::None,
                }),
                Ok(n) => {
                    response.extend(buffer[..n].to_vec());
                    let message = Message::parse_message(&mut response);
                    if let MessageType::Continuation = message.message_type { continue; }
                    Ok(message)
                }
                Err(ref e) if e.kind() == ErrorKind::WouldBlock => {
                    continue;
                }
                Err(e) => Ok(Message {
                    mode: MessageMode::Client,
                    message_type: MessageType::Error,
                    payload: vec![],
                    text: e.to_string(),
                    close: CloseCode::Other,
                    error: ErrorCode::Unknown,
                })
            };
        }
    }

    pub fn client_ip(&mut self) -> String {
        match self {
            Scheme::Http(stream) => {
                match stream.lock().unwrap().peer_addr() {
                    Ok(e) => e.ip().to_string(),
                    Err(_) => "未知".to_string(),
                }
            }
            Scheme::Https(stream) => {
                stream.lock().unwrap().get_mut().peer_addr().unwrap().ip().to_string()
            }
        }
    }
    pub fn server_ip(&mut self) -> String {
        match self {
            Scheme::Http(stream) => {
                stream.lock().unwrap().local_addr().unwrap().ip().to_string()
            }
            Scheme::Https(stream) => {
                stream.lock().unwrap().get_mut().local_addr().unwrap().ip().to_string()
            }
        }
    }
    /// 读取HTTP2
    pub fn http2_packet(&mut self, init_data: &mut Vec<u8>) -> Result<(Vec<u8>, FrameType, u8, u32), HttpError> {
        let bytes = init_data;
        self.read_data(bytes, 9)?;
        let headers = bytes.drain(..9).collect::<Vec<u8>>();
        let length = ((headers[0] as u32) << 16) | (u32::from(headers[1]) << 8) | u32::from(headers[2]);
        let frame_type = headers[3];
        let flags = headers[4];
        let stream_id = u32::from_be_bytes([headers[5], headers[6], headers[7], headers[8]]) & 0x7FFF_FFFF;
        self.read_data(bytes, length as usize)?;
        let payload = bytes.drain(..length as usize).collect::<Vec<u8>>();
        Ok((payload, FrameType::from(frame_type), flags, stream_id))
    }
    /// 读取HTTP2消息头
    pub fn http2_handle_header(&mut self, data: &mut Vec<u8>, request: &mut Request) -> Result<(), HttpError> {
        loop {
            let (payload, frame_type, flags, stream_id) = self.http2_packet(data)?;
            if request.config.debug {
                info!("http2_handle_header: frame_type: {frame_type:?} flags: {flags} stream_id: {stream_id} payload: {}", payload.len());
            }
            match frame_type {
                FrameType::Settings => {
                    let is_ack = flags & 0x01 != 0;
                    if !is_ack {
                        self.http2_settings_ack()?;
                    }
                }
                FrameType::WindowUpdate => {
                    if payload.len() == 4 {
                        let raw = u32::from_be_bytes(payload.clone().try_into().unwrap());
                        let increment = raw & 0x7FFF_FFFF; // 屏蔽最高位保留位
                        if request.config.debug {
                            info!("WindowUpdate: increment = {} {:?}", increment,payload);
                        }
                    } else {
                        return Err(HttpError::new(400, format!("Invalid WindowUpdate frame length: {}", payload.len()).as_str()));
                    }
                }
                FrameType::Headers => {
                    let mut decoder = Decoder::new();
                    let headers = decoder.decode(&payload).unwrap();
                    if request.config.debug {
                        println!("=================请求头 {:?}=================", thread::current().id());
                    }
                    for (name, value) in headers {
                        let header_name = String::from_utf8_lossy(name.as_slice());
                        let header_value = String::from_utf8_lossy(value.as_slice());
                        if request.config.debug {
                            println!("{header_name}: {header_value}");
                        }
                        match header_name.as_ref() {
                            ":method" => request.method = Method::from(header_value.as_ref()),
                            ":path" => request.uri = Uri::from(header_value.as_ref()),
                            ":scheme" => request.set_header("scheme", header_value.as_ref())?,
                            ":authority" => request.set_header("host", header_value.as_ref())?,
                            _ => request.set_header(&header_name, &header_value)?
                        }
                    }
                    if request.config.debug {
                        println!("====================================================");
                    }
                    return Ok(());
                }
                _ => return Err(HttpError::new(400, format!("Invalid {frame_type:?}").as_str()))
            }
        }
    }
    /// 读取HTTP2消息体
    pub fn http2_handle_body(&mut self, data: &mut Vec<u8>, request: Request) -> Result<Vec<u8>, HttpError> {
        let mut body = vec![];
        loop {
            let (payload, frame_type, flags, stream_id) = self.http2_packet(data)?;
            if request.config.debug {
                info!("http2_handle_body: frame_type: {frame_type:?} flags: {flags} stream_id: {stream_id} data: {}",payload.len());
            }
            match frame_type {
                FrameType::Data => {
                    if flags == 1 {
                        body.extend(payload);
                        return Ok(body);
                    }
                    body.extend(payload);
                }
                FrameType::Headers => {}
                FrameType::RstStream => {}
                FrameType::Settings => {
                    if !payload.is_empty() {
                        self.http2_send_server_settings()?;
                    } else {
                        self.http2_settings_ack()?;
                    }
                }
                FrameType::Ping => {}
                FrameType::Goaway => {
                    let text = unsafe { String::from_utf8_unchecked(payload) };
                    //if text.as_str() == "SETTINGS expected" {}
                    if request.config.debug {
                        warn!("Goaway: {text}");
                    }
                    return Ok(vec![]);
                }
                FrameType::WindowUpdate => {
                    if payload.len() == 4 {
                        let raw = u32::from_be_bytes(payload.clone().try_into().unwrap());
                        let increment = raw & 0x7FFF_FFFF; // 屏蔽最高位保留位
                        if request.config.debug {
                            info!("WindowUpdate: increment = {} {:?}", increment,payload);
                        }
                    } else {
                        return Err(HttpError::new(400, format!("Invalid WindowUpdate frame length: {}", payload.len()).as_str()));
                    }
                }
                FrameType::Continuation => {}
                FrameType::None => {}
            }
        }
    }
    /// 发送HTTP2认证参数
    pub fn http2_send_server_settings(&mut self) -> Result<(), HttpError> {
        let payload = {
            let mut p = Vec::new();
            // SETTINGS_ENABLE_PUSH = 0  （对浏览器禁用推送）
            p.extend_from_slice(&2u16.to_be_bytes());
            p.extend_from_slice(&0u32.to_be_bytes());
            // SETTINGS_INITIAL_WINDOW_SIZE = 65535
            p.extend_from_slice(&4u16.to_be_bytes());
            p.extend_from_slice(&65_535u32.to_be_bytes());
            // SETTINGS_MAX_FRAME_SIZE = 16384
            p.extend_from_slice(&5u16.to_be_bytes());
            p.extend_from_slice(&16_384u32.to_be_bytes());
            // 可按需再加 MAX_CONCURRENT_STREAMS 等
            p
        };
        let len = payload.len();
        let mut f = Vec::with_capacity(9 + len);
        f.extend_from_slice(&[(len >> 16) as u8, (len >> 8) as u8, len as u8]); // ✅ 正确长度
        f.push(0x04);             // type = SETTINGS
        f.push(0x00);                    // flags = none
        f.extend_from_slice(&0u32.to_be_bytes()); // sid=0
        f.extend_from_slice(&payload);
        self.write_all(&f)?;
        Ok(())
    }
    /// 发送对“对端 SETTINGS”的 ACK（len=0, flags=ACK）
    pub fn http2_settings_ack(&mut self) -> Result<(), HttpError> {
        let f = [0x00, 0x00, 0x00, 0x04, 0x01, 0x00, 0x00, 0x00, 0x00];
        self.write_all(&f)?;
        Ok(())
    }
    pub fn http2_goaway(&mut self, last_stream_id: u32, error_code: u32) -> Result<(), HttpError> {
        // 构造帧头
        let mut frame = Vec::new();
        frame.extend_from_slice(&[0x00, 0x00, 0x08]); // length
        frame.push(0x07); // type = GOAWAY
        frame.push(0x00); // flags = none
        frame.extend_from_slice(&[0x00, 0x00, 0x00, 0x00]); // stream id = 0
        frame.extend_from_slice(&last_stream_id.to_be_bytes());
        frame.extend_from_slice(&error_code.to_be_bytes());
        self.write_all(frame.as_slice())?;
        Ok(())
    }
}
#[derive(Debug)]
pub enum FrameType {
    Data,
    Headers,
    RstStream,
    Settings,
    Ping,
    Goaway,
    WindowUpdate,
    Continuation,
    None,
}
impl FrameType {
    fn from(code: u8) -> FrameType {
        match code {
            0x00 => FrameType::Data,
            0x01 => FrameType::Headers,
            0x03 => FrameType::RstStream,
            0x04 => FrameType::Settings,
            0x06 => FrameType::Ping,
            0x07 => FrameType::Goaway,
            0x08 => FrameType::WindowUpdate,
            0x09 => FrameType::Continuation,
            _ => FrameType::None,
        }
    }
}


pub enum ClientStream {
    Http(TcpStream),
    Https(Box<StreamOwned<ClientConnection, TcpStream>>),
}
impl ClientStream {
    pub fn write_all(&mut self, data: &[u8]) -> std::io::Result<()> {
        match self {
            ClientStream::Http(e) => e.write_all(data),
            ClientStream::Https(e) => e.write_all(data)
        }
    }
    pub fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
        match self {
            ClientStream::Http(e) => e.read(buf),
            ClientStream::Https(e) => e.read(buf)
        }
    }
    pub fn read_data(&mut self, buffer: &mut Vec<u8>) -> Result<(), String> {
        let mut tmp = [0u8; 1024];
        let n = self.read(&mut tmp).map_err(|e| e.to_string())?;
        if n == 0 {
            return Err("unexpected EOF while reading chunk data".to_string());
        }
        buffer.extend_from_slice(&tmp[..n]);
        Ok(())
    }
    /// 强制输出
    pub fn flush(&mut self) -> std::io::Result<()> {
        match self {
            ClientStream::Http(e) => e.flush(),
            ClientStream::Https(e) => e.flush()
        }
    }
}