Line data Source code
1 : use std::collections::HashMap;
2 : use std::future::Future;
3 : use std::str::FromStr;
4 : use std::sync::atomic::AtomicUsize;
5 : use std::sync::Arc;
6 : use std::sync::RwLock;
7 : use std::result::Result;
8 : use std::time::Duration;
9 :
10 : use tonic::codegen::http::uri::PathAndQuery;
11 : use http::HeaderMap;
12 : use http::HeaderValue;
13 : use tonic::metadata::AsciiMetadataKey;
14 : use tonic::metadata::AsciiMetadataValue;
15 : use tonic::metadata::MetadataMap;
16 : use tonic::transport::channel::Endpoint;
17 : use http::HeaderName;
18 :
19 : use tonic::transport::ClientTlsConfig;
20 : use tonic::transport::Certificate;
21 : use tonic::transport::Identity;
22 :
23 : use crate::grpc_client::GenericGrpcClient;
24 : use crate::grpc_stream_manager::GrpcServerStreamBatch;
25 : use crate::grpc_stream_manager::GrpcStreamManager;
26 : use crate::proxy_headers::HEADER_NAME_BATCH_BYTES;
27 : use crate::proxy_headers::HEADER_NAME_BATCH_TIMEOUT;
28 : use crate::proxy_headers::HEADER_NAME_ENDPOINT;
29 : use crate::proxy_headers::HEADER_NAME_PATH;
30 : use crate::proxy_headers::HEADER_NAME_READ_TIMEOUT;
31 : use crate::proxy_headers::HEADER_NAME_TLS;
32 : use crate::proxy_headers::HEADER_NAME_TLS_CACERTS;
33 : use crate::proxy_headers::HEADER_NAME_TLS_CLIENT_CERT;
34 : use crate::proxy_headers::HEADER_NAME_TLS_CLIENT_KEY;
35 : use crate::proxy_headers::HEADER_PREFIX;
36 : use crate::status::Status;
37 :
38 : /// The TLS settings of a gRPC channel, read from the proxy request headers
39 : struct GrpcRequestTlsConfig {
40 : /// The path of the PEM client private key, required together with the client certificate
41 : client_key: Option<String>,
42 : /// The path of the PEM client certificate, required together with the client key
43 : client_cert: Option<String>,
44 : /// The path of a PEM file with extra CA certificates to trust on top of the OS roots
45 : cacerts: Option<String>,
46 : }
47 :
48 : /// The parameters for opening a gRPC channel, read from the proxy request headers
49 : struct GrpcChannelParams {
50 : /// The endpoint to connect to
51 : endpoint: Endpoint,
52 : /// The TLS settings, `None` for a plaintext channel
53 : tls: Option<GrpcRequestTlsConfig>,
54 : }
55 :
56 2 : async fn read_tls_file(path: &str, label: &'static str) -> Result<Vec<u8>, Status> {
57 2 : tokio::fs::read(path)
58 2 : .await
59 2 : .map_err(|e| Status::GrpcTlsConfigInvalid {
60 2 : message: format!("failed to read {} from '{}': {}", label, path, e),
61 2 : })
62 2 : }
63 :
64 13 : fn connect_endpoint(
65 13 : endpoint: Endpoint,
66 13 : ) -> impl Future<Output = Result<tonic::transport::Channel, Status>> {
67 13 : async move {
68 13 : endpoint
69 13 : .connect()
70 13 : .await
71 13 : .map_err(|e| {
72 2 : log::warn!("Failed creating channel: {:?}", e);
73 2 : Status::GrpcEndpointConnectFailed{ message: e.to_string() }
74 2 : })
75 13 : }
76 13 : }
77 :
78 5 : async fn connect_tls_endpoint(
79 5 : endpoint: Endpoint,
80 5 : tls: GrpcRequestTlsConfig,
81 5 : ) -> Result<tonic::transport::Channel, Status> {
82 : // Seed with the OS trust store so servers using publicly-trusted (or
83 : // corporate) CAs validate out of the box. An explicit cacerts header can
84 : // still add its own trust anchor on top.
85 5 : let mut tls_config = ClientTlsConfig::new().with_native_roots();
86 :
87 5 : if let Some(cacerts_path) = tls.cacerts.as_deref() {
88 1 : let pem = read_tls_file(cacerts_path, "ca certificates").await?;
89 0 : let cert = Certificate::from_pem(pem);
90 0 : tls_config = tls_config.ca_certificate(cert);
91 4 : }
92 :
93 4 : match (tls.client_cert.as_deref(), tls.client_key.as_deref()) {
94 1 : (Some(client_cert_path), Some(client_key_path)) => {
95 1 : let client_cert = read_tls_file(client_cert_path, "client certificate").await?;
96 0 : let client_key = read_tls_file(client_key_path, "client private key").await?;
97 0 : let identity = Identity::from_pem(client_cert, client_key);
98 0 : tls_config = tls_config.identity(identity);
99 : }
100 : (Some(_), None) => {
101 1 : return Err(Status::HeaderRequiredButMissing { header: HEADER_NAME_TLS_CLIENT_KEY });
102 : }
103 : (None, Some(_)) => {
104 1 : return Err(Status::HeaderRequiredButMissing { header: HEADER_NAME_TLS_CLIENT_CERT });
105 : }
106 1 : (None, None) => {}
107 : }
108 :
109 1 : let endpoint = endpoint
110 1 : .tls_config(tls_config)
111 1 : .map_err(|e| Status::GrpcTlsConfigInvalid { message: e.to_string() })?;
112 1 : connect_endpoint(endpoint).await
113 5 : }
114 :
115 : /// An open gRPC channel registered in the [`GrpcProxy`]
116 : pub struct GrpcChannelEntry {
117 : /// The connected tonic channel
118 : pub channel: tonic::transport::Channel,
119 : }
120 :
121 : /// The gRPC proxy that manages open channels and their server streams
122 : pub struct GrpcProxy {
123 : /// The id to assign to the next created channel, starting at 1
124 : pub next_channel_id: AtomicUsize,
125 : /// The open channels, keyed by channel id
126 : pub channels: RwLock<HashMap<usize, Arc<GrpcChannelEntry>>>,
127 : /// The server streams opened on the channels
128 : pub streams: Arc<GrpcStreamManager>,
129 : }
130 :
131 : impl Default for GrpcProxy {
132 1 : fn default() -> Self {
133 1 : Self {
134 1 : next_channel_id: AtomicUsize::new(1),
135 1 : channels: Default::default(),
136 1 : streams: Default::default(),
137 1 : }
138 1 : }
139 : }
140 :
141 : /// Helper to copy incoming metadata
142 8 : fn prepare_metadata(out: &mut MetadataMap, headers: &mut HeaderMap) {
143 11 : for (key, value) in headers.drain() {
144 11 : let key = match &key {
145 11 : Some(k) => k.as_str(),
146 0 : None => continue,
147 : };
148 11 : if !key.starts_with(HEADER_PREFIX) {
149 3 : if let Ok(value) = value.to_str() {
150 3 : let k = AsciiMetadataKey::from_str(key);
151 3 : let v = AsciiMetadataValue::from_str(value);
152 3 : if let (Ok(k), Ok(v)) = (k, v) {
153 3 : out.insert(k, v);
154 3 : } else {
155 0 : log::warn!("Failed adding extra metadata with key: {}", key);
156 : }
157 0 : }
158 8 : }
159 : }
160 : // TODO Data Cloud is returning PermissionDenied if the origin header is set, find out why
161 24 : for masked in vec!["origin", "accept", "referrer"] {
162 24 : out.remove(masked);
163 24 : }
164 8 : }
165 :
166 : /// Helper to unpack parameters for a gRPC channel
167 19 : fn read_channel_params(headers: &mut HeaderMap) -> Result<GrpcChannelParams, Status> {
168 19 : let mut host = None;
169 19 : let mut tls = false;
170 19 : let mut tls_client_key = None;
171 19 : let mut tls_client_cert = None;
172 19 : let mut tls_cacerts = None;
173 :
174 : // Helper to unpack a header value
175 23 : let read_header_value = |value: HeaderValue, header: &'static str| -> Result<String, Status> {
176 23 : Ok(value
177 23 : .to_str()
178 23 : .map_err(|e| Status::HeaderHasInvalidEncoding{ header, message: e.to_string() })?
179 23 : .to_string())
180 23 : };
181 :
182 : // Read all headers in the request, pick up the one from us and declare the remaining as extra
183 24 : for (key, value) in headers.drain() {
184 24 : let key = match &key {
185 24 : Some(k) => k.as_str(),
186 0 : None => continue,
187 : };
188 24 : match key {
189 24 : HEADER_NAME_ENDPOINT => {
190 18 : host = Some(read_header_value(value, HEADER_NAME_ENDPOINT)?);
191 : }
192 6 : HEADER_NAME_TLS => {
193 1 : tls = true;
194 1 : }
195 5 : HEADER_NAME_TLS_CLIENT_KEY => {
196 2 : tls = true;
197 2 : tls_client_key = Some(read_header_value(value, HEADER_NAME_TLS_CLIENT_KEY)?);
198 : }
199 3 : HEADER_NAME_TLS_CLIENT_CERT => {
200 2 : tls = true;
201 2 : tls_client_cert = Some(read_header_value(value, HEADER_NAME_TLS_CLIENT_CERT)?);
202 : }
203 1 : HEADER_NAME_TLS_CACERTS => {
204 1 : tls = true;
205 1 : tls_cacerts = Some(read_header_value(value, HEADER_NAME_TLS_CACERTS)?);
206 : }
207 : _ => {
208 0 : if !key.starts_with(HEADER_PREFIX) && HeaderName::try_from(key.to_string()).is_err() {
209 0 : log::warn!("Failed adding extra metadata with key: {}", key);
210 0 : }
211 : }
212 : }
213 : }
214 :
215 : // Make sure the user provided an endpoint
216 19 : let endpoint = if let Some(host) = &host {
217 18 : Endpoint::from_str(host).map_err(|e| Status::HeaderIsNotAValidEndpoint { header: HEADER_NAME_ENDPOINT, message: e.to_string() })
218 1 : ?
219 : } else {
220 1 : return Err(Status::HeaderRequiredButMissing { header: HEADER_NAME_ENDPOINT });
221 : };
222 17 : let tls_config = if tls {
223 5 : Some(GrpcRequestTlsConfig {
224 5 : client_key: tls_client_key,
225 5 : client_cert: tls_client_cert,
226 5 : cacerts: tls_cacerts,
227 5 : })
228 : } else {
229 12 : None
230 : };
231 :
232 17 : Ok(GrpcChannelParams { endpoint, tls: tls_config })
233 19 : }
234 :
235 : /// Helper to read a string from request headers
236 15 : fn require_string_header(headers: &HeaderMap, header_name: &'static str) -> Result<String, Status> {
237 15 : if let Some(header) = headers.get(header_name) {
238 13 : let header = header
239 13 : .to_str()
240 13 : .map_err(|e| Status::HeaderHasInvalidEncoding{ header: header_name, message: e.to_string() })?
241 13 : .to_string();
242 13 : Ok(header)
243 : } else {
244 2 : Err(Status::HeaderRequiredButMissing { header: header_name })
245 : }
246 15 : }
247 :
248 : /// Helper to read a string from request headers
249 29 : fn require_usize_header(headers: &HeaderMap, header_name: &'static str) -> Result<usize, Status> {
250 29 : if let Some(header) = headers.get(header_name) {
251 28 : let header = header
252 28 : .to_str()
253 28 : .map_err(|e| Status::HeaderHasInvalidEncoding{ header: header_name, message: e.to_string() })?
254 28 : .to_string();
255 28 : let header: usize = header.parse::<usize>().map_err(|e| Status::HeaderIsNotAnUsize { header: header_name, message: e.to_string() })?;
256 27 : Ok(header)
257 : } else {
258 1 : Err(Status::HeaderRequiredButMissing { header: header_name })
259 : }
260 29 : }
261 :
262 : impl GrpcProxy {
263 : /// Create a channel
264 19 : pub async fn create_channel(&self, headers: &mut HeaderMap) -> Result<usize, Status> {
265 19 : let channel_id = self
266 19 : .next_channel_id
267 19 : .fetch_add(1, std::sync::atomic::Ordering::SeqCst);
268 19 : let params = read_channel_params(headers)?;
269 17 : let channel = if let Some(tls) = params.tls {
270 5 : connect_tls_endpoint(params.endpoint, tls).await?
271 : } else {
272 12 : connect_endpoint(params.endpoint).await?
273 : };
274 11 : if let Ok(mut channels) = self.channels.write() {
275 11 : channels.insert(channel_id, Arc::new(GrpcChannelEntry { channel }));
276 11 : }
277 11 : Ok(channel_id)
278 19 : }
279 : /// Destroy a channel
280 1 : pub async fn destroy_channel(&self, channel_id: usize) -> Result<(), Status> {
281 1 : if let Ok(mut channels) = self.channels.write() {
282 1 : channels.remove(&channel_id);
283 1 : }
284 1 : Ok(())
285 1 : }
286 :
287 : /// Call a unary gRPC function
288 7 : pub async fn call_unary(&self, channel_id: usize, headers: &mut HeaderMap, body: Vec<u8>) -> Result<(Vec<u8>, MetadataMap), Status> {
289 7 : let path = require_string_header(headers, HEADER_NAME_PATH)?;
290 6 : let channel_entry = if let Some(channel) = self.channels.read().unwrap().get(&channel_id) {
291 4 : channel.clone()
292 : } else {
293 2 : return Err(Status::GrpcChannelIdIsUnknown { channel_id });
294 : };
295 4 : let mut client = GenericGrpcClient::new(channel_entry.channel.clone());
296 4 : let path = PathAndQuery::from_str(&path)
297 4 : .map_err(|e| {
298 1 : Status::HeaderPathIsInvalid { header: HEADER_NAME_PATH, path: path.to_string(), message: e.to_string() }
299 1 : })?;
300 :
301 3 : let mut request = tonic::Request::new(body);
302 3 : prepare_metadata(request.metadata_mut(), headers);
303 3 : let mut response = client.call_unary(request, path).await
304 3 : .map_err(|status| {
305 2 : log::warn!("Grpc call failed: {:?}", status);
306 2 : Status::GrpcCallFailed { status }
307 2 : })?;
308 1 : let metadata = std::mem::take(response.metadata_mut());
309 1 : let body = response.into_inner();
310 1 : Ok((body, metadata))
311 7 : }
312 :
313 : /// Call a gRPC function with results streamed from the server
314 8 : pub async fn start_server_stream(&self, channel_id: usize, headers: &mut HeaderMap, body: Vec<u8>) -> Result<(usize, MetadataMap), Status> {
315 8 : let path = require_string_header(headers, HEADER_NAME_PATH)?;
316 7 : let channel_entry = if let Some(channel) = self.channels.read().unwrap().get(&channel_id) {
317 6 : channel.clone()
318 : } else {
319 1 : return Err(Status::GrpcChannelIdIsUnknown { channel_id });
320 : };
321 :
322 : // Send the gRPC request
323 6 : let mut client = GenericGrpcClient::new(channel_entry.channel.clone());
324 6 : let path = PathAndQuery::from_str(&path)
325 6 : .map_err(|e| {
326 1 : Status::HeaderPathIsInvalid { header: HEADER_NAME_PATH, path: path.to_string(), message: e.to_string() }
327 1 : })?;
328 5 : let mut request = tonic::Request::new(body);
329 5 : prepare_metadata(request.metadata_mut(), headers);
330 5 : log::debug!("Sending gRPC request: channel_id={}, path={}, request={:?}", channel_id, path.to_string(), &request);
331 :
332 5 : let mut response = client.call_server_streaming(request, path).await
333 5 : .map_err(|status| {
334 0 : log::warn!("Grpc call failed: {:?}", status);
335 0 : Status::GrpcCallFailed { status }
336 0 : })?;
337 :
338 : // Save the response headers
339 5 : let response_headers = std::mem::take(response.metadata_mut());
340 :
341 : // Register the output stream
342 5 : let streaming = response.into_inner();
343 5 : let stream_id = self.streams.start_server_stream(channel_id, streaming)?;
344 :
345 5 : Ok((stream_id, response_headers))
346 8 : }
347 :
348 : /// Read from a result stream
349 12 : pub async fn read_server_stream(&self, channel_id: usize, stream_id: usize, headers: &mut HeaderMap) -> Result<GrpcServerStreamBatch, Status> {
350 : // We don't need the channel id to resolve the stream today since the stream id is unique across all channels.
351 : // We still check if the channel id exists so that we can still maintain streams per channel later.
352 12 : if self.channels.read().unwrap().get(&channel_id).is_none() {
353 1 : return Err(Status::GrpcChannelIdIsUnknown { channel_id });
354 11 : }
355 : // Read limits from request headers
356 11 : let read_timeout = require_usize_header(headers, HEADER_NAME_READ_TIMEOUT)?;
357 9 : let batch_timeout = require_usize_header(headers, HEADER_NAME_BATCH_TIMEOUT)?;
358 9 : let batch_bytes = require_usize_header(headers, HEADER_NAME_BATCH_BYTES)?;
359 :
360 : // Read from the stream
361 9 : let read_timeout_duration = Duration::from_millis(read_timeout as u64);
362 9 : let batch_timeout_duration = Duration::from_millis(batch_timeout as u64);
363 9 : let read_result = self.streams.read_server_stream(channel_id, stream_id, read_timeout_duration, batch_timeout_duration, batch_bytes).await?;
364 4 : Ok(read_result)
365 12 : }
366 :
367 : /// Destroy a result steram
368 2 : pub async fn destroy_server_stream(&self, _channel_id: usize, stream_id: usize) -> Result<(), Status> {
369 2 : self.streams.destroy_server_stream(stream_id).await;
370 2 : Ok(())
371 2 : }
372 : }
|
