Line data    Source code

       1              : use std::io::Write;
       2              : 
       3              : use http::Request;
       4              : use http::Response;
       5              : use lazy_static::lazy_static;
       6              : 
       7              : use crate::http_proxy::HttpProxy;
       8              : use crate::proxy_headers::HEADER_NAME_BATCH_BYTES;
       9              : use crate::proxy_headers::HEADER_NAME_BATCH_EVENT;
      10              : use crate::proxy_headers::HEADER_NAME_RESPONSE_STARTED;
      11              : use crate::proxy_headers::HEADER_NAME_STREAM_ID;
      12              : use crate::proxy_headers::HEADER_PREFIX;
      13              : 
      14              : lazy_static! {
      15              :     static ref HTTP_PROXY: HttpProxy = HttpProxy::default();
      16              : }
      17              : 
      18           24 : pub async fn start_http_server_stream(mut req: Request<Vec<u8>>) -> Response<Vec<u8>> {
      19           24 :     let body = std::mem::take(req.body_mut());
      20           24 :     match HTTP_PROXY.start_server_stream(req.headers_mut(), body).await {
      21           16 :         Ok(stream_id) => {
      22           16 :             let response = Response::builder()
      23           16 :                 .status(200)
      24           16 :                 .header(HEADER_NAME_STREAM_ID, stream_id)
      25           16 :                 .body(Vec::new())
      26           16 :                 .unwrap();
      27           16 :             response
      28              :         },
      29            8 :         Err(e) => Response::from(&e)
      30              :     }
      31           24 : }
      32              : 
      33           20 : pub async fn read_http_server_stream(stream_id: usize, req: Request<Vec<u8>>) -> Response<Vec<u8>> {
      34           20 :     match HTTP_PROXY.read_server_stream(stream_id, req.headers()).await {
      35           10 :         Ok(batches) => {
      36           10 :             let mut buffer: Vec<u8> = Vec::with_capacity(batches.total_body_bytes);
      37           10 :             for message in batches.body_chunks.iter() {
      38            9 :                 buffer.write(&message).unwrap();
      39            9 :             }
      40           10 :             let status_code = match batches.status {
      41            8 :                 Some(status) => status.as_u16(),
      42            2 :                 None => 200,
      43              :             };
      44           10 :             let mut response = Response::builder()
      45           10 :                 .status(status_code)
      46           10 :                 .header(HEADER_NAME_STREAM_ID, stream_id)
      47           10 :                 .header(HEADER_NAME_BATCH_EVENT, batches.event.to_str())
      48           10 :                 .header(HEADER_NAME_BATCH_BYTES, batches.total_body_bytes);
      49           10 :             let headers = &mut response.headers_mut().unwrap();
      50           10 :             if batches.status.is_some() {
      51            8 :                 headers.insert(HEADER_NAME_RESPONSE_STARTED, "true".parse().unwrap());
      52            8 :             }
      53              :             // The `dashql-` headers belong to the proxy, a remote must not be able to replace them.
      54              :             // A header may repeat, e.g. Set-Cookie.
      55           15 :             for (key, value) in batches.headers.iter() {
      56           15 :                 if !key.as_str().starts_with(HEADER_PREFIX) {
      57           12 :                     headers.append(key, value.clone());
      58           12 :                 }
      59              :             }
      60           10 :             response
      61           10 :                 .body(buffer)
      62           10 :                 .unwrap()
      63              :         },
      64           10 :         Err(e) => Response::from(&e)
      65              :     }
      66           20 : }
      67              : 
      68            5 : pub async fn delete_http_server_stream(stream_id: usize, _req: Request<Vec<u8>>) -> Response<Vec<u8>> {
      69            5 :     match HTTP_PROXY.destroy_server_stream(stream_id).await {
      70              :         Ok(()) => {
      71            5 :             Response::builder()
      72            5 :                 .status(200)
      73            5 :                 .body(Vec::new())
      74            5 :                 .unwrap()
      75              :         }
      76            0 :         Err(e) => Response::from(&e)
      77              :     }
      78            5 : }