kafka-ts 0.0.3-beta → 0.0.3

package/src/api/index.ts

@@ -37,7 +37,9 @@ export const API = {
     SYNC_GROUP,
 };
 
-export const getApiName = (api: Api<unknown, unknown>) => Object.entries(API).find(([, v]) => v === api)?.[0];
+const apiNameByKey = Object.fromEntries(Object.entries(API).map(([k, v]) => [v.apiKey, k]));
+
+export const getApiName = <Request, Response>(api: Api<Request, Response>) => apiNameByKey[api.apiKey];
 
 export const API_ERROR = {
     UNKNOWN_SERVER_ERROR: -1,

package/src/api/metadata.ts

@@ -1,7 +1,7 @@
 import { createApi } from '../utils/api';
 import { KafkaTSApiError } from '../utils/error';
 
-export type Metadata = ReturnType<(typeof METADATA)['response']>;
+export type Metadata = Awaited<ReturnType<(typeof METADATA)['response']>>;
 
 export const METADATA = createApi({
     apiKey: 3,

package/src/api/offset-commit.ts

@@ -3,7 +3,7 @@ import { KafkaTSApiError } from '../utils/error';
 
 export const OFFSET_COMMIT = createApi({
     apiKey: 8,
-    apiVersion: 9,
+    apiVersion: 8,
     request: (
         encoder,
         data: {

package/src/api/offset-fetch.ts

@@ -3,14 +3,12 @@ import { KafkaTSApiError } from '../utils/error';
 
 export const OFFSET_FETCH = createApi({
     apiKey: 9,
-    apiVersion: 9,
+    apiVersion: 8,
     request: (
         encoder,
         data: {
             groups: {
                 groupId: string;
-                memberId: string | null;
-                memberEpoch: number;
                 topics: {
                     name: string;
                     partitionIndexes: number[];
@@ -24,8 +22,6 @@ export const OFFSET_FETCH = createApi({
             .writeCompactArray(data.groups, (encoder, group) =>
                 encoder
                     .writeCompactString(group.groupId)
-                    .writeCompactString(group.memberId)
-                    .writeInt32(group.memberEpoch)
                     .writeCompactArray(group.topics, (encoder, topic) =>
                         encoder
                             .writeCompactString(topic.name)

package/src/api/produce.ts

@@ -4,7 +4,7 @@ import { KafkaTSApiError } from '../utils/error.js';
 
 export const PRODUCE = createApi({
     apiKey: 0,
-    apiVersion: 10,
+    apiVersion: 9,
     request: (
         encoder,
         data: {
@@ -28,11 +28,11 @@ export const PRODUCE = createApi({
                         attributes: number;
                         timestampDelta: bigint;
                         offsetDelta: number;
-                        key: string | null;
-                        value: string | null;
+                        key: Buffer | null;
+                        value: Buffer | null;
                         headers: {
-                            key: string;
-                            value: string;
+                            key: Buffer;
+                            value: Buffer;
                         }[];
                     }[];
                 }[];
@@ -61,14 +61,13 @@ export const PRODUCE = createApi({
                                     .writeInt8(record.attributes)
                                     .writeVarLong(record.timestampDelta)
                                     .writeVarInt(record.offsetDelta)
-                                    .writeVarIntString(record.key)
-                                    .writeVarIntString(record.value)
+                                    .writeVarIntBuffer(record.key)
+                                    .writeVarIntBuffer(record.value)
                                     .writeVarIntArray(record.headers, (encoder, header) =>
-                                        encoder.writeVarIntString(header.key).writeVarIntString(header.value),
-                                    )
-                                    .value();
+                                        encoder.writeVarIntBuffer(header.key).writeVarIntBuffer(header.value),
+                                    );
 
-                                return encoder.writeVarInt(recordBody.length).write(recordBody);
+                                return encoder.writeVarInt(recordBody.getByteLength()).writeEncoder(recordBody);
                             })
                             .value();
 
@@ -76,19 +75,17 @@ export const PRODUCE = createApi({
                             .writeInt32(partition.partitionLeaderEpoch)
                             .writeInt8(2) // magic byte
                             .writeUInt32(unsigned(crc32C(batchBody)))
-                            .write(batchBody)
-                            .value();
+                            .write(batchBody);
 
                         const batch = new Encoder()
                             .writeInt64(partition.baseOffset)
-                            .writeInt32(batchHeader.length)
-                            .write(batchHeader)
-                            .value();
+                            .writeInt32(batchHeader.getByteLength())
+                            .writeEncoder(batchHeader);
 
                         return encoder
                             .writeInt32(partition.index)
-                            .writeUVarInt(batch.length + 1) // batch size
-                            .write(batch)
+                            .writeUVarInt(batch.getByteLength() + 1) // batch size
+                            .writeEncoder(batch)
                             .writeUVarInt(0);
                     })
                     .writeUVarInt(0),

package/src/auth/index.ts

@@ -0,0 +1,2 @@
+export { saslPlain } from './plain';
+export { saslScramSha256, saslScramSha512 } from './scram';

package/src/auth/plain.ts

@@ -0,0 +1,10 @@
+import { API } from "../api";
+import { SASLProvider } from "../broker";
+
+export const saslPlain = ({ username, password }: { username: string; password: string }): SASLProvider => ({
+    mechanism: 'PLAIN',
+    authenticate: async ({ sendRequest }) => {
+        const authBytes = [null, username, password].join('\u0000');
+        await sendRequest(API.SASL_AUTHENTICATE, { authBytes: Buffer.from(authBytes) });
+    },
+});

package/src/auth/scram.ts

@@ -0,0 +1,52 @@
+import { API } from '../api';
+import { SASLProvider } from '../broker';
+import { base64Decode, base64Encode, generateNonce, hash, hmac, saltPassword, xor } from '../utils/crypto';
+import { KafkaTSError } from '../utils/error';
+
+const saslScram =
+    ({ mechanism, keyLength, digest }: { mechanism: string; keyLength: number; digest: string }) =>
+    ({ username, password }: { username: string; password: string }): SASLProvider => ({
+        mechanism,
+        authenticate: async ({ sendRequest }) => {
+            const nonce = generateNonce();
+            const firstMessage = `n=${username},r=${nonce}`;
+
+            const { authBytes } = await sendRequest(API.SASL_AUTHENTICATE, {
+                authBytes: Buffer.from(`n,,${firstMessage}`),
+            });
+            if (!authBytes) {
+                throw new KafkaTSError('No auth response');
+            }
+
+            const response = Object.fromEntries(
+                authBytes
+                    .toString()
+                    .split(',')
+                    .map((pair) => pair.split('=')),
+            ) as { r: string; s: string; i: string };
+
+            const rnonce = response.r;
+            if (!rnonce.startsWith(nonce)) {
+                throw new KafkaTSError('Invalid nonce');
+            }
+            const iterations = parseInt(response.i);
+            const salt = base64Decode(response.s);
+
+            const saltedPassword = await saltPassword(password, salt, iterations, keyLength, digest);
+            const clientKey = hmac(saltedPassword, 'Client Key', digest);
+            const clientKeyHash = hash(clientKey, digest);
+
+            let finalMessage = `c=${base64Encode('n,,')},r=${rnonce}`;
+
+            const fullMessage = `${firstMessage},${authBytes.toString()},${finalMessage}`;
+            const clientSignature = hmac(clientKeyHash, fullMessage, digest);
+            const clientProof = base64Encode(xor(clientKey, clientSignature));
+
+            finalMessage += `,p=${clientProof}`;
+
+            await sendRequest(API.SASL_AUTHENTICATE, { authBytes: Buffer.from(finalMessage) });
+        },
+    });
+
+export const saslScramSha256 = saslScram({ mechanism: 'SCRAM-SHA-256', keyLength: 32, digest: 'sha256' });
+export const saslScramSha512 = saslScram({ mechanism: 'SCRAM-SHA-512', keyLength: 64, digest: 'sha512' });

package/src/broker.ts

@@ -5,12 +5,15 @@ import { Connection, SendRequest } from './connection';
 import { KafkaTSError } from './utils/error';
 import { memo } from './utils/memo';
 
-export type SASLOptions = { mechanism: 'PLAIN'; username: string; password: string };
+export type SASLProvider = {
+    mechanism: string;
+    authenticate: (context: { sendRequest: SendRequest }) => Promise<void>;
+};
 
 type BrokerOptions = {
     clientId: string | null;
     options: TcpSocketConnectOpts;
-    sasl: SASLOptions | null;
+    sasl: SASLProvider | null;
     ssl: TLSSocketOptions | null;
 };
 
@@ -51,7 +54,7 @@ export class Broker {
             }
             const { apiVersion } = apiByKey[apiKey];
             if (apiVersion < minVersion || apiVersion > maxVersion) {
-                throw new KafkaTSError(`API ${apiKey} version ${apiVersion} is not supported by the broker`);
+                throw new KafkaTSError(`API ${apiKey} version ${apiVersion} is not supported by the broker (minVersion=${minVersion}, maxVersion=${maxVersion})`);
             }
         });
     }
@@ -64,11 +67,6 @@ export class Broker {
     }
 
     private async saslAuthenticate() {
-        if (this.options.sasl?.mechanism !== 'PLAIN') {
-            return;
-        }
-        const { username, password } = this.options.sasl;
-        const authBytes = [null, username, password].join('\u0000');
-        await this.sendRequest(API.SASL_AUTHENTICATE, { authBytes: Buffer.from(authBytes) });
+        await this.options.sasl?.authenticate({ sendRequest: this.sendRequest });
     }
 }

package/src/client.ts

@@ -1,6 +1,6 @@
 import { TcpSocketConnectOpts } from 'net';
 import { TLSSocketOptions } from 'tls';
-import { SASLOptions } from './broker';
+import { SASLProvider } from './broker';
 import { Cluster } from './cluster';
 import { Consumer, ConsumerOptions } from './consumer/consumer';
 import { Producer, ProducerOptions } from './producer/producer';
@@ -8,7 +8,7 @@ import { Producer, ProducerOptions } from './producer/producer';
 type ClientOptions = {
     clientId?: string | null;
     bootstrapServers: TcpSocketConnectOpts[];
-    sasl?: SASLOptions | null;
+    sasl?: SASLProvider | null;
     ssl?: TLSSocketOptions | null;
 };
 

package/src/cluster.test.ts

@@ -3,24 +3,26 @@ import { readFileSync } from 'fs';
 import { afterAll, beforeAll, describe, expect, it } from 'vitest';
 import { API } from './api';
 import { KEY_TYPE } from './api/find-coordinator';
+import { saslPlain } from './auth';
 import { createKafkaClient } from './client';
 import { Cluster } from './cluster';
 import { KafkaTSApiError } from './utils/error';
 
-export const kafka = createKafkaClient({
+const kafka = createKafkaClient({
     clientId: 'kafka-ts',
     bootstrapServers: [{ host: 'localhost', port: 9092 }],
-    sasl: { mechanism: 'PLAIN', username: 'admin', password: 'admin' },
+    sasl: saslPlain({ username: 'admin', password: 'admin' }),
     ssl: { ca: readFileSync('./certs/ca.crt').toString() },
 });
 
-describe.sequential('Request handler', () => {
+describe.sequential('Low-level API', () => {
     const groupId = randomBytes(16).toString('hex');
 
     let cluster: Cluster;
 
     beforeAll(async () => {
-        cluster = await kafka.createCluster().connect();
+        cluster = await kafka.createCluster();
+        await cluster.connect();
 
         const metadataResult = await cluster.sendRequest(API.METADATA, {
             topics: null,
@@ -51,8 +53,8 @@ describe.sequential('Request handler', () => {
             topics: [
                 {
                     name: 'kafka-ts-test-topic',
-                    numPartitions: 1,
-                    replicationFactor: 1,
+                    numPartitions: 10,
+                    replicationFactor: 3,
                     assignments: [],
                     configs: [],
                 },
@@ -88,6 +90,7 @@ describe.sequential('Request handler', () => {
         expect(result).toMatchSnapshot();
     });
 
+    let partitionIndex = 0;
     let leaderId = 0;
 
     it('should request metadata for a topic', async () => {
@@ -96,6 +99,7 @@ describe.sequential('Request handler', () => {
             allowTopicAutoCreation: false,
             includeTopicAuthorizedOperations: false,
         });
+        partitionIndex = result.topics[0].partitions[0].partitionIndex;
         leaderId = result.topics[0].partitions[0].leaderId;
         result.controllerId = 0;
         result.topics.forEach((topic) => {
@@ -133,7 +137,7 @@ describe.sequential('Request handler', () => {
                     name: 'kafka-ts-test-topic',
                     partitionData: [
                         {
-                            index: 0,
+                            index: partitionIndex,
                             baseOffset: 0n,
                             partitionLeaderEpoch: 0,
                             attributes: 0,
@@ -148,12 +152,12 @@ describe.sequential('Request handler', () => {
                                     attributes: 0,
                                     offsetDelta: 0,
                                     timestampDelta: 0n,
-                                    key: 'key',
-                                    value: 'value',
+                                    key: Buffer.from('key'),
+                                    value: Buffer.from('value'),
                                     headers: [
                                         {
-                                            key: 'header-key',
-                                            value: 'header-value',
+                                            key: Buffer.from('header-key'),
+                                            value: Buffer.from('header-value'),
                                         },
                                     ],
                                 },
@@ -179,7 +183,7 @@ describe.sequential('Request handler', () => {
                     topicId,
                     partitions: [
                         {
-                            partition: 0,
+                            partition: partitionIndex,
                             currentLeaderEpoch: -1,
                             fetchOffset: 0n,
                             lastFetchedEpoch: 0,
@@ -317,8 +321,6 @@ describe.sequential('Request handler', () => {
             groups: [
                 {
                     groupId,
-                    memberId,
-                    memberEpoch: 0,
                     topics: [
                         {
                             name: 'kafka-ts-test-topic',

package/src/cluster.ts

@@ -1,85 +1,83 @@
 import { TcpSocketConnectOpts } from 'net';
 import { TLSSocketOptions } from 'tls';
 import { API } from './api';
-import { Broker, SASLOptions } from './broker';
+import { Metadata } from './api/metadata';
+import { Broker, SASLProvider } from './broker';
 import { SendRequest } from './connection';
-import { ConnectionError, KafkaTSError } from './utils/error';
+import { KafkaTSError } from './utils/error';
+import { log } from './utils/logger';
 
 type ClusterOptions = {
     clientId: string | null;
     bootstrapServers: TcpSocketConnectOpts[];
-    sasl: SASLOptions | null;
+    sasl: SASLProvider | null;
     ssl: TLSSocketOptions | null;
 };
 
 export class Cluster {
-    private seedBroker: Broker;
+    private seedBroker = new Broker({ clientId: null, sasl: null, ssl: null, options: { port: 9092 } });
     private brokerById: Record<number, Broker> = {};
+    private brokerMetadata: Record<number, Metadata['brokers'][number]> = {};
 
-    constructor(private options: ClusterOptions) {
-        this.seedBroker = new Broker({
-            clientId: this.options.clientId,
-            sasl: this.options.sasl,
-            ssl: this.options.ssl,
-            options: this.options.bootstrapServers[0],
-        });
-    }
+    constructor(private options: ClusterOptions) {}
 
     public async connect() {
-        await this.connectSeedBroker();
+        this.seedBroker = await this.findSeedBroker();
+        this.brokerById = {};
+
         const metadata = await this.sendRequest(API.METADATA, {
             allowTopicAutoCreation: false,
             includeTopicAuthorizedOperations: false,
             topics: [],
         });
-
-        this.brokerById = Object.fromEntries(
-            metadata.brokers.map(({ nodeId, ...options }) => [
-                nodeId,
-                new Broker({
-                    clientId: this.options.clientId,
-                    sasl: this.options.sasl,
-                    ssl: this.options.ssl,
-                    options,
-                }),
-            ]),
-        );
-        return this;
+        this.brokerMetadata = Object.fromEntries(metadata.brokers.map((options) => [options.nodeId, options]));
     }
 
     public async disconnect() {
-        await Promise.all([
-            this.seedBroker.disconnect(),
-            ...Object.values(this.brokerById).map((broker) => broker.disconnect()),
-        ]);
+        await Promise.all([this.seedBroker.disconnect(), ...Object.values(this.brokerById).map((x) => x.disconnect())]);
     }
 
+    public setSeedBroker = async (nodeId: number) => {
+        await this.seedBroker.disconnect();
+        this.seedBroker = await this.acquireBroker(nodeId);
+    };
+
     public sendRequest: SendRequest = (...args) => this.seedBroker.sendRequest(...args);
 
     public sendRequestToNode =
         (nodeId: number): SendRequest =>
         async (...args) => {
-            const broker = this.brokerById[nodeId];
-            if (!broker) {
-                throw new ConnectionError(`Broker ${nodeId} is not available`);
+            if (!this.brokerById[nodeId]) {
+                this.brokerById[nodeId] = await this.acquireBroker(nodeId);
             }
-            await broker.ensureConnected();
-            return broker.sendRequest(...args);
+            return this.brokerById[nodeId].sendRequest(...args);
         };
 
-    private async connectSeedBroker() {
+    public async acquireBroker(nodeId: number) {
+        const broker = new Broker({
+            clientId: this.options.clientId,
+            sasl: this.options.sasl,
+            ssl: this.options.ssl,
+            options: this.brokerMetadata[nodeId],
+        });
+        await broker.connect();
+        return broker;
+    }
+
+    private async findSeedBroker() {
         const randomizedBrokers = this.options.bootstrapServers.toSorted(() => Math.random() - 0.5);
         for (const options of randomizedBrokers) {
             try {
-                this.seedBroker = await new Broker({
+                const broker = await new Broker({
                     clientId: this.options.clientId,
                     sasl: this.options.sasl,
                     ssl: this.options.ssl,
                     options,
-                }).connect();
-                return;
+                });
+                await broker.connect();
+                return broker;
             } catch (error) {
-                console.warn(`Failed to connect to seed broker ${options.host}:${options.port}`, error);
+                log.warn(`Failed to connect to seed broker ${options.host}:${options.port}`, error);
             }
         }
         throw new KafkaTSError('No seed brokers found');

package/src/codecs/gzip.ts

@@ -0,0 +1,9 @@
+import { gzip, unzip } from 'zlib';
+import { Codec } from './types';
+
+export const GZIP: Codec = {
+    compress: async (data) =>
+        new Promise<Buffer>((resolve, reject) => gzip(data, (err, result) => (err ? reject(err) : resolve(result)))),
+    decompress: async (data) =>
+        new Promise<Buffer>((resolve, reject) => unzip(data, (err, result) => (err ? reject(err) : resolve(result)))),
+};

package/src/codecs/index.ts

@@ -0,0 +1,16 @@
+import { GZIP } from './gzip';
+import { NONE } from './none';
+import { Codec } from './types';
+
+const codecs: Record<number, Codec> = {
+    0: NONE,
+    1: GZIP,
+};
+
+export const findCodec = (type: number) => {
+    const codec = codecs[type];
+    if (!codec) {
+        throw new Error(`Unsupported codec: ${type}`);
+    }
+    return codec;
+};

package/src/codecs/none.ts

@@ -0,0 +1,6 @@
+import { Codec } from './types';
+
+export const NONE: Codec = {
+    compress: (data: Buffer) => Promise.resolve(data),
+    decompress: (data: Buffer) => Promise.resolve(data),
+};

package/src/codecs/types.ts

@@ -0,0 +1,4 @@
+export type Codec = {
+    compress: (data: Buffer) => Promise<Buffer>;
+    decompress: (data: Buffer) => Promise<Buffer>;
+};

package/src/connection.ts

@@ -6,11 +6,12 @@ import { Api } from './utils/api';
 import { Decoder } from './utils/decoder';
 import { Encoder } from './utils/encoder';
 import { ConnectionError } from './utils/error';
+import { log } from './utils/logger';
 import { createTracer } from './utils/tracer';
 
 const trace = createTracer('Connection');
 
-export type ConnectionOptions = {
+type ConnectionOptions = {
     clientId: string | null;
     connection: TcpSocketConnectOpts;
     ssl: TLSSocketOptions | null;
@@ -24,14 +25,14 @@ export class Connection {
         [correlationId: number]: { resolve: (response: RawResonse) => void; reject: (error: Error) => void };
     } = {};
     private lastCorrelationId = 0;
-    private buffer: Buffer | null = null;
+    private chunks: Buffer[] = [];
 
     constructor(private options: ConnectionOptions) {}
 
     @trace()
     public async connect() {
         this.queue = {};
-        this.buffer = null;
+        this.chunks = [];
 
         await new Promise<void>((resolve, reject) => {
             const { ssl, connection } = this.options;
@@ -59,6 +60,7 @@ export class Connection {
         });
     }
 
+    @trace()
     public disconnect() {
         this.socket.removeAllListeners();
         return new Promise<void>((resolve) => {
@@ -69,9 +71,10 @@ export class Connection {
         });
     }
 
-    @trace((api, body) => ({ apiName: getApiName(api), body }))
+    @trace((api, body) => ({ message: getApiName(api), body }))
     public async sendRequest<Request, Response>(api: Api<Request, Response>, body: Request): Promise<Response> {
         const correlationId = this.nextCorrelationId();
+        const apiName = getApiName(api);
 
         const encoder = new Encoder()
             .writeInt16(api.apiKey)
@@ -79,18 +82,25 @@ export class Connection {
             .writeInt32(correlationId)
             .writeString(this.options.clientId);
 
-        const request = api.request(encoder, body).value();
-        const requestEncoder = new Encoder().writeInt32(request.length).write(request);
+        const request = api.request(encoder, body);
+        const requestEncoder = new Encoder().writeInt32(request.getByteLength()).writeEncoder(request);
 
+        let timeout: NodeJS.Timeout | undefined;
         const { responseDecoder, responseSize } = await new Promise<RawResonse>(async (resolve, reject) => {
+            timeout = setTimeout(() => {
+                delete this.queue[correlationId];
+                reject(new ConnectionError(`${apiName} timed out`));
+            }, 30_000);
+
             try {
-                await this.write(requestEncoder.value());
                 this.queue[correlationId] = { resolve, reject };
+                await this.write(requestEncoder.value());
             } catch (error) {
                 reject(error);
             }
         });
-        const response = api.response(responseDecoder);
+        clearTimeout(timeout);
+        const response = await api.response(responseDecoder);
 
         assert(
             responseDecoder.getOffset() - 4 === responseSize,
@@ -115,12 +125,13 @@ export class Connection {
     }
 
     private handleData(buffer: Buffer) {
-        this.buffer = this.buffer ? Buffer.concat([this.buffer, buffer]) : buffer;
-        if (this.buffer.length < 4) {
+        this.chunks.push(buffer);
+
+        const decoder = new Decoder(Buffer.concat(this.chunks));
+        if (decoder.getBufferLength() < 4) {
             return;
         }
 
-        const decoder = new Decoder(this.buffer);
         const size = decoder.readInt32();
         if (size !== decoder.getBufferLength() - 4) {
             return;
@@ -128,15 +139,18 @@ export class Connection {
 
         const correlationId = decoder.readInt32();
 
-        const { resolve } = this.queue[correlationId];
-        delete this.queue[correlationId];
-
-        resolve({ responseDecoder: decoder, responseSize: size });
-        this.buffer = null;
+        const context = this.queue[correlationId];
+        if (context) {
+            delete this.queue[correlationId];
+            context.resolve({ responseDecoder: decoder, responseSize: size });
+        } else {
+            log.debug('Could not find pending request for correlationId', { correlationId });
+        }
+        this.chunks = [];
     }
 
     private nextCorrelationId() {
-        return (this.lastCorrelationId = (this.lastCorrelationId + 1) % 2 ** 31);
+        return this.lastCorrelationId++;
     }
 }