mongreldb_protocol/
native_transport.rs

1//! Native gRPC transport over multiplexed HTTP/2 with certificate validation.
2
3use std::future::Future;
4use std::io::{BufReader, Error, ErrorKind};
5use std::net::SocketAddr;
6use std::pin::Pin;
7use std::sync::atomic::{AtomicUsize, Ordering};
8use std::sync::Arc;
9use std::task::{Context, Poll};
10use std::time::Duration;
11
12use tokio::io::{AsyncRead, AsyncWrite, ReadBuf};
13use tokio::net::TcpStream;
14use tokio::sync::OwnedSemaphorePermit;
15use tokio_rustls::rustls::pki_types::CertificateDer;
16use tokio_rustls::rustls::server::WebPkiClientVerifier;
17use tokio_rustls::rustls::{RootCertStore, ServerConfig};
18use tokio_rustls::TlsAcceptor;
19use tonic::transport::{
20    server::Connected, Certificate, Channel, ClientTlsConfig, Endpoint, Identity, Server,
21};
22
23use crate::native;
24
25#[derive(Debug, thiserror::Error)]
26pub enum NativeRpcTransportError {
27    #[error("native RPC I/O error: {0}")]
28    Io(#[from] std::io::Error),
29    #[error("native RPC transport error: {0}")]
30    Tonic(#[from] tonic::transport::Error),
31}
32
33#[derive(Debug, Clone)]
34pub struct NativeRpcServerConfig {
35    pub address: SocketAddr,
36    pub certificate_pem: Vec<u8>,
37    pub private_key_pem: Vec<u8>,
38    pub client_ca_pem: Option<Vec<u8>>,
39    pub max_connections: usize,
40    pub max_concurrent_streams: u32,
41    pub max_in_flight_per_connection: usize,
42    pub request_timeout: Duration,
43    pub idle_timeout: Duration,
44    pub keepalive_interval: Duration,
45    pub keepalive_timeout: Duration,
46}
47
48impl NativeRpcServerConfig {
49    fn tls13_acceptor(&self) -> Result<TlsAcceptor, Error> {
50        let _ = tokio_rustls::rustls::crypto::ring::default_provider().install_default();
51        let certificates =
52            rustls_pemfile::certs(&mut BufReader::new(self.certificate_pem.as_slice()))
53                .collect::<Result<Vec<CertificateDer<'static>>, _>>()?;
54        let private_key =
55            rustls_pemfile::private_key(&mut BufReader::new(self.private_key_pem.as_slice()))?
56                .ok_or_else(|| Error::new(ErrorKind::InvalidInput, "TLS private key is missing"))?;
57        let builder =
58            ServerConfig::builder_with_protocol_versions(&[&tokio_rustls::rustls::version::TLS13]);
59        let builder = match &self.client_ca_pem {
60            None => builder.with_no_client_auth(),
61            Some(client_ca) => {
62                let mut roots = RootCertStore::empty();
63                roots.add_parsable_certificates(
64                    rustls_pemfile::certs(&mut BufReader::new(client_ca.as_slice()))
65                        .collect::<Result<Vec<_>, _>>()?,
66                );
67                let verifier = WebPkiClientVerifier::builder(roots.into())
68                    .build()
69                    .map_err(|error| Error::new(ErrorKind::InvalidInput, error))?;
70                builder.with_client_cert_verifier(verifier)
71            }
72        };
73        let mut config = builder
74            .with_single_cert(certificates, private_key)
75            .map_err(|error| Error::new(ErrorKind::InvalidInput, error))?;
76        config.alpn_protocols = vec![b"h2".to_vec()];
77        Ok(TlsAcceptor::from(Arc::new(config)))
78    }
79}
80
81struct Tls13Connection {
82    stream: tokio_rustls::server::TlsStream<TcpStream>,
83    peer: SocketAddr,
84    idle_timeout: Duration,
85    idle: Pin<Box<tokio::time::Sleep>>,
86    _permit: OwnedSemaphorePermit,
87}
88
89impl Tls13Connection {
90    fn poll_idle(&mut self, cx: &mut Context<'_>) -> std::io::Result<()> {
91        if self.idle.as_mut().poll(cx).is_ready() {
92            Err(Error::new(
93                ErrorKind::TimedOut,
94                "native RPC connection idle timeout",
95            ))
96        } else {
97            Ok(())
98        }
99    }
100
101    fn reset_idle(&mut self) {
102        self.idle
103            .as_mut()
104            .reset(tokio::time::Instant::now() + self.idle_timeout);
105    }
106}
107
108impl Connected for Tls13Connection {
109    type ConnectInfo = SocketAddr;
110
111    fn connect_info(&self) -> Self::ConnectInfo {
112        self.peer
113    }
114}
115
116impl AsyncRead for Tls13Connection {
117    fn poll_read(
118        mut self: Pin<&mut Self>,
119        cx: &mut Context<'_>,
120        buffer: &mut ReadBuf<'_>,
121    ) -> Poll<std::io::Result<()>> {
122        if let Err(error) = self.poll_idle(cx) {
123            return Poll::Ready(Err(error));
124        }
125        let before = buffer.filled().len();
126        let result = Pin::new(&mut self.stream).poll_read(cx, buffer);
127        if matches!(&result, Poll::Ready(Ok(()))) && buffer.filled().len() > before {
128            self.reset_idle();
129        }
130        result
131    }
132}
133
134impl AsyncWrite for Tls13Connection {
135    fn poll_write(
136        mut self: Pin<&mut Self>,
137        cx: &mut Context<'_>,
138        buffer: &[u8],
139    ) -> Poll<std::io::Result<usize>> {
140        if let Err(error) = self.poll_idle(cx) {
141            return Poll::Ready(Err(error));
142        }
143        let result = Pin::new(&mut self.stream).poll_write(cx, buffer);
144        if matches!(&result, Poll::Ready(Ok(written)) if *written > 0) {
145            self.reset_idle();
146        }
147        result
148    }
149
150    fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
151        Pin::new(&mut self.stream).poll_flush(cx)
152    }
153
154    fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
155        Pin::new(&mut self.stream).poll_shutdown(cx)
156    }
157}
158
159pub struct NativeRpcServices<A, S, Q, T, C, D, H> {
160    pub auth: A,
161    pub session: S,
162    pub query: Q,
163    pub transaction: T,
164    pub catalog: C,
165    pub admin: D,
166    pub health: H,
167}
168
169/// Real native listener. All seven adapters are supplied from one canonical
170/// runtime so HTTP/Kit/native transports can share behavior.
171pub struct NativeRpcServer {
172    config: NativeRpcServerConfig,
173}
174
175impl NativeRpcServer {
176    pub fn new(config: NativeRpcServerConfig) -> Self {
177        Self { config }
178    }
179
180    pub async fn serve_with_shutdown<A, S, Q, T, C, D, H, F>(
181        self,
182        services: NativeRpcServices<A, S, Q, T, C, D, H>,
183        shutdown: F,
184    ) -> Result<(), NativeRpcTransportError>
185    where
186        A: native::auth_service_server::AuthService,
187        S: native::session_service_server::SessionService,
188        Q: native::query_service_server::QueryService,
189        T: native::transaction_service_server::TransactionService,
190        C: native::catalog_service_server::CatalogService,
191        D: native::admin_service_server::AdminService,
192        H: native::health_service_server::HealthService,
193        F: Future<Output = ()>,
194    {
195        let listener = tokio::net::TcpListener::bind(self.config.address).await?;
196        let acceptor = self.config.tls13_acceptor()?;
197        let idle_timeout = self.config.idle_timeout;
198        let permits = Arc::new(tokio::sync::Semaphore::new(
199            self.config.max_connections.max(1),
200        ));
201        let handshake_timeout = self.config.keepalive_timeout;
202        let (connections, incoming) =
203            tokio::sync::mpsc::channel(self.config.max_connections.max(1));
204        let accept_task = tokio::spawn(async move {
205            loop {
206                let (stream, peer) = listener.accept().await?;
207                let permit = match Arc::clone(&permits).try_acquire_owned() {
208                    Ok(permit) => permit,
209                    Err(_) => continue,
210                };
211                let acceptor = acceptor.clone();
212                let connections = connections.clone();
213                tokio::spawn(async move {
214                    let accepted = tokio::time::timeout(handshake_timeout, acceptor.accept(stream))
215                        .await
216                        .map_err(|_| Error::new(ErrorKind::TimedOut, "TLS handshake timed out"))
217                        .and_then(|result| result.map_err(Error::other))
218                        .map(|stream| Tls13Connection {
219                            stream,
220                            peer,
221                            idle_timeout,
222                            idle: Box::pin(tokio::time::sleep(idle_timeout)),
223                            _permit: permit,
224                        });
225                    let _ = connections.send(accepted).await;
226                });
227            }
228            #[allow(unreachable_code)]
229            Ok::<(), Error>(())
230        });
231        let result = Server::builder()
232            .max_concurrent_streams(Some(self.config.max_concurrent_streams.max(1)))
233            .concurrency_limit_per_connection(self.config.max_in_flight_per_connection.max(1))
234            .timeout(self.config.request_timeout)
235            .http2_keepalive_interval(Some(self.config.keepalive_interval))
236            .http2_keepalive_timeout(Some(self.config.keepalive_timeout))
237            .add_service(native::auth_service_server::AuthServiceServer::new(
238                services.auth,
239            ))
240            .add_service(native::session_service_server::SessionServiceServer::new(
241                services.session,
242            ))
243            .add_service(native::query_service_server::QueryServiceServer::new(
244                services.query,
245            ))
246            .add_service(
247                native::transaction_service_server::TransactionServiceServer::new(
248                    services.transaction,
249                ),
250            )
251            .add_service(native::catalog_service_server::CatalogServiceServer::new(
252                services.catalog,
253            ))
254            .add_service(native::admin_service_server::AdminServiceServer::new(
255                services.admin,
256            ))
257            .add_service(native::health_service_server::HealthServiceServer::new(
258                services.health,
259            ))
260            .serve_with_incoming_shutdown(
261                tokio_stream::wrappers::ReceiverStream::new(incoming),
262                shutdown,
263            )
264            .await;
265        accept_task.abort();
266        result.map_err(Into::into)
267    }
268}
269
270#[derive(Clone)]
271pub struct NativeRpcClientConfig {
272    pub endpoint: String,
273    pub domain_name: String,
274    pub ca_certificate_pem: Vec<u8>,
275    pub client_identity_pem: Option<(Vec<u8>, Vec<u8>)>,
276    pub connect_timeout: Duration,
277    pub request_timeout: Duration,
278    pub max_in_flight: usize,
279    pub tcp_keepalive: Duration,
280    pub http2_keepalive_interval: Duration,
281}
282
283#[derive(Clone)]
284pub struct NativeRpcConnection {
285    channel: Channel,
286}
287
288impl NativeRpcConnection {
289    pub async fn connect(config: &NativeRpcClientConfig) -> Result<Self, tonic::transport::Error> {
290        let mut tls = ClientTlsConfig::new()
291            .ca_certificate(Certificate::from_pem(config.ca_certificate_pem.clone()))
292            .domain_name(config.domain_name.clone());
293        if let Some((certificate, key)) = &config.client_identity_pem {
294            tls = tls.identity(Identity::from_pem(certificate.clone(), key.clone()));
295        }
296        let channel = Endpoint::from_shared(config.endpoint.clone())?
297            .tls_config(tls)?
298            .connect_timeout(config.connect_timeout)
299            .timeout(config.request_timeout)
300            .concurrency_limit(config.max_in_flight.max(1))
301            .tcp_keepalive(Some(config.tcp_keepalive))
302            .http2_keep_alive_interval(config.http2_keepalive_interval)
303            .keep_alive_while_idle(false)
304            .connect()
305            .await?;
306        Ok(Self { channel })
307    }
308
309    pub fn client(&self) -> NativeRpcClient {
310        NativeRpcClient::new(self.channel.clone())
311    }
312}
313
314/// Seven generated clients sharing one multiplexed HTTP/2 channel.
315#[derive(Clone)]
316pub struct NativeRpcClient {
317    channel: Channel,
318}
319
320impl NativeRpcClient {
321    pub fn new(channel: Channel) -> Self {
322        Self { channel }
323    }
324
325    pub fn auth(&self) -> native::auth_service_client::AuthServiceClient<Channel> {
326        native::auth_service_client::AuthServiceClient::new(self.channel.clone())
327    }
328
329    pub fn session(&self) -> native::session_service_client::SessionServiceClient<Channel> {
330        native::session_service_client::SessionServiceClient::new(self.channel.clone())
331    }
332
333    pub fn query(&self) -> native::query_service_client::QueryServiceClient<Channel> {
334        native::query_service_client::QueryServiceClient::new(self.channel.clone())
335    }
336
337    pub fn transaction(
338        &self,
339    ) -> native::transaction_service_client::TransactionServiceClient<Channel> {
340        native::transaction_service_client::TransactionServiceClient::new(self.channel.clone())
341    }
342
343    pub fn catalog(&self) -> native::catalog_service_client::CatalogServiceClient<Channel> {
344        native::catalog_service_client::CatalogServiceClient::new(self.channel.clone())
345    }
346
347    pub fn admin(&self) -> native::admin_service_client::AdminServiceClient<Channel> {
348        native::admin_service_client::AdminServiceClient::new(self.channel.clone())
349    }
350
351    pub fn health(&self) -> native::health_service_client::HealthServiceClient<Channel> {
352        native::health_service_client::HealthServiceClient::new(self.channel.clone())
353    }
354}
355
356/// Small round-robin connection pool. Each connection still multiplexes many
357/// concurrent streams.
358pub struct NativeRpcClientPool {
359    clients: Vec<NativeRpcClient>,
360    next: AtomicUsize,
361}
362
363impl NativeRpcClientPool {
364    pub async fn connect(
365        config: NativeRpcClientConfig,
366        connections: usize,
367    ) -> Result<Arc<Self>, tonic::transport::Error> {
368        let mut clients = Vec::with_capacity(connections.max(1));
369        for _ in 0..connections.max(1) {
370            clients.push(NativeRpcConnection::connect(&config).await?.client());
371        }
372        Ok(Arc::new(Self {
373            clients,
374            next: AtomicUsize::new(0),
375        }))
376    }
377
378    pub fn client(&self) -> NativeRpcClient {
379        let index = self.next.fetch_add(1, Ordering::Relaxed) % self.clients.len();
380        self.clients[index].clone()
381    }
382}
mongreldb_protocol/native_transport.rs

mongreldb_protocol/
native_transport.rs
