1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165
// Copyright Materialize, Inc. and contributors. All rights reserved.
//
// Use of this software is governed by the Business Source License
// included in the LICENSE file.
//
// As of the Change Date specified in that file, in accordance with
// the Business Source License, use of this software will be governed
// by the Apache License, Version 2.0.
use std::sync::Arc;
use async_trait::async_trait;
use futures::stream::{Stream, StreamExt};
use tokio::io::{self, AsyncWriteExt};
use tokio::net::TcpStream;
use tracing::{debug, error};
use mz_ore::netio::{self, SniffedStream, SniffingStream};
use mz_ore::task;
use crate::http;
type Handlers = Vec<Box<dyn ConnectionHandler + Send + Sync>>;
/// A mux routes incoming TCP connections to a dynamic set of connection
/// handlers. It enables serving multiple protocols over the same port.
///
/// Connections are routed by sniffing the first several bytes sent over the
/// wire and matching them against each handler, in order. The first handler
/// to match the connection will be invoked.
pub struct Mux {
handlers: Handlers,
}
impl Mux {
/// Constructs a new `Mux`.
pub fn new() -> Mux {
Mux { handlers: vec![] }
}
/// Adds a new connection handler to this mux.
pub fn add_handler<H>(&mut self, handler: H)
where
H: ConnectionHandler + Send + Sync + 'static,
{
self.handlers.push(Box::new(handler));
}
/// Serves incoming TCP traffic from `listener`.
pub async fn serve<S>(self, mut incoming: S)
where
S: Stream<Item = io::Result<TcpStream>> + Unpin,
{
let handlers = Arc::new(self.handlers);
while let Some(conn) = incoming.next().await {
let conn = match conn {
Ok(conn) => conn,
Err(err) => {
error!("error accepting connection: {}", err);
continue;
}
};
// Set TCP_NODELAY to disable tinygram prevention (Nagle's
// algorithm), which forces a 40ms delay between each query
// on linux. According to John Nagle [0], the true problem
// is delayed acks, but disabling those is a receive-side
// operation (TCP_QUICKACK), and we can't always control the
// client. PostgreSQL sets TCP_NODELAY on both sides of its
// sockets, so it seems sane to just do the same.
//
// If set_nodelay fails, it's a programming error, so panic.
//
// [0]: https://news.ycombinator.com/item?id=10608356
conn.set_nodelay(true).expect("set_nodelay failed");
task::spawn(
|| "mux_serve",
handle_connection(Arc::clone(&handlers), conn),
);
}
}
}
async fn handle_connection(handlers: Arc<Handlers>, conn: TcpStream) {
// Sniff out what protocol we've received. Choosing how many bytes to
// sniff is a delicate business. Read too many bytes and you'll stall
// out protocols with small handshakes, like pgwire. Read too few bytes
// and you won't be able to tell what protocol you have. For now, eight
// bytes is the magic number, but this may need to change if we learn to
// speak new protocols.
let mut ss = SniffingStream::new(conn);
let mut buf = [0; 8];
let nread = match netio::read_exact_or_eof(&mut ss, &mut buf).await {
Ok(nread) => nread,
Err(err) => {
error!("error handling request: {}", err);
return;
}
};
let buf = &buf[..nread];
for handler in &*handlers {
if handler.match_handshake(buf) {
if let Err(e) = handler.handle_connection(ss.into_sniffed()).await {
error!("error handling connection in {}: {:#}", handler.name(), e);
}
return;
}
}
debug!(
"dropped connection using unknown protocol (sniffed: 0x{})",
hex::encode(buf)
);
let _ = ss.into_sniffed().write_all(b"unknown protocol\n").await;
}
/// A connection handler manages an incoming network connection.
#[async_trait]
pub trait ConnectionHandler {
/// Returns the name of the connection handler for use in e.g. log messages.
fn name(&self) -> &str;
/// Determines whether this handler can accept the connection based on the
/// first several bytes in the stream.
fn match_handshake(&self, buf: &[u8]) -> bool;
/// Handles the connection.
async fn handle_connection(&self, conn: SniffedStream<TcpStream>) -> Result<(), anyhow::Error>;
}
#[async_trait]
impl ConnectionHandler for mz_pgwire::Server {
fn name(&self) -> &str {
"pgwire server"
}
fn match_handshake(&self, buf: &[u8]) -> bool {
mz_pgwire::match_handshake(buf)
}
async fn handle_connection(&self, conn: SniffedStream<TcpStream>) -> Result<(), anyhow::Error> {
// Using fully-qualified syntax means we won't accidentally call
// ourselves (i.e., silently infinitely recurse) if the name or type of
// `pgwire::Server::handle_connection` changes.
mz_pgwire::Server::handle_connection(self, conn).await
}
}
#[async_trait]
impl ConnectionHandler for http::Server {
fn name(&self) -> &str {
"http server"
}
fn match_handshake(&self, buf: &[u8]) -> bool {
self.match_handshake(buf)
}
async fn handle_connection(&self, conn: SniffedStream<TcpStream>) -> Result<(), anyhow::Error> {
// Using fully-qualified syntax means we won't accidentally call
// ourselves (i.e., silently infinitely recurse) if the name or type of
// `http::Server::handle_connection` changes.
http::Server::handle_connection(self, conn).await
}
}