Line data Source code
1 : use std::io::Write;
2 : use std::str::FromStr;
3 :
4 : use base64::Engine;
5 : use byteorder::LittleEndian;
6 : use byteorder::WriteBytesExt;
7 : use http::HeaderMap;
8 : use http::HeaderName;
9 : use http::HeaderValue;
10 : use http::Request;
11 : use http::Response;
12 : use tonic::metadata::KeyAndValueRef;
13 : use tonic::metadata::MetadataMap;
14 : use lazy_static::lazy_static;
15 : use tracing::instrument;
16 :
17 : use crate::grpc_proxy::GrpcProxy;
18 : use crate::proxy_headers::HEADER_NAME_BATCH_BYTES;
19 : use crate::proxy_headers::HEADER_NAME_BATCH_EVENT;
20 : use crate::proxy_headers::HEADER_NAME_BATCH_MESSAGES;
21 : use crate::proxy_headers::HEADER_NAME_CHANNEL_ID;
22 : use crate::proxy_headers::HEADER_NAME_STREAM_ID;
23 : use crate::proxy_headers::HEADER_PREFIX;
24 :
25 : lazy_static! {
26 : static ref GRPC_PROXY: GrpcProxy = GrpcProxy::default();
27 : }
28 :
29 : #[instrument(skip(req))]
30 19 : pub async fn create_grpc_channel(mut req: Request<Vec<u8>>) -> Response<Vec<u8>> {
31 : match GRPC_PROXY.create_channel(req.headers_mut()).await {
32 : Ok(channel_id) => {
33 : Response::builder()
34 : .status(200)
35 : .header("X-Custom-Foo", "bar")
36 : .header(HEADER_NAME_CHANNEL_ID, channel_id)
37 : .body(Vec::new())
38 : .unwrap()
39 : },
40 : Err(e) => Response::from(&e)
41 : }
42 19 : }
43 :
44 1 : pub async fn delete_grpc_channel(channel_id: usize) -> Response<Vec<u8>> {
45 1 : match GRPC_PROXY.destroy_channel(channel_id).await {
46 : Ok(()) => {
47 1 : let response = Response::builder()
48 1 : .status(200)
49 1 : .body(Vec::new())
50 1 : .unwrap();
51 1 : response
52 : },
53 0 : Err(e) => Response::from(&e)
54 :
55 : }
56 1 : }
57 :
58 : /// Copy the metadata or the trailers of a call into response headers.
59 : /// The `dashql-` headers belong to the proxy, a server must not be able to replace them.
60 10 : fn copy_metadata(metadata: &MetadataMap, headers: &mut HeaderMap) {
61 23 : for entry in metadata.iter() {
62 23 : match entry {
63 19 : KeyAndValueRef::Ascii(key, value) => {
64 19 : if key.as_str().starts_with(HEADER_PREFIX) {
65 2 : continue;
66 17 : }
67 17 : let key = HeaderName::from_str(key.as_str());
68 17 : let value = HeaderValue::from_str(value.to_str().unwrap_or_default());
69 17 : if let (Ok(key), Ok(value)) = (key, value) {
70 17 : headers.insert(key, value);
71 17 : }
72 : }
73 : // Binary metadata (e.g. `grpc-status-details-bin`) cannot be
74 : // put into an HTTP header verbatim. Forward it base64-encoded
75 : // under the same `-bin` key so the frontend can decode it.
76 4 : KeyAndValueRef::Binary(key, value) => {
77 4 : if key.as_str().starts_with(HEADER_PREFIX) {
78 2 : continue;
79 2 : }
80 2 : let encoded = base64::engine::general_purpose::STANDARD
81 2 : .encode(value.to_bytes().unwrap_or_default());
82 2 : let key = HeaderName::from_str(key.as_str());
83 2 : let value = HeaderValue::from_str(&encoded);
84 2 : if let (Ok(key), Ok(value)) = (key, value) {
85 2 : headers.insert(key, value);
86 2 : }
87 : }
88 : }
89 : }
90 10 : }
91 :
92 : #[instrument(skip(req), fields(body_size))]
93 7 : pub async fn call_grpc_unary(channel_id: usize, mut req: Request<Vec<u8>>) -> Response<Vec<u8>> {
94 : let body = std::mem::take(req.body_mut());
95 : tracing::Span::current().record("body_size", body.len());
96 : match GRPC_PROXY.call_unary(channel_id, req.headers_mut(), body).await {
97 : Ok((body, metadata)) => {
98 : let mut response = Response::builder()
99 : .status(200)
100 : .header(HEADER_NAME_CHANNEL_ID, channel_id)
101 : .body(body)
102 : .unwrap();
103 : copy_metadata(&metadata, response.headers_mut());
104 : response
105 : },
106 : Err(e) => Response::from(&e)
107 : }
108 7 : }
109 :
110 : #[instrument(skip(req))]
111 8 : pub async fn start_grpc_server_stream(channel_id: usize, mut req: Request<Vec<u8>>) -> Response<Vec<u8>> {
112 : let body = std::mem::take(req.body_mut());
113 : match GRPC_PROXY.start_server_stream(channel_id, req.headers_mut(), body).await {
114 : Ok((stream_id, metadata)) => {
115 : let mut response = Response::builder()
116 : .status(200)
117 : .header(HEADER_NAME_CHANNEL_ID, channel_id)
118 : .header(HEADER_NAME_STREAM_ID, stream_id)
119 : .body(Vec::new())
120 : .unwrap();
121 : copy_metadata(&metadata, response.headers_mut());
122 : response
123 : },
124 : Err(e) => Response::from(&e)
125 : }
126 8 : }
127 :
128 12 : pub async fn read_grpc_server_stream(channel_id: usize, stream_id: usize, mut req: Request<Vec<u8>>) -> Response<Vec<u8>> {
129 12 : match GRPC_PROXY.read_server_stream(channel_id, stream_id, req.headers_mut()).await {
130 4 : Ok(batches) => {
131 4 : let mut buffer: Vec<u8> = Vec::with_capacity(batches.total_message_bytes + 4 * batches.messages.len());
132 4 : for message in batches.messages.iter() {
133 4 : buffer.write_u32::<LittleEndian>(message.len() as u32).unwrap();
134 4 : buffer.write(&message).unwrap();
135 4 : }
136 4 : let mut response = Response::builder()
137 4 : .status(200)
138 4 : .header(HEADER_NAME_CHANNEL_ID, channel_id)
139 4 : .header(HEADER_NAME_STREAM_ID, stream_id)
140 4 : .header(HEADER_NAME_BATCH_EVENT, batches.event.to_str())
141 4 : .header(HEADER_NAME_BATCH_MESSAGES, batches.messages.len())
142 4 : .header(HEADER_NAME_BATCH_BYTES, batches.total_message_bytes);
143 4 : if let Some(trailers) = batches.trailers {
144 2 : log::debug!("{:?}", trailers);
145 2 : copy_metadata(&trailers, response.headers_mut().unwrap());
146 2 : }
147 4 : response
148 4 : .body(buffer)
149 4 : .unwrap()
150 : },
151 8 : Err(e) => Response::from(&e)
152 : }
153 12 : }
154 :
155 2 : pub async fn delete_grpc_server_stream(channel_id: usize, stream_id: usize, _req: Request<Vec<u8>>) -> Response<Vec<u8>> {
156 2 : match GRPC_PROXY.destroy_server_stream(channel_id, stream_id).await {
157 : Ok(()) => {
158 2 : Response::builder()
159 2 : .status(200)
160 2 : .header(HEADER_NAME_CHANNEL_ID, channel_id)
161 2 : .body(Vec::new())
162 2 : .unwrap()
163 : }
164 0 : Err(e) => Response::from(&e)
165 : }
166 2 : }
167 :
168 : #[cfg(test)]
169 : mod test {
170 : use super::*;
171 : use tonic::metadata::{AsciiMetadataValue, BinaryMetadataValue};
172 :
173 2 : fn metadata() -> MetadataMap {
174 2 : let mut metadata = MetadataMap::new();
175 2 : metadata.insert("x-trace", AsciiMetadataValue::from_static("abc"));
176 2 : metadata.insert("dashql-batch-event", AsciiMetadataValue::from_static("StreamFinished"));
177 2 : metadata.insert_bin("grpc-status-details-bin", BinaryMetadataValue::from_bytes(&[0, 1, 255]));
178 2 : metadata.insert_bin("dashql-error-bin", BinaryMetadataValue::from_bytes(b"spoofed"));
179 2 : metadata
180 2 : }
181 :
182 : #[test]
183 1 : fn metadata_is_forwarded_and_binary_values_stay_base64() {
184 1 : let mut headers = HeaderMap::new();
185 1 : copy_metadata(&metadata(), &mut headers);
186 1 : assert_eq!(headers.get("x-trace").unwrap(), "abc");
187 1 : assert_eq!(headers.get("grpc-status-details-bin").unwrap(), "AAH/");
188 1 : }
189 :
190 : #[test]
191 1 : fn a_server_cannot_replace_headers_of_the_proxy() {
192 1 : let mut headers = HeaderMap::new();
193 1 : headers.insert(HEADER_NAME_BATCH_EVENT, HeaderValue::from_static("FlushAfterBytes"));
194 1 : copy_metadata(&metadata(), &mut headers);
195 1 : assert_eq!(headers.get(HEADER_NAME_BATCH_EVENT).unwrap(), "FlushAfterBytes");
196 1 : assert!(headers.get("dashql-error-bin").is_none());
197 1 : }
198 : }
|
