Skip to content

Commit 548a597

Browse files
mish-ellen1ru4l
mish-elle
and
n1ru4l
authored
feat: support for MSK AWS IAM auth (#8049)
Co-authored-by: Laurin Quast <laurinquast@googlemail.com> Co-authored-by: Michelle Song <mish-elle@users.noreply.github.com>
1 parent 8d2909f commit 548a597

12 files changed

Lines changed: 609 additions & 37 deletions

File tree

‎.changeset/silly-flowers-yell.md‎

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,27 @@
1+
---
2+
'hive': minor
3+
---
4+
5+
Add opt-in AWS IAM authentication for MSK (Kafka) connections. When enabled, services authenticate
6+
to Kafka using AWS IAM (SigV4) via the OAUTHBEARER SASL mechanism.
7+
8+
### New environment variables
9+
10+
| Variable | Services | Description |
11+
| ---------------------------- | --------------------- | ------------------------------------------------------------------------- |
12+
| `AWS_REGION` | usage, usage-ingestor | Default AWS region for the service for all AWS connections. |
13+
| `KAFKA_AWS_IAM_AUTH_ENABLED` | usage, usage-ingestor | Set to `1` to enable IAM authentication. |
14+
| `KAFKA_AWS_REGION` | usage, usage-ingestor | Optional override for the Kafka broker region (defaults to `AWS_REGION`). |
15+
16+
### To enable
17+
18+
- `KAFKA_AWS_IAM_AUTH_ENABLED=1`
19+
- `KAFKA_SSL=1` must be set (IAM authentication requires TLS).
20+
- `KAFKA_AWS_REGION` or `AWS_REGION` must be set.
21+
- The pod/instance must have AWS credentials available (e.g. IRSA, EKS Pod Identity, instance
22+
profile) with the appropriate MSK IAM permissions.
23+
24+
### Other changes
25+
26+
- `KAFKA_BROKER` now accepts a comma-separated list of broker addresses (e.g.
27+
`broker1:9092,broker2:9092,broker3:9092`).

‎packages/services/service-common/package.json‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,7 @@
3030
"@sentry/node": "7.120.2",
3131
"@sentry/types": "7.120.2",
3232
"@sentry/utils": "7.120.2",
33+
"aws-msk-iam-sasl-signer-js": "1.0.3",
3334
"fastify": "5.8.5",
3435
"fastify-plugin": "5.1.0",
3536
"opentelemetry-instrumentation-fetch-node": "1.2.3",

‎packages/services/service-common/src/iam-msk.spec.ts‎

Lines changed: 40 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,40 @@
1+
import { generateAuthToken } from 'aws-msk-iam-sasl-signer-js';
2+
import { createMskIamTokenProvider } from './iam-msk';
3+
4+
vi.mock('aws-msk-iam-sasl-signer-js', () => ({
5+
generateAuthToken: vi.fn(async () => ({ token: 'abc' })),
6+
}));
7+
8+
const mockedGenerateAuthToken = vi.mocked(generateAuthToken);
9+
10+
test('returns kafkajs-compatible oauth bearer provider', async () => {
11+
const provider = createMskIamTokenProvider('us-east-1');
12+
await expect(provider()).resolves.toEqual({ value: 'abc' });
13+
});
14+
15+
test('passes the region to generateAuthToken', async () => {
16+
const provider = createMskIamTokenProvider('eu-west-1');
17+
await provider();
18+
expect(mockedGenerateAuthToken).toHaveBeenCalledWith({ region: 'eu-west-1' });
19+
});
20+
21+
test('propagates errors from generateAuthToken', async () => {
22+
mockedGenerateAuthToken.mockRejectedValueOnce(new Error('credentials not found'));
23+
const provider = createMskIamTokenProvider('us-east-1');
24+
await expect(provider()).rejects.toThrow('credentials not found');
25+
});
26+
27+
test('returns undefined token when generateAuthToken returns undefined', async () => {
28+
mockedGenerateAuthToken.mockResolvedValueOnce({ token: undefined } as any);
29+
const provider = createMskIamTokenProvider('us-east-1');
30+
await expect(provider()).resolves.toEqual({ value: undefined });
31+
});
32+
33+
test('can be called multiple times (token refresh)', async () => {
34+
mockedGenerateAuthToken
35+
.mockResolvedValueOnce({ token: 'token-1' } as any)
36+
.mockResolvedValueOnce({ token: 'token-2' } as any);
37+
const provider = createMskIamTokenProvider('us-east-1');
38+
await expect(provider()).resolves.toEqual({ value: 'token-1' });
39+
await expect(provider()).resolves.toEqual({ value: 'token-2' });
40+
});

‎packages/services/service-common/src/iam-msk.ts‎

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,18 @@
1+
import { generateAuthToken } from 'aws-msk-iam-sasl-signer-js';
2+
3+
/**
4+
* AWS MSK IAM SASL token provider for KafkaJS
5+
* Generates OAuth bearer tokens for MSK IAM authentication
6+
*/
7+
8+
/**
9+
* Creates an OAuth bearer token provider for AWS MSK IAM authentication
10+
* @param region AWS region (e.g., 'us-east-1')
11+
* @returns Async function that returns { value: token }
12+
*/
13+
export function createMskIamTokenProvider(region: string) {
14+
return async () => {
15+
const token = await generateAuthToken({ region });
16+
return { value: token.token };
17+
};
18+
}

‎packages/services/service-common/src/index.ts‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,3 +10,4 @@ export { registerShutdown } from './graceful-shutdown';
1010
export { cleanRequestId, maskToken } from './helpers';
1111
export { sentryInit } from './sentry';
1212
export { scrubBasicAuth } from './scrub';
13+
export { createMskIamTokenProvider } from './iam-msk';

‎packages/services/usage-ingestor/README.md‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,7 @@ This service takes care of feeding usage data into the ClickHouse instance.
88
| ----------------------------------- | ---------------------------------------------- | ---------------------------------------------------------------------------------------------------------------------------------------------------- | ---------------------------------------------------- |
99
| `ENVIRONMENT` | No | The environment of your Hive app. (**Note:** This will be used for Sentry reporting.) | `staging` |
1010
| `KAFKA_TOPIC` | **Yes** | The kafka topic. | `usage_reports_v2` |
11-
| `KAFKA_BROKER` | **Yes** | The address of the Kafka broker. | `127.0.0.1:29092` |
11+
| `KAFKA_BROKER` | **Yes** | The address of the Kafka broker. To use multiple, separate the addresses with commas. | `127.0.0.1:29092` |
1212
| `KAFKA_CONCURRENCY` | **Yes** | The concurrency of the Kafka connection. | `3` |
1313
| `KAFKA_SSL` | No | Whether an SSL connection should be established to the kafka service. | `1` (enabled) or `0` (disabled) |
1414
| `KAFKA_SSL_CA_PATH` | No | Refer to [TLS create secure context](https://nodejs.org/dist/latest-v8.x/docs/api/tls.html#tls_tls_createsecurecontext_options) for more information | `./path_to_ca` |
@@ -17,6 +17,9 @@ This service takes care of feeding usage data into the ClickHouse instance.
1717
| `KAFKA_SASL_MECHANISM` | No | The mechanism used for doing SASL authentication | `plain` or `scram-sha-256` or `scram-sha-512` |
1818
| `KAFKA_SASL_USERNAME` | No (Yes, if `KAFKA_SASL_MECHANISM` is defined) | The username for the SASL authentication | `letmein` |
1919
| `KAFKA_SASL_PASSWORD` | No (Yes, if `KAFKA_SASL_MECHANISM` is defined) | Whether an SSL connection should be established to the kafka service. | `letmein` |
20+
| `AWS_REGION` | No | The global AWS region for the service. Used as the default region for AWS connections. | `us-east-1` |
21+
| `KAFKA_AWS_IAM_AUTH_ENABLED` | No | Use AWS IAM to authenticate to the AWS MSK instance. Requires `AWS_REGION` to be set. | `1` (enabled) or `0` (disabled) |
22+
| `KAFKA_AWS_REGION` | No | Override for `AWS_REGION` specific to the Kafka/MSK connection. | `us-east-2` |
2023
| `CLICKHOUSE_PROTOCOL` | **Yes** | The ClickHouse protocol. | `http` or `https` |
2124
| `CLICKHOUSE_HOST` | **Yes** | The ClickHouse host. | `127.0.0.1` |
2225
| `CLICKHOUSE_PORT` | **Yes** | The ClickHouse port. | `8443` |

‎packages/services/usage-ingestor/src/environment.ts‎

Lines changed: 37 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -19,13 +19,18 @@ const emptyString = <T extends zod.ZodType>(input: T) => {
1919
}, input);
2020
};
2121

