use std::{collections::VecDeque, fmt, sync::Arc, time::Duration};
use tokio::{
io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt},
net::TcpStream,
sync::{Mutex, mpsc, oneshot},
time,
};
use tokio_rustls::{TlsConnector, rustls::pki_types::ServerName};
use crate::native::{
KafkaClientError, KafkaClientResult,
protocol::{
API_KEY_SASL_AUTHENTICATE, API_KEY_SASL_HANDSHAKE, Decoder, Encoder,
SaslAuthenticateResponse, SaslHandshakeResponse, encode_sasl_authenticate_request,
encode_sasl_handshake_request, request_header_version, response_header_version,
},
security::{NativeSecurityConfig, SaslConfig, SaslMechanism, ScramClient},
};
const DEFAULT_COMMAND_BUFFER: usize = 64;
const MAX_RESPONSE_BYTES: usize = 512 * 1024 * 1024;
#[derive(Debug, Clone)]
pub(crate) struct BrokerConnection {
sender: mpsc::Sender<Command>,
broker: String,
sasl: Option<SaslConfig>,
auth: Arc<Mutex<AuthState>>,
}
trait BrokerIo: AsyncRead + AsyncWrite + Unpin + Send {}
impl<T> BrokerIo for T where T: AsyncRead + AsyncWrite + Unpin + Send {}
#[derive(Debug, Default)]
struct AuthState {
handshake_version: Option<i16>,
authenticate_version: Option<i16>,
authenticated: bool,
reauthenticate_at: Option<time::Instant>,
}
#[derive(Debug)]
pub(crate) struct BrokerResponse {
receiver: oneshot::Receiver<KafkaClientResult<Vec<u8>>>,
}
#[derive(Debug)]
struct Command {
api_key: i16,
api_version: i16,
request_header_version: i16,
response_header_version: i16,
body: Vec<u8>,
expect_response: bool,
respond: oneshot::Sender<KafkaClientResult<Vec<u8>>>,
}
#[derive(Debug)]
struct Pending {
correlation_id: i32,
response_header_version: i16,
respond: oneshot::Sender<KafkaClientResult<Vec<u8>>>,
}
impl BrokerConnection {
pub(crate) async fn connect(
addr: String,
broker: String,
client_id: String,
max_in_flight: usize,
request_timeout: Duration,
security: &NativeSecurityConfig,
) -> KafkaClientResult<Self> {
let stream = time::timeout(request_timeout, TcpStream::connect(&addr))
.await
.map_err(|_| {
KafkaClientError::Io(std::io::Error::new(
std::io::ErrorKind::TimedOut,
format!("Kafka TCP connect to {addr} timed out"),
))
})??;
stream.set_nodelay(true)?;
let stream: Box<dyn BrokerIo> = if let Some(tls) = &security.tls {
let server_name =
ServerName::try_from(broker.clone()).map_err(|error| KafkaClientError::Tls {
broker: broker.clone(),
message: format!("invalid TLS SNI name: {error}"),
})?;
let connector = TlsConnector::from(Arc::clone(tls));
let tls_stream = time::timeout(request_timeout, connector.connect(server_name, stream))
.await
.map_err(|_| KafkaClientError::Tls {
broker: broker.clone(),
message: "TLS handshake timed out".to_owned(),
})?
.map_err(|error| KafkaClientError::Tls {
broker: broker.clone(),
message: error.to_string(),
})?;
Box::new(tls_stream)
} else {
Box::new(stream)
};
let (sender, receiver) = mpsc::channel(DEFAULT_COMMAND_BUFFER);
tokio::spawn(connection_task(
stream,
receiver,
client_id,
max_in_flight.max(1),
request_timeout,
));
Ok(Self {
sender,
broker,
sasl: security.sasl.clone(),
auth: Arc::new(Mutex::new(AuthState::default())),
})
}
pub(crate) async fn initialize_sasl(
&self,
handshake_version: i16,
authenticate_version: i16,
) -> KafkaClientResult<()> {
let Some(sasl) = self.sasl.clone() else {
return Ok(());
};
let mut auth = self.auth.lock().await;
auth.handshake_version = Some(handshake_version);
auth.authenticate_version = Some(authenticate_version);
self.authenticate(&sasl, &mut auth).await
}
pub(crate) async fn request_before_auth(
&self,
api_key: i16,
api_version: i16,
request_header_version: i16,
response_header_version: i16,
body: Vec<u8>,
) -> KafkaClientResult<Vec<u8>> {
self.raw_request(
api_key,
api_version,
request_header_version,
response_header_version,
body,
)
.await
}
pub(crate) async fn request(
&self,
api_key: i16,
api_version: i16,
request_header_version: i16,
response_header_version: i16,
body: Vec<u8>,
) -> KafkaClientResult<Vec<u8>> {
self.ensure_authenticated().await?;
self.raw_request(
api_key,
api_version,
request_header_version,
response_header_version,
body,
)
.await
}
pub(crate) async fn begin_request(
&self,
api_key: i16,
api_version: i16,
request_header_version: i16,
response_header_version: i16,
body: Vec<u8>,
) -> KafkaClientResult<BrokerResponse> {
self.ensure_authenticated().await?;
self.raw_begin_request(
api_key,
api_version,
request_header_version,
response_header_version,
body,
)
.await
}
async fn raw_begin_request(
&self,
api_key: i16,
api_version: i16,
request_header_version: i16,
response_header_version: i16,
body: Vec<u8>,
) -> KafkaClientResult<BrokerResponse> {
let receiver = self
.begin(
api_key,
api_version,
request_header_version,
response_header_version,
body,
true,
)
.await?;
Ok(BrokerResponse { receiver })
}
pub(crate) async fn send_one_way(
&self,
api_key: i16,
api_version: i16,
request_header_version: i16,
body: Vec<u8>,
) -> KafkaClientResult<()> {
self.ensure_authenticated().await?;
self.begin(api_key, api_version, request_header_version, 0, body, false)
.await?
.await
.map_err(|_| KafkaClientError::ChannelClosed)?
.map(|_| ())
}
async fn raw_request(
&self,
api_key: i16,
api_version: i16,
request_header_version: i16,
response_header_version: i16,
body: Vec<u8>,
) -> KafkaClientResult<Vec<u8>> {
self.raw_begin_request(
api_key,
api_version,
request_header_version,
response_header_version,
body,
)
.await?
.receive()
.await
}
async fn ensure_authenticated(&self) -> KafkaClientResult<()> {
let Some(sasl) = self.sasl.clone() else {
return Ok(());
};
let mut auth = self.auth.lock().await;
if auth.authenticated
&& auth
.reauthenticate_at
.is_none_or(|deadline| deadline > time::Instant::now())
{
return Ok(());
}
self.authenticate(&sasl, &mut auth).await
}
async fn authenticate(&self, sasl: &SaslConfig, auth: &mut AuthState) -> KafkaClientResult<()> {
let handshake_version = auth.handshake_version.ok_or_else(|| {
self.sasl_error(sasl.mechanism, "SASL API versions were not negotiated")
})?;
let authenticate_version = auth.authenticate_version.ok_or_else(|| {
self.sasl_error(sasl.mechanism, "SASL API versions were not negotiated")
})?;
let body = encode_sasl_handshake_request(sasl.mechanism.name())
.map_err(|error| self.sasl_error(sasl.mechanism, error))?;
let response = self
.raw_request(
API_KEY_SASL_HANDSHAKE,
handshake_version,
request_header_version(handshake_version),
response_header_version(API_KEY_SASL_HANDSHAKE, handshake_version),
body,
)
.await
.map_err(|error| self.sasl_error(sasl.mechanism, error))?;
let handshake = SaslHandshakeResponse::decode(handshake_version, &response)
.map_err(|error| self.sasl_error(sasl.mechanism, error))?;
if handshake.error_code != 0 {
return Err(self.sasl_error(
sasl.mechanism,
format!(
"SaslHandshake returned broker code {} (advertised: {})",
handshake.error_code,
handshake.mechanisms.join(", ")
),
));
}
if !handshake
.mechanisms
.iter()
.any(|mechanism| mechanism == sasl.mechanism.name())
{
return Err(self.sasl_error(
sasl.mechanism,
format!(
"broker did not advertise the selected mechanism (advertised: {})",
handshake.mechanisms.join(", ")
),
));
}
let session_lifetime_ms = match sasl.mechanism {
SaslMechanism::Plain => {
if sasl.username.contains('\0') || sasl.password.contains('\0') {
return Err(self.sasl_error(
sasl.mechanism,
"PLAIN credentials must not contain NUL bytes",
));
}
let mut token = Vec::with_capacity(sasl.username.len() + sasl.password.len() + 2);
token.push(0);
token.extend_from_slice(sasl.username.as_bytes());
token.push(0);
token.extend_from_slice(sasl.password.as_bytes());
self.authenticate_token(sasl.mechanism, authenticate_version, &token)
.await?
.session_lifetime_ms
}
SaslMechanism::ScramSha256 | SaslMechanism::ScramSha512 => {
let client = ScramClient::new(sasl)
.map_err(|error| self.sasl_error(sasl.mechanism, error))?;
let first = self
.authenticate_token(
sasl.mechanism,
authenticate_version,
client.first_message().as_bytes(),
)
.await?;
let server_first = std::str::from_utf8(&first.auth_bytes).map_err(|_| {
self.sasl_error(sasl.mechanism, "SCRAM server-first message was not UTF-8")
})?;
let (client_final, verifier) = client
.handle_server_first(server_first)
.map_err(|error| self.sasl_error(sasl.mechanism, error))?;
let final_response = self
.authenticate_token(
sasl.mechanism,
authenticate_version,
client_final.as_bytes(),
)
.await?;
let server_final =
std::str::from_utf8(&final_response.auth_bytes).map_err(|_| {
self.sasl_error(sasl.mechanism, "SCRAM server-final message was not UTF-8")
})?;
verifier
.verify(server_final)
.map_err(|error| self.sasl_error(sasl.mechanism, error))?;
final_response.session_lifetime_ms
}
};
auth.authenticated = true;
auth.reauthenticate_at = reauthentication_deadline(session_lifetime_ms);
Ok(())
}
async fn authenticate_token(
&self,
mechanism: SaslMechanism,
version: i16,
token: &[u8],
) -> KafkaClientResult<SaslAuthenticateResponse> {
let body = encode_sasl_authenticate_request(token)
.map_err(|error| self.sasl_error(mechanism, error))?;
let response = self
.raw_request(
API_KEY_SASL_AUTHENTICATE,
version,
request_header_version(version),
response_header_version(API_KEY_SASL_AUTHENTICATE, version),
body,
)
.await
.map_err(|error| self.sasl_error(mechanism, error))?;
let response = SaslAuthenticateResponse::decode(version, &response)
.map_err(|error| self.sasl_error(mechanism, error))?;
if response.error_code != 0 {
return Err(self.sasl_error(
mechanism,
response.error_message.as_deref().map_or_else(
|| {
format!(
"SaslAuthenticate returned broker code {}",
response.error_code
)
},
|message| {
format!(
"SaslAuthenticate returned broker code {}: {message}",
response.error_code
)
},
),
));
}
Ok(response)
}
fn sasl_error(&self, mechanism: SaslMechanism, message: impl fmt::Display) -> KafkaClientError {
KafkaClientError::Sasl {
broker: self.broker.clone(),
mechanism: mechanism.name().to_owned(),
message: message.to_string(),
}
}
async fn begin(
&self,
api_key: i16,
api_version: i16,
request_header_version: i16,
response_header_version: i16,
body: Vec<u8>,
expect_response: bool,
) -> KafkaClientResult<oneshot::Receiver<KafkaClientResult<Vec<u8>>>> {
let (respond, receiver) = oneshot::channel();
let command = Command {
api_key,
api_version,
request_header_version,
response_header_version,
body,
expect_response,
respond,
};
self.sender
.send(command)
.await
.map_err(|_| KafkaClientError::ChannelClosed)?;
Ok(receiver)
}
}
fn reauthentication_deadline(session_lifetime_ms: i64) -> Option<time::Instant> {
let lifetime_ms = u64::try_from(session_lifetime_ms)
.ok()
.filter(|value| *value > 0)?;
let margin_ms = (lifetime_ms / 10).clamp(1, 30_000);
Some(time::Instant::now() + Duration::from_millis(lifetime_ms.saturating_sub(margin_ms)))
}
impl BrokerResponse {
pub(crate) async fn receive(self) -> KafkaClientResult<Vec<u8>> {
self.receiver
.await
.map_err(|_| KafkaClientError::ChannelClosed)?
}
}
async fn connection_task(
mut stream: Box<dyn BrokerIo>,
mut receiver: mpsc::Receiver<Command>,
client_id: String,
max_in_flight: usize,
request_timeout: Duration,
) {
let mut correlation_id = 1_i32;
let mut pending = VecDeque::<Pending>::new();
let mut receiving_commands = true;
loop {
if !receiving_commands {
if pending.is_empty() {
break;
}
match read_response(&mut stream, &mut pending, request_timeout).await {
Ok(()) => continue,
Err(error) => {
fail_all(&mut pending, error);
break;
}
}
}
if pending.len() >= max_in_flight {
match read_response(&mut stream, &mut pending, request_timeout).await {
Ok(()) => continue,
Err(error) => {
fail_all(&mut pending, error);
break;
}
}
}
tokio::select! {
command = receiver.recv(), if receiving_commands => {
let Some(command) = command else {
receiving_commands = false;
continue;
};
let request_correlation = correlation_id;
correlation_id = if correlation_id == i32::MAX { 1 } else { correlation_id + 1 };
match write_request(&mut stream, &client_id, request_correlation, &command).await {
Ok(()) => {
if command.expect_response {
pending.push_back(Pending {
correlation_id: request_correlation,
response_header_version: command.response_header_version,
respond: command.respond,
});
} else {
let _ = command.respond.send(Ok(Vec::new()));
}
}
Err(error) => {
let _ = command.respond.send(Err(error));
fail_all(&mut pending, KafkaClientError::ChannelClosed);
break;
}
}
}
result = read_response(&mut stream, &mut pending, request_timeout), if !pending.is_empty() => {
if let Err(error) = result {
fail_all(&mut pending, error);
break;
}
}
}
}
}
async fn write_request<S>(
stream: &mut S,
client_id: &str,
correlation_id: i32,
command: &Command,
) -> KafkaClientResult<()>
where
S: AsyncWrite + Unpin + ?Sized,
{
let mut frame = Encoder::with_capacity(command.body.len() + 64);
frame.put_i16(command.api_key);
frame.put_i16(command.api_version);
frame.put_i32(correlation_id);
frame.put_nullable_string(Some(client_id))?;
if command.request_header_version >= 2 {
frame.put_empty_tags();
}
frame.put_raw(&command.body);
let frame = frame.into_inner();
let len = i32::try_from(frame.len())
.map_err(|_| KafkaClientError::protocol("Kafka request frame exceeds int32 length"))?;
stream.write_all(&len.to_be_bytes()).await?;
stream.write_all(&frame).await?;
Ok(())
}
async fn read_response<S>(
stream: &mut S,
pending: &mut VecDeque<Pending>,
request_timeout: Duration,
) -> KafkaClientResult<()>
where
S: AsyncRead + Unpin + ?Sized,
{
let mut len_buf = [0_u8; 4];
time::timeout(request_timeout, stream.read_exact(&mut len_buf))
.await
.map_err(|_| KafkaClientError::protocol("Kafka response read timed out"))??;
let len = i32::from_be_bytes(len_buf);
if len < 4 {
return Err(KafkaClientError::protocol(format!(
"invalid Kafka response frame length {len}"
)));
}
let len = usize::try_from(len)
.map_err(|_| KafkaClientError::protocol("Kafka response length conversion failed"))?;
if len > MAX_RESPONSE_BYTES {
return Err(KafkaClientError::protocol(format!(
"Kafka response frame too large: {len} bytes"
)));
}
let mut frame = vec![0_u8; len];
time::timeout(request_timeout, stream.read_exact(&mut frame))
.await
.map_err(|_| KafkaClientError::protocol("Kafka response body read timed out"))??;
let Some(pending_request) = pending.pop_front() else {
return Err(KafkaClientError::protocol(
"Kafka response arrived with no pending request",
));
};
let mut decoder = Decoder::new(&frame);
let correlation_id = decoder.get_i32()?;
if pending_request.response_header_version >= 1 {
decoder.skip_tags()?;
}
if correlation_id != pending_request.correlation_id {
let _ = pending_request
.respond
.send(Err(KafkaClientError::protocol(format!(
"Kafka response correlation mismatch: expected {}, got {correlation_id}",
pending_request.correlation_id
))));
return Err(KafkaClientError::protocol(
"Kafka connection correlation-id mismatch",
));
}
let body = decoder.take(decoder.remaining())?.to_vec();
let _ = pending_request.respond.send(Ok(body));
Ok(())
}
fn fail_all(pending: &mut VecDeque<Pending>, error: KafkaClientError) {
let retriable = error.retriable();
let message = error.to_string();
while let Some(pending) = pending.pop_front() {
let error = if retriable {
KafkaClientError::ChannelClosed
} else {
KafkaClientError::protocol(message.clone())
};
let _ = pending.respond.send(Err(error));
}
}