Line data Source code
1 : use std::{str::FromStr, sync::Arc};
2 :
3 : use crate::{http_stream_manager::HttpServerStreamBatch, proxy_headers::HEADER_NAME_SEARCH_PARAMS};
4 : use http::{HeaderMap, HeaderName, HeaderValue};
5 : use std::time::Duration;
6 : use url::Url;
7 :
8 : use crate::{http_stream_manager::HttpStreamManager, proxy_headers::{HEADER_NAME_BATCH_BYTES, HEADER_NAME_BATCH_TIMEOUT, HEADER_NAME_ENDPOINT, HEADER_NAME_METHOD, HEADER_NAME_PATH, HEADER_NAME_READ_TIMEOUT, HEADER_PREFIX}, status::Status};
9 :
10 : /// The parameters of a proxied HTTP request, read from the proxy request headers
11 : #[derive(Debug)]
12 : struct HttpRequestParams {
13 : /// The HTTP method
14 : method: http::Method,
15 : /// The target URL built from endpoint, path and search params
16 : url: Url,
17 : /// The read timeout of the HTTP client in milliseconds
18 : read_timeout: usize,
19 : /// The headers to forward, i.e. all request headers without the proxy prefix
20 : headers: HeaderMap<HeaderValue>
21 : }
22 :
23 : /// Helper to read a string from request headers
24 90 : fn require_string_header(headers: &HeaderMap, header_name: &'static str) -> Result<String, Status> {
25 90 : if let Some(header) = headers.get(header_name) {
26 86 : let header = header
27 86 : .to_str()
28 86 : .map_err(|e| Status::HeaderHasInvalidEncoding{ header: header_name, message: e.to_string() })?
29 86 : .to_string();
30 86 : Ok(header)
31 : } else {
32 4 : Err(Status::HeaderRequiredButMissing { header: header_name })
33 : }
34 90 : }
35 :
36 : /// Helper to read a string from request headers
37 77 : fn require_usize_header(headers: &HeaderMap, header_name: &'static str) -> Result<usize, Status> {
38 77 : if let Some(header) = headers.get(header_name) {
39 73 : let header = header
40 73 : .to_str()
41 73 : .map_err(|e| Status::HeaderHasInvalidEncoding{ header: header_name, message: e.to_string() })?
42 73 : .to_string();
43 73 : let header: usize = header.parse::<usize>().map_err(|e| Status::HeaderIsNotAnUsize { header: header_name, message: e.to_string() })?;
44 72 : Ok(header)
45 : } else {
46 4 : Err(Status::HeaderRequiredButMissing { header: header_name })
47 : }
48 77 : }
49 :
50 : /// Helper to unpack request parameters
51 24 : fn read_request_params(headers: &mut HeaderMap) -> Result<HttpRequestParams, Status> {
52 24 : let mut extra_metadata = HeaderMap::with_capacity(headers.len());
53 24 : let method = require_string_header(headers, HEADER_NAME_METHOD)?;
54 23 : let endpoint = require_string_header(headers, HEADER_NAME_ENDPOINT)?;
55 22 : let path = require_string_header(headers, HEADER_NAME_PATH)?;
56 21 : let search_params = require_string_header(headers, HEADER_NAME_SEARCH_PARAMS)?;
57 20 : let read_timeout = require_usize_header(headers, HEADER_NAME_READ_TIMEOUT)?;
58 :
59 : // Copy all headers in the request that don't start with dashql-
60 93 : for (key, value) in headers.drain() {
61 93 : let key = match &key {
62 93 : Some(k) => k.as_str(),
63 0 : None => continue,
64 : };
65 93 : if !key.starts_with(HEADER_PREFIX) {
66 2 : if let Ok(header) = HeaderName::try_from(key.to_string()) {
67 2 : extra_metadata.insert(header, value);
68 2 : } else {
69 0 : log::warn!("Failed adding extra metadata with key: {}", key);
70 : }
71 91 : }
72 : }
73 : // Remove some headers set by the webviews fetch call
74 18 : extra_metadata.remove("origin");
75 :
76 18 : let url_text = if search_params.is_empty() {
77 17 : format!("{}{}", &endpoint, &path)
78 : } else {
79 1 : format!("{}{}?{}", &endpoint, &path, &search_params)
80 : };
81 18 : let url = url::Url::parse(&url_text)
82 18 : .map_err(|e| {
83 1 : Status::HttpEndpointIsInvalid { header: HEADER_NAME_ENDPOINT, endpoint: url_text, message: e.to_string() }
84 1 : })?;
85 17 : let method = http::Method::from_str(&method)
86 17 : .map_err(|e| {
87 1 : Status::HttpMethodIsInvalid { header: HEADER_NAME_METHOD, method: method.to_string(), message: e.to_string() }
88 1 : })?;
89 :
90 16 : Ok(HttpRequestParams { method, url, read_timeout, headers: extra_metadata })
91 24 : }
92 :
93 : /// The HTTP proxy that streams responses of remote HTTP requests
94 : #[derive(Default)]
95 : pub struct HttpProxy {
96 : /// The response streams of the started requests
97 : pub streams: Arc<HttpStreamManager>,
98 : }
99 :
100 : impl HttpProxy {
101 : /// Start a server stream function
102 24 : pub async fn start_server_stream(&self, headers: &mut HeaderMap, body: Vec<u8>) -> Result<usize, Status> {
103 24 : let mut params = read_request_params(headers)?;
104 16 : log::debug!("Starting stream to remote: {:?}", params.url.to_string());
105 16 : let client = reqwest::Client::builder()
106 16 : .read_timeout(Duration::from_millis(params.read_timeout as u64))
107 16 : .build()
108 16 : .map_err(|e| {
109 0 : Status::HttpClientConfigInvalid { message: e.to_string() }
110 0 : })?;
111 16 : let mut request = reqwest::Request::new(
112 16 : params.method,
113 16 : params.url
114 : );
115 16 : *request.body_mut() = Some(body.into());
116 16 : std::mem::swap(&mut params.headers, &mut request.headers_mut());
117 16 : self.streams.start_server_stream(client, request)
118 24 : }
119 :
120 : /// Read from a result stream
121 20 : pub async fn read_server_stream(&self, stream_id: usize, headers: &HeaderMap) -> Result<HttpServerStreamBatch, Status> {
122 20 : log::debug!("Reading from http stream {}", stream_id);
123 :
124 : // Read limits from request headers
125 20 : let read_timeout = require_usize_header(headers, HEADER_NAME_READ_TIMEOUT)?;
126 19 : let batch_timeout = require_usize_header(headers, HEADER_NAME_BATCH_TIMEOUT)?;
127 18 : let batch_bytes = require_usize_header(headers, HEADER_NAME_BATCH_BYTES)?;
128 :
129 : // Read from the stream
130 17 : let read_timeout_duration = Duration::from_millis(read_timeout as u64);
131 17 : let batch_timeout_duration = Duration::from_millis(batch_timeout as u64);
132 17 : let read_result = self.streams.read_server_stream(stream_id, read_timeout_duration, batch_timeout_duration, batch_bytes).await?;
133 10 : Ok(read_result)
134 20 : }
135 :
136 : /// Destroy a server stream
137 5 : pub async fn destroy_server_stream(&self, stream_id: usize) -> Result<(), Status> {
138 5 : self.streams.destroy_server_stream(stream_id).await;
139 5 : Ok(())
140 5 : }
141 : }
|