22+
function raiseInvariant(reason: string): never {
23+
throw new Error(reason);
24+
}
25+
2226
const EnvironmentModel = zod.object({
2327
PORT: emptyString(NumberFromString.optional()),
2428
SERVER_HOST: emptyString(zod.string().optional()),
2529
SERVER_HOST_IPV6_ONLY: emptyString(zod.union([zod.literal('1'), zod.literal('0')]).optional()),
2630
ENVIRONMENT: emptyString(zod.string().optional()),
2731
RELEASE: emptyString(zod.string().optional()),
2832
HEARTBEAT_ENDPOINT: emptyString(zod.string().url().optional()),
33+
AWS_REGION: emptyString(zod.string().optional()),
2934
});
3035

3136
const SentryModel = zod.union([
@@ -47,6 +52,10 @@ const KafkaBaseModel = zod.object({
4752
KAFKA_CONCURRENCY: NumberFromString,
4853
KAFKA_CONSUMER_GROUP: zod.string(),
4954
KAFKA_TOPIC: zod.string(),
55+
KAFKA_AWS_REGION: emptyString(zod.string().optional()),
56+
KAFKA_AWS_IAM_AUTH_ENABLED: emptyString(
57+
zod.union([zod.literal('0'), zod.literal('1')]).optional(),
58+
),
5059
});
5160

5261
const KafkaModel = zod.union([
@@ -139,6 +148,19 @@ for (const config of Object.values(configs)) {
139148
}
140149
}
141150

151+
if (configs.kafka.success && configs.kafka.data.KAFKA_AWS_IAM_AUTH_ENABLED === '1') {
152+
const missingKafkaIamVars: string[] = [];
153+
if (configs.kafka.data.KAFKA_SSL !== '1')
154+
missingKafkaIamVars.push('KAFKA_SSL must be enabled (MSK IAM requires TLS)');
155+
if (!configs.kafka.data.KAFKA_AWS_REGION && !configs.base.data?.AWS_REGION)
156+
missingKafkaIamVars.push('KAFKA_AWS_REGION or AWS_REGION');
157+
if (missingKafkaIamVars.length > 0) {
158+
environmentErrors.push(
159+
`KAFKA_AWS_IAM_AUTH_ENABLED is enabled but the following required variables are missing or invalid: ${missingKafkaIamVars.join(', ')}`,
160+
);
161+
}
162+
}
163+
142164
if (environmentErrors.length) {
143165
const fullError = environmentErrors.join(`\n`);
144166
console.error('❌ Invalid environment variables:', fullError);
@@ -189,13 +211,23 @@ export const env = {
189211
: true
190212
: false,
191213
sasl:
192-
kafka.KAFKA_SASL_MECHANISM != null
214+
kafka.KAFKA_AWS_IAM_AUTH_ENABLED === '1'
193215
? {
194-
mechanism: kafka.KAFKA_SASL_MECHANISM,
195-
username: kafka.KAFKA_SASL_USERNAME,
196-
password: kafka.KAFKA_SASL_PASSWORD,
216+
mechanism: 'aws-iam' as const,
217+
region:
218+
kafka.KAFKA_AWS_REGION ??
219+
base.AWS_REGION ??
220+
raiseInvariant(
221+
'KAFKA_AWS_REGION or AWS_REGION must be set when KAFKA_AWS_IAM_AUTH_ENABLED is enabled',
222+
),
197223
}
198-
: null,
224+
: kafka.KAFKA_SASL_MECHANISM != null
225+
? {
226+
mechanism: kafka.KAFKA_SASL_MECHANISM,
227+
username: kafka.KAFKA_SASL_USERNAME,
228+
password: kafka.KAFKA_SASL_PASSWORD,
229+
}
230+
: null,
199231
},
200232
},
201233
clickhouse: {

‎packages/services/usage-ingestor/src/ingestor.ts‎

Lines changed: 11 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
import { Kafka, KafkaMessage, logLevel } from 'kafkajs';
22
import type { ServiceLogger } from '@hive/service-common';
3+
import { createMskIamTokenProvider } from '@hive/service-common';
34
import type { RawReport } from '@hive/usage-common';
45
import { decompress } from '@hive/usage-common';
56
import type { KafkaEnvironment } from './environment';
@@ -51,7 +52,7 @@ export function createIngestor(config: {
5152

5253
const kafka = new Kafka({
5354
clientId: 'usage-ingestor',
54-
brokers: [config.kafka.connection.broker],
55+
brokers: config.kafka.connection.broker.split(',').map(b => b.trim()),
5556
ssl: config.kafka.connection.ssl,
5657
sasl:
5758
config.kafka.connection.sasl?.mechanism === 'plain'
@@ -72,7 +73,14 @@ export function createIngestor(config: {
7273
username: config.kafka.connection.sasl.username,
7374
password: config.kafka.connection.sasl.password,
7475
}
75-
: undefined,
76+
: config.kafka.connection.sasl?.mechanism === 'aws-iam'
77+
? {
78+
mechanism: 'oauthbearer',
79+
oauthBearerProvider: createMskIamTokenProvider(
80+
config.kafka.connection.sasl.region,
81+
),
82+
}
83+
: undefined,
7684
logLevel: logLevel.INFO,
7785
logCreator() {
7886
return entry => {
@@ -141,6 +149,7 @@ export function createIngestor(config: {
141149
changeStatus(Status.Waiting);
142150

143151
logger.info('Connecting Kafka Consumer');
152+
logger.debug(`Kafka SASL mechanism: ${config.kafka.connection.sasl?.mechanism ?? 'none'}`);
144153
await consumer.connect();
145154

146155
changeStatus(Status.Connected);

‎packages/services/usage/README.md‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -14,14 +14,17 @@ The data is written to a Kafka broker, form Kafka the data is feed into clickhou
1414
| `COMMERCE_ENDPOINT` | No | The endpoint of the commerce service. | `http://127.0.0.1:4012` |
1515
| `KAFKA_TOPIC` | **Yes** | The kafka topic. | `usage_reports_v2` |
1616
| `KAFKA_CONSUMER_GROUP` | **Yes** | The kafka consumer group. | `usage_reports_v2` |
17-
| `KAFKA_BROKER` | **Yes** | The address of the Kafka broker. | `127.0.0.1:29092` |
17+
| `KAFKA_BROKER` | **Yes** | The address of the Kafka broker. To use multiple, separate the addresses with commas. | `127.0.0.1:29092` |
1818
| `KAFKA_SSL` | No | Whether an SSL connection should be established to the kafka service. | `1` (enabled) or `0` (disabled) |
1919
| `KAFKA_SSL_CA_PATH` | No | Refer to [TLS create secure context](https://nodejs.org/dist/latest-v8.x/docs/api/tls.html#tls_tls_createsecurecontext_options) for more information | `./path_to_ca` |
2020
| `KAFKA_SSL_CERT_PATH` | No | Refer to [TLS create secure context](https://nodejs.org/dist/latest-v8.x/docs/api/tls.html#tls_tls_createsecurecontext_options) for more information | `./path_to_cert` |
2121
| `KAFKA_SSL_KEY_PATH` | No | Refer to [TLS create secure context](https://nodejs.org/dist/latest-v8.x/docs/api/tls.html#tls_tls_createsecurecontext_options) for more information | `./path_to_key` |
2222
| `KAFKA_SASL_MECHANISM` | No | The mechanism used for doing SASL authentication | `plain` or `scram-sha-256` or `scram-sha-512` |
2323
| `KAFKA_SASL_USERNAME` | No | The username for the SASL authentication | `letmein` |
2424
| `KAFKA_SASL_PASSWORD` | No | Whether an SSL connection should be established to the kafka service. | `letmein` |
25+
| `AWS_REGION` | No | The global AWS region for the service. Used as the default region for AWS connections. | `us-east-1` |
26+
| `KAFKA_AWS_IAM_AUTH_ENABLED` | No | Use AWS IAM to authenticate to the Kafka instance. Requires `AWS_REGION` to be set. | `1` (enabled) or `0` (disabled) |
27+
| `KAFKA_AWS_REGION` | No | Override for `AWS_REGION` specific to the Kafka/MSK connection. | `us-east-2` |
2528
| `KAFKA_BUFFER_SIZE` | No | The buffer size ??? | `12` |
2629
| `KAFKA_BUFFER_INTERVAL` | No | The buffer interval ??? | `1` |
2730
| `KAFKA_BUFFER_DYNAMIC` | No | The buffer interval ??? | `1` |

0 commit comments

Comments
 (0)