use bytes::{Buf, BufMut, Bytes, BytesMut};
use crate::error::{ErrorCode, KrafkaError, ProtocolErrorKind, Result};
use crate::protocol::{
Decode, Encode, KafkaBytes, KafkaString, TaggedField, TaggedFields, TryEncode,
};
const MAX_ARRAY_LEN: usize = 100_000;
pub(crate) fn read_string(buf: &mut impl Buf) -> Result<String> {
KafkaString::decode(buf)?.0.ok_or_else(|| {
KrafkaError::protocol_kind(ProtocolErrorKind::Malformed, "unexpected null string")
})
}
pub(crate) fn read_nullable_string(buf: &mut impl Buf) -> Result<Option<String>> {
Ok(KafkaString::decode(buf)?.0)
}
pub(crate) fn read_nullable_bytes(buf: &mut impl Buf) -> Result<Option<Bytes>> {
Ok(KafkaBytes::decode(buf)?.0)
}
pub(crate) fn read_bytes(buf: &mut impl Buf) -> Result<Bytes> {
read_nullable_bytes(buf)?.ok_or_else(|| {
KrafkaError::protocol_kind(ProtocolErrorKind::Malformed, "unexpected null bytes")
})
}
pub(crate) fn read_array_len(buf: &mut impl Buf) -> Result<usize> {
let len = i32::decode(buf)?;
if len < 0 {
return Err(KrafkaError::protocol_kind(
ProtocolErrorKind::Malformed,
format!("null array length {len} where a non-nullable array was expected"),
));
}
check_len(len as usize)
}
pub(crate) fn read_nullable_array_len(buf: &mut impl Buf) -> Result<Option<usize>> {
let len = i32::decode(buf)?;
if len < 0 {
return Ok(None);
}
Ok(Some(check_len(len as usize)?))
}
fn check_len(len: usize) -> Result<usize> {
if len > MAX_ARRAY_LEN {
return Err(KrafkaError::protocol_kind(
ProtocolErrorKind::InvalidLength,
format!("array length {len} exceeds the fake broker's safety limit {MAX_ARRAY_LEN}"),
));
}
Ok(len)
}
pub(crate) fn write_string(buf: &mut impl BufMut, value: &str) -> Result<()> {
KafkaString::new(value).try_encode(buf)
}
pub(crate) fn write_nullable_string(buf: &mut impl BufMut, value: Option<&str>) -> Result<()> {
match value {
Some(v) => KafkaString::new(v).try_encode(buf),
None => KafkaString::null().try_encode(buf),
}
}
pub(crate) fn write_nullable_bytes(buf: &mut impl BufMut, value: Option<&Bytes>) -> Result<()> {
match value {
Some(v) => KafkaBytes::new(v.clone()).try_encode(buf),
None => KafkaBytes::null().try_encode(buf),
}
}
pub(crate) fn write_array_len(buf: &mut impl BufMut, len: usize) -> Result<()> {
let len = i32::try_from(len).map_err(|_| {
KrafkaError::protocol_kind(ProtocolErrorKind::InvalidLength, "array length exceeds i32")
})?;
buf.put_i32(len);
Ok(())
}
pub(crate) fn write_error(buf: &mut impl BufMut, code: ErrorCode) {
code.to_i16().encode(buf);
}
pub(crate) fn read_compact_array_len(buf: &mut impl Buf) -> Result<usize> {
let raw = crate::util::varint::decode_unsigned_varint(buf)?;
if raw == 0 {
return Ok(0); }
check_len((raw - 1) as usize)
}
pub(crate) fn write_compact_array_len(buf: &mut impl BufMut, len: usize) -> Result<()> {
let len = u32::try_from(len.saturating_add(1)).map_err(|_| {
KrafkaError::protocol_kind(ProtocolErrorKind::InvalidLength, "array length exceeds u32")
})?;
crate::util::varint::encode_unsigned_varint(len, buf);
Ok(())
}
pub(crate) fn read_compact_string(buf: &mut impl Buf) -> Result<String> {
KafkaString::decode_compact(buf)?.0.ok_or_else(|| {
KrafkaError::protocol_kind(
ProtocolErrorKind::Malformed,
"unexpected null compact string",
)
})
}
pub(crate) fn write_compact_string(buf: &mut impl BufMut, value: &str) -> Result<()> {
KafkaString::new(value).try_encode_compact(buf)
}
pub(crate) fn write_compact_nullable_string(
buf: &mut impl BufMut,
value: Option<&str>,
) -> Result<()> {
match value {
Some(v) => KafkaString::new(v).try_encode_compact(buf),
None => KafkaString::null().try_encode_compact(buf),
}
}
pub(crate) fn read_compact_nullable_bytes(buf: &mut impl Buf) -> Result<Option<Bytes>> {
Ok(KafkaBytes::decode_compact(buf)?.0)
}
pub(crate) fn write_compact_nullable_bytes(
buf: &mut impl BufMut,
value: Option<&Bytes>,
) -> Result<()> {
match value {
Some(v) => KafkaBytes(Some(v.clone())).try_encode_compact(buf),
None => KafkaBytes(None).try_encode_compact(buf),
}
}
pub(crate) fn write_empty_tagged_fields(buf: &mut impl BufMut) -> Result<()> {
TaggedFields::default().try_encode(buf)
}
pub(crate) fn write_tagged_fields(buf: &mut impl BufMut, fields: Vec<TaggedField>) -> Result<()> {
TaggedFields(fields).try_encode(buf)
}
pub(crate) fn skip_tagged_fields(buf: &mut impl Buf) -> Result<()> {
TaggedFields::decode(buf)?;
Ok(())
}
pub(crate) fn current_leader_field(leader_id: i32, leader_epoch: i32) -> TaggedField {
let mut data = BytesMut::new();
data.put_i32(leader_id);
data.put_i32(leader_epoch);
data.put_u8(0); TaggedField {
tag: 0,
data: data.freeze(),
}
}
pub(crate) fn node_endpoints_field(endpoints: &[(i32, &str, i32)]) -> Result<TaggedField> {
let mut data = BytesMut::new();
write_compact_array_len(&mut data, endpoints.len())?;
for (node_id, host, port) in endpoints {
data.put_i32(*node_id);
KafkaString::new(*host).try_encode_compact(&mut data)?;
data.put_i32(*port);
KafkaString::null().try_encode_compact(&mut data)?; write_empty_tagged_fields(&mut data)?;
}
Ok(TaggedField {
tag: 0,
data: data.freeze(),
})
}
#[derive(Debug, Clone)]
pub(crate) struct MetadataReq {
pub topics: Option<Vec<String>>,
pub allow_auto_topic_creation: bool,
}
impl MetadataReq {
pub(crate) fn read_v12(buf: &mut impl Buf) -> Result<Self> {
let topics = match read_compact_nullable_array_len(buf)? {
None => None,
Some(count) => {
let mut names = Vec::with_capacity(count);
for _ in 0..count {
if buf.remaining() < 16 {
return Err(crate::error::KrafkaError::protocol_kind(
crate::error::ProtocolErrorKind::TruncatedFrame,
"fake broker: truncated topic id in Metadata v12 request",
));
}
let mut topic_id = [0u8; 16];
buf.copy_to_slice(&mut topic_id);
let name = read_compact_nullable_string(buf)?.ok_or_else(|| {
crate::error::KrafkaError::protocol_kind(
crate::error::ProtocolErrorKind::InvalidValue,
"fake broker: Metadata v12 topic lookup by UUID is not modelled; \
krafka always requests by name",
)
})?;
skip_tagged_fields(buf)?;
names.push(name);
}
Some(names)
}
};
let allow_auto_topic_creation = bool::decode(buf)?;
let _ = bool::decode(buf)?;
skip_tagged_fields(buf)?;
Ok(Self {
topics,
allow_auto_topic_creation,
})
}
}
#[derive(Debug, Clone)]
pub(crate) struct ProduceReqPartition {
pub index: i32,
pub records: Option<Bytes>,
}
#[derive(Debug, Clone)]
pub(crate) struct ProduceReqTopic {
pub name: String,
pub partitions: Vec<ProduceReqPartition>,
}
#[derive(Debug, Clone)]
pub(crate) struct ProduceReq {
pub transactional_id: Option<String>,
pub topics: Vec<ProduceReqTopic>,
}
impl ProduceReq {
pub(crate) fn read(buf: &mut impl Buf) -> Result<Self> {
let transactional_id = KafkaString::decode_compact(buf)?
.0
.filter(|s| !s.is_empty());
let _acks = i16::decode(buf)?;
let _timeout_ms = i32::decode(buf)?;
let topic_count = read_compact_array_len(buf)?;
let mut topics = Vec::with_capacity(topic_count);
for _ in 0..topic_count {
let name = read_compact_string(buf)?;
let partition_count = read_compact_array_len(buf)?;
let mut partitions = Vec::with_capacity(partition_count);
for _ in 0..partition_count {
let index = i32::decode(buf)?;
let records = read_compact_nullable_bytes(buf)?;
skip_tagged_fields(buf)?;
partitions.push(ProduceReqPartition { index, records });
}
skip_tagged_fields(buf)?;
topics.push(ProduceReqTopic { name, partitions });
}
skip_tagged_fields(buf)?;
Ok(Self {
transactional_id,
topics,
})
}
}
#[derive(Debug, Clone)]
pub(crate) struct FetchReqPartition {
pub partition: i32,
pub current_leader_epoch: i32,
pub fetch_offset: i64,
}
#[derive(Debug, Clone)]
pub(crate) struct FetchReqTopic {
pub topic: String,
pub partitions: Vec<FetchReqPartition>,
}
#[derive(Debug, Clone)]
pub(crate) struct FetchReq {
pub session_id: i32,
pub isolation_level: i8,
pub topics: Vec<FetchReqTopic>,
}
impl FetchReq {
pub(crate) fn read(buf: &mut impl Buf) -> Result<Self> {
let _replica_id = i32::decode(buf)?;
let _max_wait_ms = i32::decode(buf)?;
let _min_bytes = i32::decode(buf)?;
let _max_bytes = i32::decode(buf)?;
let isolation_level = i8::decode(buf)?;
let session_id = i32::decode(buf)?;
let _session_epoch = i32::decode(buf)?;
let topic_count = read_array_len(buf)?;
let mut topics = Vec::with_capacity(topic_count);
for _ in 0..topic_count {
let topic = read_string(buf)?;
let partition_count = read_array_len(buf)?;
let mut partitions = Vec::with_capacity(partition_count);
for _ in 0..partition_count {
let partition = i32::decode(buf)?;
let current_leader_epoch = i32::decode(buf)?;
let fetch_offset = i64::decode(buf)?;
let _log_start_offset = i64::decode(buf)?;
let _partition_max_bytes = i32::decode(buf)?;
partitions.push(FetchReqPartition {
partition,
current_leader_epoch,
fetch_offset,
});
}
topics.push(FetchReqTopic { topic, partitions });
}
let forgotten_count = read_array_len(buf)?;
for _ in 0..forgotten_count {
let _ = read_string(buf)?;
let partition_count = read_array_len(buf)?;
for _ in 0..partition_count {
let _ = i32::decode(buf)?;
}
}
let _ = read_nullable_string(buf)?;
Ok(Self {
session_id,
isolation_level,
topics,
})
}
}
pub(crate) fn read_compact_nullable_string(buf: &mut impl Buf) -> Result<Option<String>> {
Ok(KafkaString::decode_compact(buf)?.0)
}
pub(crate) fn read_compact_nullable_array_len(buf: &mut impl Buf) -> Result<Option<usize>> {
let raw = crate::util::varint::decode_unsigned_varint(buf)?;
if raw == 0 {
return Ok(None);
}
Ok(Some(crate::protocol::check_compact_array_len(raw)?))
}
#[derive(Debug, Clone)]
pub(crate) struct HeartbeatTopicPartitions {
pub topic_id: [u8; 16],
pub partitions: Vec<i32>,
}
#[derive(Debug, Clone)]
pub(crate) struct ConsumerGroupHeartbeatReq {
pub group_id: String,
pub member_id: String,
pub member_epoch: i32,
pub instance_id: Option<String>,
#[allow(dead_code)]
pub rack_id: Option<String>,
#[allow(dead_code)]
pub rebalance_timeout_ms: i32,
pub subscribed_topic_names: Option<Vec<String>>,
#[allow(dead_code)]
pub subscribed_topic_regex: Option<String>,
#[allow(dead_code)]
pub server_assignor: Option<String>,
#[allow(dead_code)]
pub topic_partitions: Option<Vec<HeartbeatTopicPartitions>>,
}
impl ConsumerGroupHeartbeatReq {
pub(crate) fn read(buf: &mut impl Buf) -> Result<Self> {
let group_id = read_compact_string(buf)?;
let member_id = read_compact_string(buf)?;
let member_epoch = i32::decode(buf)?;
let instance_id = read_compact_nullable_string(buf)?;
let rack_id = read_compact_nullable_string(buf)?;
let rebalance_timeout_ms = i32::decode(buf)?;
let subscribed_topic_names = match read_compact_nullable_array_len(buf)? {
None => None,
Some(count) => {
let mut names = Vec::with_capacity(count);
for _ in 0..count {
names.push(read_compact_string(buf)?);
}
Some(names)
}
};
let subscribed_topic_regex = read_compact_nullable_string(buf)?;
let server_assignor = read_compact_nullable_string(buf)?;
let topic_partitions = match read_compact_nullable_array_len(buf)? {
None => None,
Some(count) => {
let mut entries = Vec::with_capacity(count);
for _ in 0..count {
let mut topic_id = [0u8; 16];
if buf.remaining() < 16 {
return Err(crate::error::KrafkaError::protocol_kind(
crate::error::ProtocolErrorKind::TruncatedFrame,
"fake broker: truncated topic id in ConsumerGroupHeartbeat",
));
}
buf.copy_to_slice(&mut topic_id);
let partition_count = read_compact_array_len(buf)?;
let mut partitions = Vec::with_capacity(partition_count);
for _ in 0..partition_count {
partitions.push(i32::decode(buf)?);
}
skip_tagged_fields(buf)?;
entries.push(HeartbeatTopicPartitions {
topic_id,
partitions,
});
}
Some(entries)
}
};
skip_tagged_fields(buf)?;
Ok(Self {
group_id,
member_id,
member_epoch,
instance_id,
rack_id,
rebalance_timeout_ms,
subscribed_topic_names,
subscribed_topic_regex,
server_assignor,
topic_partitions,
})
}
}
pub(crate) fn write_heartbeat_assignment(
out: &mut BytesMut,
assignment: Option<&[HeartbeatTopicPartitions]>,
) -> Result<()> {
match assignment {
None => {
out.put_i8(-1);
Ok(())
}
Some(entries) => {
out.put_i8(1);
write_compact_array_len(out, entries.len())?;
for entry in entries {
out.put_slice(&entry.topic_id);
write_compact_array_len(out, entry.partitions.len())?;
for partition in &entry.partitions {
out.put_i32(*partition);
}
write_empty_tagged_fields(out)?;
}
write_empty_tagged_fields(out)
}
}
}
#[derive(Debug, Clone)]
pub(crate) struct ListOffsetsReqPartition {
pub partition_index: i32,
pub current_leader_epoch: i32,
pub timestamp: i64,
}
#[derive(Debug, Clone)]
pub(crate) struct ListOffsetsReqTopic {
pub name: String,
pub partitions: Vec<ListOffsetsReqPartition>,
}
#[derive(Debug, Clone)]
pub(crate) struct ListOffsetsReq {
pub topics: Vec<ListOffsetsReqTopic>,
}
impl ListOffsetsReq {
pub(crate) fn read(buf: &mut impl Buf) -> Result<Self> {
let _replica_id = i32::decode(buf)?;
let _isolation_level = i8::decode(buf)?;
let topic_count = read_array_len(buf)?;
let mut topics = Vec::with_capacity(topic_count);
for _ in 0..topic_count {
let name = read_string(buf)?;
let partition_count = read_array_len(buf)?;
let mut partitions = Vec::with_capacity(partition_count);
for _ in 0..partition_count {
let partition_index = i32::decode(buf)?;
let current_leader_epoch = i32::decode(buf)?;
let timestamp = i64::decode(buf)?;
partitions.push(ListOffsetsReqPartition {
partition_index,
current_leader_epoch,
timestamp,
});
}
topics.push(ListOffsetsReqTopic { name, partitions });
}
Ok(Self { topics })
}
}
#[derive(Debug, Clone)]
pub(crate) struct FindCoordinatorReq {
pub key: String,
pub key_type: i8,
}
impl FindCoordinatorReq {
pub(crate) fn read(buf: &mut impl Buf) -> Result<Self> {
let key = read_string(buf)?;
let key_type = i8::decode(buf)?;
Ok(Self { key, key_type })
}
}
#[derive(Debug, Clone)]
pub(crate) struct JoinGroupReqProtocol {
pub name: String,
pub metadata: Bytes,
}
#[derive(Debug, Clone)]
pub(crate) struct JoinGroupReq {
pub group_id: String,
pub member_id: String,
pub group_instance_id: Option<String>,
pub protocol_type: String,
pub protocols: Vec<JoinGroupReqProtocol>,
}
impl JoinGroupReq {
pub(crate) fn read(buf: &mut impl Buf) -> Result<Self> {
let group_id = read_string(buf)?;
let _session_timeout_ms = i32::decode(buf)?;
let _rebalance_timeout_ms = i32::decode(buf)?;
let member_id = read_string(buf)?;
let group_instance_id = read_nullable_string(buf)?;
let protocol_type = read_string(buf)?;
let protocol_count = read_array_len(buf)?;
let mut protocols = Vec::with_capacity(protocol_count);
for _ in 0..protocol_count {
let name = read_string(buf)?;
let metadata = read_bytes(buf)?;
protocols.push(JoinGroupReqProtocol { name, metadata });
}
Ok(Self {
group_id,
member_id,
group_instance_id,
protocol_type,
protocols,
})
}
}
#[derive(Debug, Clone)]
pub(crate) struct SyncGroupReqAssignment {
pub member_id: String,
pub assignment: Bytes,
}
#[derive(Debug, Clone)]
pub(crate) struct SyncGroupReq {
pub group_id: String,
pub generation_id: i32,
pub member_id: String,
pub assignments: Vec<SyncGroupReqAssignment>,
}
impl SyncGroupReq {
pub(crate) fn read(buf: &mut impl Buf) -> Result<Self> {
let group_id = read_string(buf)?;
let generation_id = i32::decode(buf)?;
let member_id = read_string(buf)?;
let _group_instance_id = read_nullable_string(buf)?;
let count = read_array_len(buf)?;
let mut assignments = Vec::with_capacity(count);
for _ in 0..count {
let member_id = read_string(buf)?;
let assignment = read_bytes(buf)?;
assignments.push(SyncGroupReqAssignment {
member_id,
assignment,
});
}
Ok(Self {
group_id,
generation_id,
member_id,
assignments,
})
}
}
#[derive(Debug, Clone)]
pub(crate) struct HeartbeatReq {
pub group_id: String,
pub generation_id: i32,
pub member_id: String,
}
impl HeartbeatReq {
pub(crate) fn read(buf: &mut impl Buf) -> Result<Self> {
let group_id = read_string(buf)?;
let generation_id = i32::decode(buf)?;
let member_id = read_string(buf)?;
let _group_instance_id = read_nullable_string(buf)?;
Ok(Self {
group_id,
generation_id,
member_id,
})
}
}
#[derive(Debug, Clone)]
pub(crate) struct LeaveGroupReq {
pub group_id: String,
pub members: Vec<(String, Option<String>)>,
}
impl LeaveGroupReq {
pub(crate) fn read(buf: &mut impl Buf) -> Result<Self> {
let group_id = read_string(buf)?;
let count = read_array_len(buf)?;
let mut members = Vec::with_capacity(count);
for _ in 0..count {
let member_id = read_string(buf)?;
let instance = read_nullable_string(buf)?;
members.push((member_id, instance));
}
Ok(Self { group_id, members })
}
}
#[derive(Debug, Clone)]
pub(crate) struct OffsetCommitReqPartition {
pub partition_index: i32,
pub committed_offset: i64,
pub committed_leader_epoch: i32,
pub committed_metadata: Option<String>,
}
#[derive(Debug, Clone)]
pub(crate) struct OffsetCommitReqTopic {
pub name: String,
pub partitions: Vec<OffsetCommitReqPartition>,
}
#[derive(Debug, Clone)]
pub(crate) struct OffsetCommitReq {
pub group_id: String,
pub generation_id: i32,
pub member_id: String,
pub topics: Vec<OffsetCommitReqTopic>,
}
impl OffsetCommitReq {
pub(crate) fn read(buf: &mut impl Buf) -> Result<Self> {
let group_id = read_string(buf)?;
let generation_id = i32::decode(buf)?;
let member_id = read_string(buf)?;
let _group_instance_id = read_nullable_string(buf)?;
let topic_count = read_array_len(buf)?;
let mut topics = Vec::with_capacity(topic_count);
for _ in 0..topic_count {
let name = read_string(buf)?;
let partition_count = read_array_len(buf)?;
let mut partitions = Vec::with_capacity(partition_count);
for _ in 0..partition_count {
let partition_index = i32::decode(buf)?;
let committed_offset = i64::decode(buf)?;
let committed_leader_epoch = i32::decode(buf)?;
let committed_metadata = read_nullable_string(buf)?;
partitions.push(OffsetCommitReqPartition {
partition_index,
committed_offset,
committed_leader_epoch,
committed_metadata,
});
}
topics.push(OffsetCommitReqTopic { name, partitions });
}
Ok(Self {
group_id,
generation_id,
member_id,
topics,
})
}
}
#[derive(Debug, Clone)]
pub(crate) struct OffsetFetchReq {
pub group_id: String,
pub topics: Option<Vec<(String, Vec<i32>)>>,
}
impl OffsetFetchReq {
pub(crate) fn read(buf: &mut impl Buf) -> Result<Self> {
let group_id = read_string(buf)?;
let topics = match read_nullable_array_len(buf)? {
None => None,
Some(count) => {
let mut topics = Vec::with_capacity(count);
for _ in 0..count {
let name = read_string(buf)?;
let partition_count = read_array_len(buf)?;
let mut partitions = Vec::with_capacity(partition_count);
for _ in 0..partition_count {
partitions.push(i32::decode(buf)?);
}
topics.push((name, partitions));
}
Some(topics)
}
};
Ok(Self { group_id, topics })
}
}
#[derive(Debug, Clone)]
pub(crate) struct InitProducerIdReq {
pub transactional_id: Option<String>,
}
impl InitProducerIdReq {
pub(crate) fn read(buf: &mut impl Buf) -> Result<Self> {
let transactional_id = read_nullable_string(buf)?.filter(|s| !s.is_empty());
let _transaction_timeout_ms = i32::decode(buf)?;
Ok(Self { transactional_id })
}
}
#[derive(Debug, Clone)]
pub(crate) struct AddPartitionsToTxnReq {
pub transactional_id: String,
pub producer_id: i64,
pub producer_epoch: i16,
pub partitions: Vec<(String, i32)>,
}
impl AddPartitionsToTxnReq {
pub(crate) fn read(buf: &mut impl Buf) -> Result<Self> {
let transactional_id = read_string(buf)?;
let producer_id = i64::decode(buf)?;
let producer_epoch = i16::decode(buf)?;
let topic_count = read_array_len(buf)?;
let mut partitions = Vec::with_capacity(topic_count);
for _ in 0..topic_count {
let name = read_string(buf)?;
let partition_count = read_array_len(buf)?;
for _ in 0..partition_count {
partitions.push((name.clone(), i32::decode(buf)?));
}
}
Ok(Self {
transactional_id,
producer_id,
producer_epoch,
partitions,
})
}
}
#[derive(Debug, Clone)]
pub(crate) struct AddOffsetsToTxnReq {
pub transactional_id: String,
pub producer_id: i64,
pub producer_epoch: i16,
pub group_id: String,
}
impl AddOffsetsToTxnReq {
pub(crate) fn read(buf: &mut impl Buf) -> Result<Self> {
Ok(Self {
transactional_id: read_string(buf)?,
producer_id: i64::decode(buf)?,
producer_epoch: i16::decode(buf)?,
group_id: read_string(buf)?,
})
}
}
#[derive(Debug, Clone)]
pub(crate) struct EndTxnReq {
pub transactional_id: String,
pub producer_id: i64,
pub producer_epoch: i16,
pub committed: bool,
}
impl EndTxnReq {
pub(crate) fn read(buf: &mut impl Buf) -> Result<Self> {
let transactional_id = read_compact_string(buf)?;
let producer_id = i64::decode(buf)?;
let producer_epoch = i16::decode(buf)?;
let committed = i8::decode(buf)? != 0;
skip_tagged_fields(buf)?;
Ok(Self {
transactional_id,
producer_id,
producer_epoch,
committed,
})
}
}
#[derive(Debug, Clone)]
pub(crate) struct TxnOffsetCommitReqPartition {
pub topic: String,
pub partition: i32,
pub committed_offset: i64,
pub committed_leader_epoch: i32,
pub metadata: Option<String>,
}
#[derive(Debug, Clone)]
pub(crate) struct TxnOffsetCommitReq {
pub transactional_id: String,
pub group_id: String,
pub producer_id: i64,
pub producer_epoch: i16,
pub generation_id: i32,
pub member_id: String,
pub group_instance_id: Option<String>,
pub offsets: Vec<TxnOffsetCommitReqPartition>,
}
impl TxnOffsetCommitReq {
pub(crate) fn read(buf: &mut impl Buf) -> Result<Self> {
let transactional_id = read_compact_string(buf)?;
let group_id = read_compact_string(buf)?;
let producer_id = i64::decode(buf)?;
let producer_epoch = i16::decode(buf)?;
let generation_id = i32::decode(buf)?;
let member_id = read_compact_string(buf)?;
let group_instance_id = read_compact_nullable_string(buf)?;
let topic_count = read_compact_array_len(buf)?;
let mut offsets = Vec::with_capacity(topic_count);
for _ in 0..topic_count {
let topic = read_compact_string(buf)?;
let partition_count = read_compact_array_len(buf)?;
for _ in 0..partition_count {
let partition = i32::decode(buf)?;
let committed_offset = i64::decode(buf)?;
let committed_leader_epoch = i32::decode(buf)?;
let metadata = read_compact_nullable_string(buf)?;
skip_tagged_fields(buf)?;
offsets.push(TxnOffsetCommitReqPartition {
topic: topic.clone(),
partition,
committed_offset,
committed_leader_epoch,
metadata,
});
}
skip_tagged_fields(buf)?;
}
skip_tagged_fields(buf)?;
Ok(Self {
transactional_id,
group_id,
producer_id,
producer_epoch,
generation_id,
member_id,
group_instance_id,
offsets,
})
}
}
#[derive(Debug, Clone)]
pub(crate) struct CreateTopicsReqTopic {
pub name: String,
pub num_partitions: i32,
}
#[derive(Debug, Clone)]
pub(crate) struct CreateTopicsReq {
pub topics: Vec<CreateTopicsReqTopic>,
pub validate_only: bool,
}
impl CreateTopicsReq {
pub(crate) fn read(buf: &mut impl Buf) -> Result<Self> {
let topic_count = read_array_len(buf)?;
let mut topics = Vec::with_capacity(topic_count);
for _ in 0..topic_count {
let name = read_string(buf)?;
let num_partitions = i32::decode(buf)?;
let _replication_factor = i16::decode(buf)?;
let assignment_count = read_array_len(buf)?;
for _ in 0..assignment_count {
let _ = i32::decode(buf)?;
let broker_count = read_array_len(buf)?;
for _ in 0..broker_count {
let _ = i32::decode(buf)?;
}
}
let config_count = read_array_len(buf)?;
for _ in 0..config_count {
let _ = read_string(buf)?;
let _ = read_nullable_string(buf)?;
}
topics.push(CreateTopicsReqTopic {
name,
num_partitions,
});
}
let _timeout_ms = i32::decode(buf)?;
let validate_only = bool::decode(buf)?;
Ok(Self {
topics,
validate_only,
})
}
}
#[derive(Debug, Clone)]
pub(crate) struct DeleteTopicsReq {
pub topic_names: Vec<String>,
}
impl DeleteTopicsReq {
pub(crate) fn read(buf: &mut impl Buf) -> Result<Self> {
let count = read_array_len(buf)?;
let mut topic_names = Vec::with_capacity(count);
for _ in 0..count {
topic_names.push(read_string(buf)?);
}
let _timeout_ms = i32::decode(buf)?;
Ok(Self { topic_names })
}
}
const BATCH_BASE_OFFSET_POS: usize = 0;
const BATCH_LEADER_EPOCH_POS: usize = 12;
const BATCH_RECORDS_COUNT_POS: usize = 57;
const BATCH_HEADER_LEN: usize = 61;
pub(crate) fn batch_record_count(batch: &[u8]) -> Option<i64> {
if batch.len() < BATCH_HEADER_LEN {
return None;
}
let bytes = batch.get(BATCH_RECORDS_COUNT_POS..BATCH_RECORDS_COUNT_POS + 4)?;
let count = i32::from_be_bytes([
*bytes.first()?,
*bytes.get(1)?,
*bytes.get(2)?,
*bytes.get(3)?,
]);
Some(i64::from(count).max(0))
}
pub(crate) fn stamp_batch(batch: &Bytes, base_offset: i64, leader_epoch: i32) -> Bytes {
if batch.len() < BATCH_HEADER_LEN {
return batch.clone();
}
let mut out = BytesMut::from(&batch[..]);
out[BATCH_BASE_OFFSET_POS..BATCH_BASE_OFFSET_POS + 8]
.copy_from_slice(&base_offset.to_be_bytes());
out[BATCH_LEADER_EPOCH_POS..BATCH_LEADER_EPOCH_POS + 4]
.copy_from_slice(&leader_epoch.to_be_bytes());
out.freeze()
}
pub(crate) fn batch_base_offset(batch: &[u8]) -> Option<i64> {
if batch.len() < BATCH_HEADER_LEN {
return None;
}
let bytes = batch.get(BATCH_BASE_OFFSET_POS..BATCH_BASE_OFFSET_POS + 8)?;
let mut arr = [0u8; 8];
arr.copy_from_slice(bytes);
Some(i64::from_be_bytes(arr))
}
#[cfg(test)]
#[allow(clippy::unwrap_used, clippy::expect_used, clippy::panic)]
mod tests {
use super::*;
use crate::protocol::{Record, RecordBatch};
#[test]
fn stamping_a_batch_preserves_its_crc() {
let mut batch = RecordBatch::new();
batch.records = vec![
Record::new(None, Some(Bytes::from_static(b"a"))).with_offset_delta(0),
Record::new(None, Some(Bytes::from_static(b"b"))).with_offset_delta(1),
];
let encoded = batch.encode().unwrap();
assert_eq!(batch_record_count(&encoded), Some(2));
let stamped = stamp_batch(&encoded, 41, 7);
assert_eq!(batch_base_offset(&stamped), Some(41));
let mut cursor = stamped.clone();
let decoded = RecordBatch::decode(&mut cursor).expect("stamped batch must still decode");
assert_eq!(decoded.base_offset, 41);
assert_eq!(decoded.partition_leader_epoch, 7);
assert_eq!(decoded.records.len(), 2);
}
#[test]
fn short_buffers_are_rejected_rather_than_indexed() {
assert_eq!(batch_record_count(&[0u8; 4]), None);
assert_eq!(batch_base_offset(&[0u8; 4]), None);
let short = Bytes::from_static(&[0u8; 4]);
assert_eq!(stamp_batch(&short, 1, 1), short);
}
#[test]
fn array_lengths_beyond_the_safety_limit_are_rejected() {
let mut buf = BytesMut::new();
buf.put_i32(MAX_ARRAY_LEN as i32 + 1);
assert!(read_array_len(&mut buf.freeze()).is_err());
}
#[test]
fn null_array_length_is_rejected_for_non_nullable_arrays() {
let mut buf = BytesMut::new();
buf.put_i32(-1);
assert!(read_array_len(&mut buf.freeze()).is_err());
}
}
#[derive(Debug)]
pub(crate) struct ShareGroupHeartbeatReq {
pub group_id: String,
pub member_id: String,
pub member_epoch: i32,
pub subscribed_topic_names: Option<Vec<String>>,
}
impl ShareGroupHeartbeatReq {
pub(crate) fn read(buf: &mut impl Buf) -> Result<Self> {
let group_id = read_compact_string(buf)?;
let member_id = read_compact_string(buf)?;
let member_epoch = i32::decode(buf)?;
let _rack_id = read_compact_nullable_string(buf)?;
let subscribed_topic_names = match read_compact_nullable_array_len(buf)? {
None => None,
Some(n) => {
let mut topics = Vec::with_capacity(n);
for _ in 0..n {
topics.push(read_compact_string(buf)?);
}
Some(topics)
}
};
skip_tagged_fields(buf)?;
Ok(Self {
group_id,
member_id,
member_epoch,
subscribed_topic_names,
})
}
}
#[derive(Debug)]
pub(crate) struct ShareAckBatch {
pub first_offset: i64,
pub last_offset: i64,
pub acknowledge_types: Vec<i8>,
}
#[derive(Debug)]
pub(crate) struct ShareTopicPartitionAcks {
pub topic_id: [u8; 16],
pub partition_index: i32,
pub acknowledgement_batches: Vec<ShareAckBatch>,
}
fn read_share_topics(buf: &mut impl Buf) -> Result<Vec<ShareTopicPartitionAcks>> {
let topic_count = read_compact_array_len(buf)?;
let mut out = Vec::new();
for _ in 0..topic_count {
let mut topic_id = [0u8; 16];
if buf.remaining() < 16 {
return Err(KrafkaError::protocol_kind(
ProtocolErrorKind::TruncatedFrame,
"not enough bytes for share topic_id",
));
}
buf.copy_to_slice(&mut topic_id);
let part_count = read_compact_array_len(buf)?;
for _ in 0..part_count {
let partition_index = i32::decode(buf)?;
let batch_count = read_compact_array_len(buf)?;
let mut batches = Vec::with_capacity(batch_count);
for _ in 0..batch_count {
let first_offset = i64::decode(buf)?;
let last_offset = i64::decode(buf)?;
let type_count = read_compact_array_len(buf)?;
let mut acknowledge_types = Vec::with_capacity(type_count);
for _ in 0..type_count {
acknowledge_types.push(i8::decode(buf)?);
}
skip_tagged_fields(buf)?;
batches.push(ShareAckBatch {
first_offset,
last_offset,
acknowledge_types,
});
}
skip_tagged_fields(buf)?;
out.push(ShareTopicPartitionAcks {
topic_id,
partition_index,
acknowledgement_batches: batches,
});
}
skip_tagged_fields(buf)?;
}
Ok(out)
}
#[derive(Debug)]
pub(crate) struct ShareFetchReq {
pub group_id: Option<String>,
pub member_id: Option<String>,
pub max_records: i32,
pub topics: Vec<ShareTopicPartitionAcks>,
}
impl ShareFetchReq {
pub(crate) fn read(buf: &mut impl Buf, version: i16) -> Result<Self> {
let group_id = read_compact_nullable_string(buf)?;
let member_id = read_compact_nullable_string(buf)?;
let _share_session_epoch = i32::decode(buf)?;
let _max_wait_ms = i32::decode(buf)?;
let _min_bytes = i32::decode(buf)?;
let _max_bytes = i32::decode(buf)?;
let max_records = i32::decode(buf)?;
let _batch_size = i32::decode(buf)?;
if version >= 2 {
let _share_acquire_mode = i8::decode(buf)?;
let _is_renew_ack = i8::decode(buf)?;
}
let topics = read_share_topics(buf)?;
let forgotten_count = read_compact_array_len(buf)?;
for _ in 0..forgotten_count {
if buf.remaining() < 16 {
return Err(KrafkaError::protocol_kind(
ProtocolErrorKind::TruncatedFrame,
"not enough bytes for forgotten topic_id",
));
}
buf.advance(16);
let n = read_compact_array_len(buf)?;
for _ in 0..n {
let _ = i32::decode(buf)?;
}
skip_tagged_fields(buf)?;
}
skip_tagged_fields(buf)?;
Ok(Self {
group_id,
member_id,
max_records,
topics,
})
}
}
#[derive(Debug)]
pub(crate) struct ShareAcknowledgeReq {
pub group_id: Option<String>,
pub member_id: Option<String>,
pub topics: Vec<ShareTopicPartitionAcks>,
}
impl ShareAcknowledgeReq {
pub(crate) fn read(buf: &mut impl Buf, version: i16) -> Result<Self> {
let group_id = read_compact_nullable_string(buf)?;
let member_id = read_compact_nullable_string(buf)?;
let _share_session_epoch = i32::decode(buf)?;
if version >= 2 {
let _is_renew_ack = i8::decode(buf)?;
}
let topics = read_share_topics(buf)?;
skip_tagged_fields(buf)?;
Ok(Self {
group_id,
member_id,
topics,
})
}
}
#[derive(Debug)]
pub(crate) struct FeatureUpdate {
pub feature: String,
pub max_version_level: i16,
}
#[derive(Debug)]
pub(crate) struct UpdateFeaturesReq {
pub feature_updates: Vec<FeatureUpdate>,
pub validate_only: bool,
}
impl UpdateFeaturesReq {
pub(crate) fn read(buf: &mut impl Buf, version: i16) -> Result<Self> {
let _timeout_ms = i32::decode(buf)?;
let count = read_compact_array_len(buf)?;
let mut feature_updates = Vec::with_capacity(count);
for _ in 0..count {
let feature = read_compact_string(buf)?;
let max_version_level = i16::decode(buf)?;
let _ = i8::decode(buf)?;
skip_tagged_fields(buf)?;
feature_updates.push(FeatureUpdate {
feature,
max_version_level,
});
}
let validate_only = if version >= 1 {
i8::decode(buf)? != 0
} else {
false
};
skip_tagged_fields(buf)?;
Ok(Self {
feature_updates,
validate_only,
})
}
}
#[derive(Debug)]
pub(crate) struct StreamsGroupDescribeReq {
pub group_ids: Vec<String>,
pub include_authorized_operations: bool,
}
impl StreamsGroupDescribeReq {
pub(crate) fn read(buf: &mut impl Buf) -> Result<Self> {
let count = read_compact_array_len(buf)?;
let mut group_ids = Vec::with_capacity(count);
for _ in 0..count {
group_ids.push(read_compact_string(buf)?);
}
let include_authorized_operations = i8::decode(buf)? != 0;
skip_tagged_fields(buf)?;
Ok(Self {
group_ids,
include_authorized_operations,
})
}
}
pub(crate) fn write_presence(buf: &mut impl BufMut, present: bool) {
buf.put_i8(if present { 1 } else { -1 });
}