Line data Source code
1 : //! The gRPC proxy through the routes the app calls, against the mock of the test service.
2 :
3 : use std::net::SocketAddr;
4 : use std::time::Duration;
5 :
6 : use http::{Request, Response};
7 : use prost::Message;
8 :
9 : use crate::proto::dashql_test::{TestServerStreamingRequest, TestServerStreamingResponse, TestUnaryRequest, TestUnaryResponse};
10 : use crate::route_proxy_request;
11 : use crate::test::grpc_service_mock::{spawn_grpc_test_service_mock, GrpcServiceMock};
12 :
13 : const UNARY: &str = "/dashql.test.TestService/TestUnary";
14 : const STREAMING: &str = "/dashql.test.TestService/TestServerStreaming";
15 : const GUARD: Duration = Duration::from_secs(10);
16 :
17 51 : fn request(method: &str, path: &str, headers: &[(&str, &str)], body: &[u8]) -> Request<Vec<u8>> {
18 51 : let mut builder = Request::builder().method(method).uri(format!("dashql-native://localhost{path}"));
19 73 : for (name, value) in headers {
20 73 : builder = builder.header(*name, *value);
21 73 : }
22 51 : builder.body(body.to_vec()).unwrap()
23 51 : }
24 :
25 49 : fn header<'a>(response: &'a Response<Vec<u8>>, name: &str) -> Option<&'a str> {
26 49 : response.headers().get(name).map(|value| value.to_str().unwrap())
27 49 : }
28 :
29 20 : fn error(response: &Response<Vec<u8>>) -> serde_json::Value {
30 20 : assert_eq!(header(response, "dashql-error"), Some("true"), "{response:?}");
31 20 : serde_json::from_slice(response.body()).unwrap()
32 20 : }
33 :
34 11 : async fn create_channel(address: SocketAddr) -> String {
35 11 : let endpoint = format!("http://{address}");
36 11 : let response = route_proxy_request(request("POST", "/grpc/channels", &[("dashql-endpoint", endpoint.as_str())], &[])).await;
37 11 : assert_eq!(response.status(), 200, "{response:?}");
38 11 : header(&response, "dashql-channel-id").unwrap().to_string()
39 11 : }
40 :
41 5 : async fn start_stream(channel: &str, extra: &[(&str, &str)]) -> Response<Vec<u8>> {
42 5 : let mut headers = vec![("dashql-path", STREAMING)];
43 5 : headers.extend_from_slice(extra);
44 5 : let body = TestServerStreamingRequest { data: "request".to_string() }.encode_to_vec();
45 5 : route_proxy_request(request("POST", &format!("/grpc/channel/{channel}/streams"), &headers, &body)).await
46 5 : }
47 :
48 5 : fn stream_id(response: &Response<Vec<u8>>) -> String {
49 5 : assert_eq!(response.status(), 200, "{response:?}");
50 5 : header(response, "dashql-stream-id").unwrap().to_string()
51 5 : }
52 :
53 10 : async fn read_stream(channel: &str, stream: &str, read_ms: &str, batch_ms: &str, bytes: &str) -> Response<Vec<u8>> {
54 10 : let headers = [("dashql-read-timeout", read_ms), ("dashql-batch-timeout", batch_ms), ("dashql-batch-bytes", bytes)];
55 10 : route_proxy_request(request("GET", &format!("/grpc/channel/{channel}/stream/{stream}"), &headers, &[])).await
56 10 : }
57 :
58 : /// The messages of a batch, each with its length as a little endian prefix
59 3 : fn messages(body: &[u8]) -> Vec<String> {
60 3 : let mut messages = Vec::new();
61 3 : let mut position = 0;
62 7 : while position < body.len() {
63 4 : let length = u32::from_le_bytes(body[position..position + 4].try_into().unwrap()) as usize;
64 4 : position += 4;
65 4 : messages.push(TestServerStreamingResponse::decode(&body[position..position + length]).unwrap().data);
66 4 : position += length;
67 4 : }
68 3 : messages
69 3 : }
70 :
71 20 : fn message(data: &str) -> Result<TestServerStreamingResponse, tonic::Status> {
72 20 : Ok(TestServerStreamingResponse { data: data.to_string() })
73 20 : }
74 :
75 : #[tokio::test]
76 1 : async fn a_channel_needs_an_endpoint() {
77 1 : let response = route_proxy_request(request("POST", "/grpc/channels", &[], &[])).await;
78 1 : assert_eq!(response.status(), 400);
79 1 : let body = error(&response);
80 1 : assert_eq!(body["message"], "header is missing");
81 1 : assert_eq!(body["details"]["header"], "dashql-endpoint");
82 1 : }
83 :
84 : #[tokio::test]
85 1 : async fn rejects_an_invalid_endpoint() {
86 1 : let response = route_proxy_request(request("POST", "/grpc/channels", &[("dashql-endpoint", "not a url")], &[])).await;
87 1 : assert_eq!(error(&response)["message"], "header is not a valid endpoint");
88 1 : }
89 :
90 : #[tokio::test]
91 1 : async fn reports_an_endpoint_that_cannot_be_reached() {
92 1 : let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
93 1 : let endpoint = format!("http://{}", listener.local_addr().unwrap());
94 1 : drop(listener);
95 1 : let response = route_proxy_request(request("POST", "/grpc/channels", &[("dashql-endpoint", endpoint.as_str())], &[])).await;
96 1 : assert_eq!(error(&response)["message"], "connecting to gRPC endpoint failed");
97 1 : }
98 :
99 : #[tokio::test]
100 1 : async fn a_client_certificate_needs_its_key_and_the_key_its_certificate() {
101 1 : let endpoint = ("dashql-endpoint", "https://127.0.0.1:1");
102 1 : let response = route_proxy_request(request("POST", "/grpc/channels", &[endpoint, ("dashql-tls-client-cert", "/cert.pem")], &[])).await;
103 1 : assert_eq!(error(&response)["details"]["header"], "dashql-tls-client-key");
104 1 : let response = route_proxy_request(request("POST", "/grpc/channels", &[endpoint, ("dashql-tls-client-key", "/key.pem")], &[])).await;
105 1 : assert_eq!(error(&response)["details"]["header"], "dashql-tls-client-cert");
106 1 : }
107 :
108 : #[tokio::test]
109 1 : async fn reports_tls_files_that_cannot_be_read() {
110 1 : let endpoint = ("dashql-endpoint", "https://127.0.0.1:1");
111 2 : for headers in [
112 1 : vec![endpoint, ("dashql-tls-cacerts", "/missing/ca.pem")],
113 1 : vec![endpoint, ("dashql-tls-client-cert", "/missing/cert.pem"), ("dashql-tls-client-key", "/missing/key.pem")],
114 1 : ] {
115 2 : let body = error(&route_proxy_request(request("POST", "/grpc/channels", &headers, &[])).await);
116 2 : assert_eq!(body["message"], "gRPC tls config is invalid");
117 2 : assert!(body["details"]["error"].as_str().unwrap().contains("/missing/"), "{body}");
118 1 : }
119 1 : }
120 :
121 : #[tokio::test]
122 1 : async fn tls_without_a_server_fails_to_connect() {
123 1 : let headers = [("dashql-endpoint", "https://127.0.0.1:1"), ("dashql-tls", "1")];
124 1 : let body = error(&route_proxy_request(request("POST", "/grpc/channels", &headers, &[])).await);
125 1 : assert_eq!(body["message"], "connecting to gRPC endpoint failed");
126 1 : }
127 :
128 : #[tokio::test]
129 1 : async fn calls_a_unary_function() {
130 1 : let (mock, mut setup_unary, _setup_streaming) = GrpcServiceMock::new();
131 1 : let metadata = mock.received_metadata.clone();
132 1 : let (address, shutdown) = spawn_grpc_test_service_mock(mock).await;
133 1 : let channel = create_channel(address).await;
134 :
135 1 : let server = tokio::spawn(async move {
136 1 : let (param, sender) = setup_unary.recv().await.unwrap();
137 1 : sender.send(Ok(TestUnaryResponse { data: "response".to_string() })).await.unwrap();
138 1 : param
139 1 : });
140 1 : let body = TestUnaryRequest { data: "request".to_string() }.encode_to_vec();
141 1 : let response = route_proxy_request(request("POST", &format!("/grpc/channel/{channel}/unary"), &[
142 1 : ("dashql-path", UNARY),
143 1 : ("authorization", "Bearer token"),
144 1 : ("origin", "http://localhost"),
145 1 : ("accept", "*/*"),
146 1 : ], &body)).await;
147 :
148 1 : assert_eq!(response.status(), 200, "{response:?}");
149 1 : assert_eq!(header(&response, "dashql-channel-id"), Some(channel.as_str()));
150 1 : assert_eq!(TestUnaryResponse::decode(response.body().as_slice()).unwrap().data, "response");
151 1 : assert_eq!(server.await.unwrap().data, "request");
152 :
153 1 : let metadata = metadata.lock().unwrap()[0].clone();
154 1 : assert_eq!(metadata.get("authorization").unwrap(), "Bearer token");
155 : // Data Cloud refuses these, and our headers stop at the proxy
156 3 : for masked in ["origin", "accept", "dashql-path"] {
157 3 : assert!(metadata.get(masked).is_none(), "{masked}");
158 1 : }
159 1 : shutdown.send(()).unwrap();
160 1 : }
161 :
162 : #[tokio::test]
163 1 : async fn forwards_the_rich_error_of_a_unary_function() {
164 1 : let (mock, mut setup_unary, _setup_streaming) = GrpcServiceMock::new();
165 1 : let (address, shutdown) = spawn_grpc_test_service_mock(mock).await;
166 1 : let channel = create_channel(address).await;
167 1 : tokio::spawn(async move {
168 1 : let (_, sender) = setup_unary.recv().await.unwrap();
169 1 : let status = tonic::Status::with_details(tonic::Code::InvalidArgument, "bad query", vec![1, 2, 3].into());
170 1 : sender.send(Err(status)).await.unwrap();
171 1 : });
172 1 : let response = route_proxy_request(request("POST", &format!("/grpc/channel/{channel}/unary"), &[("dashql-path", UNARY)], &[])).await;
173 1 : assert_eq!(response.status(), 400);
174 1 : assert_eq!(header(&response, "dashql-grpc-status"), Some("3"));
175 1 : assert_eq!(header(&response, "dashql-grpc-message"), Some("bad query"));
176 1 : assert_eq!(header(&response, "dashql-grpc-status-details-bin"), Some("AQID"));
177 1 : assert_eq!(error(&response)["message"], "gRPC call failed");
178 1 : shutdown.send(()).unwrap();
179 1 : }
180 :
181 : #[tokio::test]
182 1 : async fn survives_a_database_error_that_spans_lines() {
183 1 : let (mock, mut setup_unary, _setup_streaming) = GrpcServiceMock::new();
184 1 : let (address, shutdown) = spawn_grpc_test_service_mock(mock).await;
185 1 : let channel = create_channel(address).await;
186 1 : tokio::spawn(async move {
187 1 : let (_, sender) = setup_unary.recv().await.unwrap();
188 1 : let status = tonic::Status::new(tonic::Code::InvalidArgument, "syntax error at line 1\nnear \"caf\u{e9}\"");
189 1 : sender.send(Err(status)).await.unwrap();
190 1 : });
191 1 : let response = route_proxy_request(request("POST", &format!("/grpc/channel/{channel}/unary"), &[("dashql-path", UNARY)], &[])).await;
192 1 : assert_eq!(response.status(), 400);
193 1 : assert_eq!(header(&response, "dashql-grpc-status"), Some("3"));
194 1 : assert!(error(&response)["details"]["error"].as_str().unwrap().contains("near"));
195 1 : shutdown.send(()).unwrap();
196 1 : }
197 :
198 : #[tokio::test]
199 1 : async fn validates_a_call_before_sending_it() {
200 1 : let (mock, _setup_unary, _setup_streaming) = GrpcServiceMock::new();
201 1 : let (address, shutdown) = spawn_grpc_test_service_mock(mock).await;
202 1 : let channel = create_channel(address).await;
203 :
204 2 : for route in ["unary", "streams"] {
205 2 : let path = format!("/grpc/channel/{channel}/{route}");
206 2 : let body = error(&route_proxy_request(request("POST", &path, &[], &[])).await);
207 2 : assert_eq!(body["message"], "header is missing");
208 2 : let body = error(&route_proxy_request(request("POST", &path, &[("dashql-path", "no path")], &[])).await);
209 2 : assert_eq!(body["message"], "header stores invalid path");
210 1 : }
211 1 : shutdown.send(()).unwrap();
212 1 : }
213 :
214 : #[tokio::test]
215 1 : async fn refuses_an_unknown_channel() {
216 2 : for route in ["unary", "streams"] {
217 2 : let response = route_proxy_request(request("POST", &format!("/grpc/channel/999999/{route}"), &[("dashql-path", UNARY)], &[])).await;
218 2 : assert_eq!(response.status(), 404);
219 2 : assert_eq!(error(&response)["message"], "gRPC channel id is unknown");
220 1 : }
221 1 : assert_eq!(read_stream("999999", "0", "1", "1", "1").await.status(), 404);
222 1 : }
223 :
224 : #[tokio::test]
225 1 : async fn a_deleted_channel_is_gone() {
226 1 : let (mock, _setup_unary, _setup_streaming) = GrpcServiceMock::new();
227 1 : let (address, shutdown) = spawn_grpc_test_service_mock(mock).await;
228 1 : let channel = create_channel(address).await;
229 1 : assert_eq!(route_proxy_request(request("DELETE", &format!("/grpc/channel/{channel}"), &[], &[])).await.status(), 200);
230 1 : let response = route_proxy_request(request("POST", &format!("/grpc/channel/{channel}/unary"), &[("dashql-path", UNARY)], &[])).await;
231 1 : assert_eq!(response.status(), 404);
232 1 : shutdown.send(()).unwrap();
233 1 : }
234 :
235 : #[tokio::test]
236 1 : async fn rejects_methods_a_route_does_not_have() {
237 1 : assert_eq!(route_proxy_request(request("GET", "/grpc/channels", &[], &[])).await.status(), 404);
238 1 : assert_eq!(route_proxy_request(request("GET", "/grpc/channel/1/unary", &[], &[])).await.status(), 404);
239 1 : }
240 :
241 : #[tokio::test]
242 1 : async fn streams_batches_until_the_end() {
243 1 : let (mock, _setup_unary, mut setup_streaming) = GrpcServiceMock::new();
244 1 : let (address, shutdown) = spawn_grpc_test_service_mock(mock).await;
245 1 : let channel = create_channel(address).await;
246 :
247 1 : let server = tokio::spawn(async move {
248 1 : let (param, sender) = setup_streaming.recv().await.unwrap();
249 1 : sender.send(message("one")).await.unwrap();
250 1 : sender.send(message("two")).await.unwrap();
251 1 : param
252 1 : });
253 1 : let start = start_stream(&channel, &[]).await;
254 1 : let stream = stream_id(&start);
255 1 : assert_eq!(header(&start, "dashql-channel-id"), Some(channel.as_str()));
256 1 : assert_eq!(server.await.unwrap().data, "request");
257 :
258 1 : let response = read_stream(&channel, &stream, "5000", "5000", "1000000").await;
259 1 : assert_eq!(response.status(), 200, "{response:?}");
260 1 : assert_eq!(header(&response, "dashql-batch-event"), Some("StreamFinished"));
261 1 : assert_eq!(header(&response, "dashql-batch-messages"), Some("2"));
262 1 : assert_eq!(messages(response.body()), vec!["one", "two"]);
263 :
264 1 : assert_eq!(read_stream(&channel, &stream, "10", "10", "10").await.status(), 404);
265 1 : shutdown.send(()).unwrap();
266 1 : }
267 :
268 : #[tokio::test]
269 1 : async fn flushes_a_batch_after_the_size_limit_and_after_the_timeout() {
270 1 : let (mock, _setup_unary, mut setup_streaming) = GrpcServiceMock::new();
271 1 : let (address, shutdown) = spawn_grpc_test_service_mock(mock).await;
272 1 : let channel = create_channel(address).await;
273 1 : let stream = {
274 1 : let server = tokio::spawn(async move { setup_streaming.recv().await.unwrap() });
275 1 : let stream = stream_id(&start_stream(&channel, &[]).await);
276 1 : let (_, sender) = server.await.unwrap();
277 :
278 1 : sender.send(message("a long enough message")).await.unwrap();
279 1 : let response = read_stream(&channel, &stream, "5000", "5000", "4").await;
280 1 : assert_eq!(header(&response, "dashql-batch-event"), Some("FlushAfterBytes"));
281 1 : assert_eq!(messages(response.body()), vec!["a long enough message"]);
282 :
283 1 : sender.send(message("partial")).await.unwrap();
284 1 : let response = read_stream(&channel, &stream, "5000", "100", "1000000").await;
285 1 : assert_eq!(header(&response, "dashql-batch-event"), Some("FlushAfterTimeout"));
286 1 : assert_eq!(messages(response.body()), vec!["partial"]);
287 :
288 : // Nothing arrives within the read timeout
289 1 : let response = read_stream(&channel, &stream, "50", "50", "1000").await;
290 1 : assert_eq!(error(&response)["message"], "reading from gRPC stream timed out");
291 1 : drop(sender);
292 1 : stream
293 : };
294 1 : let response = read_stream(&channel, &stream, "5000", "5000", "1000").await;
295 1 : assert_eq!(header(&response, "dashql-batch-event"), Some("StreamFinished"));
296 1 : shutdown.send(()).unwrap();
297 1 : }
298 :
299 : #[tokio::test]
300 1 : async fn forwards_the_rich_error_of_a_failing_stream() {
301 1 : let (mock, _setup_unary, mut setup_streaming) = GrpcServiceMock::new();
302 1 : let (address, shutdown) = spawn_grpc_test_service_mock(mock).await;
303 1 : let channel = create_channel(address).await;
304 1 : let server = tokio::spawn(async move { setup_streaming.recv().await.unwrap() });
305 1 : let stream = stream_id(&start_stream(&channel, &[]).await);
306 1 : let (_, sender) = server.await.unwrap();
307 1 : sender.send(message("held back")).await.unwrap();
308 1 : let status = tonic::Status::with_details(tonic::Code::Internal, "failed\nmidway", vec![9].into());
309 1 : sender.send(Err(status)).await.unwrap();
310 :
311 1 : let response = read_stream(&channel, &stream, "5000", "5000", "1000000").await;
312 1 : assert_eq!(response.status(), 400);
313 1 : assert_eq!(header(&response, "dashql-grpc-status"), Some("13"));
314 1 : assert_eq!(header(&response, "dashql-grpc-status-details-bin"), Some("CQ=="));
315 1 : assert_eq!(error(&response)["message"], "reading from gRPC stream failed");
316 1 : assert_eq!(read_stream(&channel, &stream, "10", "10", "10").await.status(), 404);
317 1 : shutdown.send(()).unwrap();
318 1 : }
319 :
320 : #[tokio::test]
321 1 : async fn requires_the_limits_of_a_read() {
322 1 : let (mock, _setup_unary, _setup_streaming) = GrpcServiceMock::new();
323 1 : let (address, shutdown) = spawn_grpc_test_service_mock(mock).await;
324 1 : let channel = create_channel(address).await;
325 1 : let path = format!("/grpc/channel/{channel}/stream/0");
326 1 : assert_eq!(error(&route_proxy_request(request("GET", &path, &[], &[])).await)["message"], "header is missing");
327 1 : let headers = [("dashql-read-timeout", "soon"), ("dashql-batch-timeout", "1"), ("dashql-batch-bytes", "1")];
328 1 : assert_eq!(error(&route_proxy_request(request("GET", &path, &headers, &[])).await)["message"], "header is not an unsigned integer");
329 1 : shutdown.send(()).unwrap();
330 1 : }
331 :
332 : #[tokio::test]
333 1 : async fn deleting_a_stream_cancels_a_query_that_is_silent() {
334 1 : let (mock, _setup_unary, mut setup_streaming) = GrpcServiceMock::new();
335 1 : let (address, shutdown) = spawn_grpc_test_service_mock(mock).await;
336 1 : let channel = create_channel(address).await;
337 1 : let server = tokio::spawn(async move { setup_streaming.recv().await.unwrap() });
338 1 : let stream = stream_id(&start_stream(&channel, &[]).await);
339 1 : let (_, sender) = server.await.unwrap();
340 :
341 1 : let response = route_proxy_request(request("DELETE", &format!("/grpc/channel/{channel}/stream/{stream}"), &[], &[])).await;
342 1 : assert_eq!(response.status(), 200);
343 1 : assert_eq!(read_stream(&channel, &stream, "10", "10", "10").await.status(), 404);
344 1 : tokio::time::timeout(GUARD, sender.closed()).await.expect("the server was not told to stop");
345 1 : shutdown.send(()).unwrap();
346 1 : }
347 :
348 : #[tokio::test]
349 1 : async fn deleting_a_stream_cancels_a_query_with_a_full_queue() {
350 1 : let (mock, _setup_unary, mut setup_streaming) = GrpcServiceMock::new();
351 1 : let (address, shutdown) = spawn_grpc_test_service_mock(mock).await;
352 1 : let channel = create_channel(address).await;
353 1 : let server = tokio::spawn(async move { setup_streaming.recv().await.unwrap() });
354 1 : let stream = stream_id(&start_stream(&channel, &[]).await);
355 1 : let (_, sender) = server.await.unwrap();
356 :
357 : // More messages than the queue holds, nobody reads them
358 15 : for i in 0..15 {
359 15 : sender.send(message(&format!("message {i}"))).await.unwrap();
360 : }
361 1 : tokio::time::sleep(Duration::from_millis(50)).await;
362 1 : route_proxy_request(request("DELETE", &format!("/grpc/channel/{channel}/stream/{stream}"), &[], &[])).await;
363 1 : tokio::time::timeout(GUARD, sender.closed()).await.expect("the server was not told to stop");
364 1 : shutdown.send(()).unwrap();
365 1 : }
|
