serverless-ircd 0.9.0 → 0.11.0

package/packages/aws-adapter/tests/nlb-stream.test.ts

@@ -9,7 +9,7 @@
  */
 
 import { CreateTableCommand, DeleteTableCommand } from '@aws-sdk/client-dynamodb';
-import { GetCommand } from '@aws-sdk/lib-dynamodb';
+import { type DynamoDBDocumentClient, GetCommand } from '@aws-sdk/lib-dynamodb';
 import { type ParsedServerConfig, StaticMotdProvider } from '@serverless-ircd/irc-core';
 import { makeTestServerConfig } from '@serverless-ircd/irc-test-support';
 import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
@@ -19,8 +19,10 @@ import { TABLE_DEFS } from '../src/cdk-table-defs.js';
 import { createDynamoDocumentClient } from '../src/dynamo.js';
 import {
   type NlbStreamEvent,
+  type NlbStreamParams,
   deriveNlbConnectionId,
   handleNlbStream,
+  parseNlbFlowHeaders,
 } from '../src/handlers/nlb-stream.js';
 import { bindMessageStore } from '../src/message-store.js';
 import type { TablesConfig } from '../src/tables.js';
@@ -61,7 +63,7 @@ async function makeFixture(): Promise<Fixture> {
     const pkName = props.partitionKey?.name;
     const skName = props.sortKey?.name;
     if (pkName === undefined) throw new Error(`table ${logical} missing partitionKey`);
-    const attributeDefinitions: Array<{ AttributeName: string; AttributeType: 'S' }> = [
+    const attributeDefinitions: Array<{ AttributeName: string; AttributeType: 'S' | 'N' }> = [
       { AttributeName: pkName, AttributeType: 'S' },
     ];
     const keySchema: Array<{ AttributeName: string; KeyType: 'HASH' | 'RANGE' }> = [
@@ -71,11 +73,35 @@ async function makeFixture(): Promise<Fixture> {
       attributeDefinitions.push({ AttributeName: skName, AttributeType: 'S' });
       keySchema.push({ AttributeName: skName, KeyType: 'RANGE' });
     }
+    // Secondary indexes declared by the CDK definition (the Connections
+    // per-IP admission GSI) are created here too — the NLB new-flow rate
+    // gate Queries the index, so a fixture without it diverges from the
+    // deployed schema.
+    const globalSecondaryIndexes = (props.globalSecondaryIndexes ?? []).map((g) => {
+      const gpk = g.partitionKey;
+      if (gpk === undefined) throw new Error(`index ${g.indexName} missing partitionKey`);
+      const gsiKeySchema: Array<{ AttributeName: string; KeyType: 'HASH' | 'RANGE' }> = [
+        { AttributeName: gpk.name, KeyType: 'HASH' },
+      ];
+      attributeDefinitions.push({ AttributeName: gpk.name, AttributeType: 'S' });
+      if (g.sortKey !== undefined) {
+        gsiKeySchema.push({ AttributeName: g.sortKey.name, KeyType: 'RANGE' });
+        attributeDefinitions.push({ AttributeName: g.sortKey.name, AttributeType: 'N' });
+      }
+      return {
+        IndexName: g.indexName,
+        KeySchema: gsiKeySchema,
+        Projection: { ProjectionType: 'ALL' as const },
+      };
+    });
     await client.send(
       new CreateTableCommand({
         TableName: `${prefix}${logical}`,
         AttributeDefinitions: attributeDefinitions,
         KeySchema: keySchema,
+        ...(globalSecondaryIndexes.length > 0
+          ? { GlobalSecondaryIndexes: globalSecondaryIndexes }
+          : {}),
         BillingMode: 'PAY_PER_REQUEST',
       }),
     );
@@ -130,6 +156,184 @@ describe.skipIf(!available)('deriveNlbConnectionId', () => {
   });
 });
 
+describe('parseNlbFlowHeaders', () => {
+  it('parses a well-formed IPv4 flow', () => {
+    expect(parseNlbFlowHeaders(nlbEventWithHeaders(headersFor('10.0.0.1', '50001')))).toEqual({
+      sourceIp: '10.0.0.1',
+      sourcePort: '50001',
+    });
+  });
+
+  it('parses a well-formed IPv6 flow', () => {
+    expect(parseNlbFlowHeaders(nlbEventWithHeaders(headersFor('2001:db8::1', '50002')))).toEqual({
+      sourceIp: '2001:db8::1',
+      sourcePort: '50002',
+    });
+  });
+
+  it('matches header names case-insensitively', () => {
+    expect(
+      parseNlbFlowHeaders(
+        nlbEventWithHeaders({ 'X-Forwarded-For': '10.0.0.2', 'X-Forwarded-Port': '50003' }),
+      ),
+    ).toEqual({ sourceIp: '10.0.0.2', sourcePort: '50003' });
+  });
+
+  it('accepts the port boundaries 1 and 65535', () => {
+    expect(parseNlbFlowHeaders(nlbEventWithHeaders(headersFor('10.0.0.3', '1')))).not.toBeNull();
+    expect(
+      parseNlbFlowHeaders(nlbEventWithHeaders(headersFor('10.0.0.3', '65535'))),
+    ).not.toBeNull();
+  });
+
+  it('rejects port 0 and ports above 65535', () => {
+    expect(parseNlbFlowHeaders(nlbEventWithHeaders(headersFor('10.0.0.4', '0')))).toBeNull();
+    expect(parseNlbFlowHeaders(nlbEventWithHeaders(headersFor('10.0.0.4', '65536')))).toBeNull();
+  });
+
+  it('rejects empty header values', () => {
+    expect(
+      parseNlbFlowHeaders(nlbEventWithHeaders({ 'x-forwarded-for': '', 'x-forwarded-port': '1' })),
+    ).toBeNull();
+    expect(
+      parseNlbFlowHeaders(
+        nlbEventWithHeaders({ 'x-forwarded-for': '10.0.0.5', 'x-forwarded-port': '' }),
+      ),
+    ).toBeNull();
+  });
+
+  it('rejects an out-of-range IPv4 octet', () => {
+    expect(parseNlbFlowHeaders(nlbEventWithHeaders(headersFor('10.0.0.256', '50005')))).toBeNull();
+  });
+
+  it('rejects a port with non-digit characters', () => {
+    expect(parseNlbFlowHeaders(nlbEventWithHeaders(headersFor('10.0.0.6', '12a')))).toBeNull();
+    expect(parseNlbFlowHeaders(nlbEventWithHeaders(headersFor('10.0.0.6', '+12')))).toBeNull();
+  });
+});
+
+/** Builds the canonical NLB flow header pair. */
+function headersFor(sourceIp: string, sourcePort: string): Record<string, string> {
+  return { 'x-forwarded-for': sourceIp, 'x-forwarded-port': sourcePort };
+}
+
+// ---------------------------------------------------------------------------
+// Flow header validation: a flow the NLB cannot identify (missing or
+// malformed x-forwarded-for / x-forwarded-port) must be rejected with a
+// 400 before any DynamoDB read, row creation, or actor invocation —
+// otherwise every such client collapses onto the single literal
+// connection id `nlb-unknown-0` and the flows overwrite each other's
+// persisted ConnectionState per frame. These rejection paths never touch
+// DynamoDB, so they run without the DDB Local fixture.
+// ---------------------------------------------------------------------------
+
+describe('NLB stream handler — flow header validation', () => {
+  it('rejects a flow missing x-forwarded-for with 400 before any DynamoDB I/O', async () => {
+    const { params, send } = rejectParams();
+    const res = await handleNlbStream(nlbEventWithHeaders({ 'x-forwarded-port': '50001' }), params);
+    expect(res.statusCode).toBe(400);
+    expect(send).not.toHaveBeenCalled();
+  });
+
+  it('rejects a flow missing x-forwarded-port with 400 before any DynamoDB I/O', async () => {
+    const { params, send } = rejectParams();
+    const res = await handleNlbStream(
+      nlbEventWithHeaders({ 'x-forwarded-for': '10.0.0.1' }),
+      params,
+    );
+    expect(res.statusCode).toBe(400);
+    expect(send).not.toHaveBeenCalled();
+  });
+
+  it('rejects a flow with no headers object at all with 400', async () => {
+    const { params, send } = rejectParams();
+    const event: NlbStreamEvent = {
+      requestContext: {
+        elb: { targetGroupArn: 'arn:aws:elasticloadbalancing:::targetgroup/test' },
+      },
+      version: '2.0',
+      body: Buffer.from('PING :tok\r\n', 'utf8').toString('base64'),
+      isBase64Encoded: true,
+    };
+    const res = await handleNlbStream(event, params);
+    expect(res.statusCode).toBe(400);
+    expect(send).not.toHaveBeenCalled();
+  });
+
+  it('rejects a flow with a malformed (non-IP) x-forwarded-for with 400', async () => {
+    const { params, send } = rejectParams();
+    const res = await handleNlbStream(
+      nlbEventWithHeaders({ 'x-forwarded-for': 'not-an-ip', 'x-forwarded-port': '50001' }),
+      params,
+    );
+    expect(res.statusCode).toBe(400);
+    expect(send).not.toHaveBeenCalled();
+  });
+
+  it('rejects a comma-separated x-forwarded-for list with 400', async () => {
+    // NLB sends exactly one client IP; a list cannot identify a flow.
+    const { params, send } = rejectParams();
+    const res = await handleNlbStream(
+      nlbEventWithHeaders({
+        'x-forwarded-for': '10.0.0.1, 192.168.0.1',
+        'x-forwarded-port': '50001',
+      }),
+      params,
+    );
+    expect(res.statusCode).toBe(400);
+    expect(send).not.toHaveBeenCalled();
+  });
+
+  it('rejects a flow with a malformed (non-numeric) x-forwarded-port with 400', async () => {
+    const { params, send } = rejectParams();
+    const res = await handleNlbStream(
+      nlbEventWithHeaders({ 'x-forwarded-for': '10.0.0.1', 'x-forwarded-port': 'port' }),
+      params,
+    );
+    expect(res.statusCode).toBe(400);
+    expect(send).not.toHaveBeenCalled();
+  });
+
+  it('never produces the literal nlb-unknown-0 connection id for a happy-path flow', async () => {
+    // A well-formed flow derives the expected id — the shared-row hazard
+    // (`nlb-unknown-0`) is unreachable by construction: every input that
+    // could produce it is rejected above.
+    expect(deriveNlbConnectionId('10.0.0.1', '50001')).toBe('nlb-10-0-0-1-50001');
+  });
+});
+
+/** Builds an NLB event with arbitrary headers (base64-encodes the body). */
+function nlbEventWithHeaders(headers: Record<string, string>): NlbStreamEvent {
+  return {
+    requestContext: { elb: { targetGroupArn: 'arn:aws:elasticloadbalancing:::targetgroup/test' } },
+    version: '2.0',
+    body: Buffer.from('PING :tok\r\n', 'utf8').toString('base64'),
+    isBase64Encoded: true,
+    headers,
+  };
+}
+
+/**
+ * Params double for rejection tests: a `send` spy that must never be
+ * invoked (rejections happen before any DynamoDB I/O).
+ */
+function rejectParams(): { params: NlbStreamParams; send: ReturnType<typeof vi.fn> } {
+  const send = vi.fn();
+  const params: NlbStreamParams = {
+    dynamo: { send } as unknown as DynamoDBDocumentClient,
+    tables: {
+      Connections: 'Connections',
+      ChannelMeta: 'ChannelMeta',
+      ChannelMembers: 'ChannelMembers',
+      Nicks: 'Nicks',
+    },
+    serverConfig: SERVER_CONFIG,
+    motd: MOTD,
+    managementApi: null,
+  };
+  return { params, send };
+}
+
 describe.skipIf(!available)('NLB stream handler — byte framing', () => {
   beforeEach(async () => {
     fx = await makeFixture();
@@ -515,10 +719,11 @@ describe.skipIf(!available)('NLB stream handler — streaming back', () => {
     expect(res.statusCode).toBe(200);
   });
 
-  it('falls back to "unknown"/"0" when NLB headers are absent', async () => {
-    // `readHeader` returns undefined when no header key matches. The
-    // caller's `?? 'unknown'` / `?? '0'` fallbacks must apply so the
-    // handler does not crash on a malformed event (defensive).
+  it('falls back to rejecting an event with unrelated headers only (no NLB flow metadata)', async () => {
+    // `readHeader` returns undefined when no header key matches; the
+    // handler must reject the flow with a 400 (see the flow-header
+    // validation block above) rather than admit it under a shared
+    // fallback id.
     const f = requireFx();
     const eventWithoutHeaders: NlbStreamEvent = {
       requestContext: {
@@ -529,6 +734,7 @@ describe.skipIf(!available)('NLB stream handler — streaming back', () => {
       isBase64Encoded: true,
       headers: { 'some-other-header': 'value' },
     };
+    const sendSpy = vi.spyOn(f.client, 'send');
     const res = await handleNlbStream(eventWithoutHeaders, {
       dynamo: f.client,
       tables: f.tables,
@@ -537,9 +743,9 @@ describe.skipIf(!available)('NLB stream handler — streaming back', () => {
       messages: bindMessageStore(SERVER_CONFIG),
       managementApi: f.mgmt,
     });
-    expect(res.statusCode).toBe(200);
-    // The fallback id (nlb-unknown-0) still received a PONG.
-    expect(res.body).toBeDefined();
+    expect(res.statusCode).toBe(400);
+    expect(sendSpy).not.toHaveBeenCalled();
+    sendSpy.mockRestore();
   });
 });
 
@@ -558,3 +764,416 @@ class LocalPostToConnection implements PostToConnection {
     return {};
   }
 }
+
+// ---------------------------------------------------------------------------
+// Input size limits: two independent guards.
+//   1. Payload pre-check — an event body decoding to more than 8 KiB is a
+//      misbehaving NLB integration (or an attacker); it is rejected with a
+//      400 BEFORE any DynamoDB read or actor invocation.
+//   2. Transport buffer cap — an unterminated tail that accumulates past
+//      8 KiB (possibly across invocations via the persisted buffer) closes
+//      the connection: the ERROR line is streamed back and the Connections
+//      row is dropped so the oversized buffer is never persisted.
+// ---------------------------------------------------------------------------
+
+describe.skipIf(!available)('NLB stream handler — input size limits', () => {
+  beforeEach(async () => {
+    fx = await makeFixture();
+  });
+  afterEach(async () => {
+    if (fx) await fx.cleanup();
+    fx = null;
+  });
+
+  it('rejects an oversized base64 event body with 400 before touching DynamoDB or the actor', async () => {
+    const f = requireFx();
+    const sendSpy = vi.spyOn(f.client, 'send');
+    const res = await handleNlbStream(nlbEvent('10.0.0.60', '50060', 'x'.repeat(9 * 1024)), {
+      dynamo: f.client,
+      tables: f.tables,
+      serverConfig: SERVER_CONFIG,
+      motd: MOTD,
+      messages: bindMessageStore(SERVER_CONFIG),
+      managementApi: f.mgmt,
+    });
+    expect(res.statusCode).toBe(400);
+    // Rejected before any bookkeeping: not a single DynamoDB command ran.
+    expect(sendSpy).not.toHaveBeenCalled();
+    sendSpy.mockRestore();
+  });
+
+  it('rejects an oversized plain (non-base64) event body with 400', async () => {
+    const f = requireFx();
+    const event: NlbStreamEvent = {
+      requestContext: {
+        elb: { targetGroupArn: 'arn:aws:elasticloadbalancing:::targetgroup/test' },
+      },
+      version: '2.0',
+      body: 'y'.repeat(9 * 1024),
+      isBase64Encoded: false,
+      headers: { 'x-forwarded-for': '10.0.0.61', 'x-forwarded-port': '50061' },
+    };
+    const sendSpy = vi.spyOn(f.client, 'send');
+    const res = await handleNlbStream(event, {
+      dynamo: f.client,
+      tables: f.tables,
+      serverConfig: SERVER_CONFIG,
+      motd: MOTD,
+      messages: bindMessageStore(SERVER_CONFIG),
+      managementApi: f.mgmt,
+    });
+    expect(res.statusCode).toBe(400);
+    expect(sendSpy).not.toHaveBeenCalled();
+    sendSpy.mockRestore();
+  });
+
+  it('accepts a payload exactly at the 8 KiB cap (regression guard on legal large chunks)', async () => {
+    const f = requireFx();
+    // 8191 bytes of a terminated line + CRLF = 8193 bytes... use exactly
+    // 8190 content bytes + CRLF so the decoded body is exactly 8192.
+    const body = `${'a'.repeat(8190)}\r\n`;
+    expect(Buffer.byteLength(body, 'utf8')).toBe(8192);
+    const res = await handleNlbStream(nlbEvent('10.0.0.62', '50062', body), {
+      dynamo: f.client,
+      tables: f.tables,
+      serverConfig: SERVER_CONFIG,
+      motd: MOTD,
+      messages: bindMessageStore(SERVER_CONFIG),
+      managementApi: f.mgmt,
+    });
+    expect(res.statusCode).toBe(200);
+  });
+
+  it('closes the connection with ERROR and drops the row when a chunk pushes the buffer past the cap', async () => {
+    const f = requireFx();
+    const opts = {
+      dynamo: f.client,
+      tables: f.tables,
+      serverConfig: SERVER_CONFIG,
+      motd: MOTD,
+      messages: bindMessageStore(SERVER_CONFIG),
+      managementApi: f.mgmt,
+    };
+    // Establish the connection with a 7 KiB unterminated tail (under both
+    // the payload pre-check and the buffer cap — buffered and persisted).
+    await handleNlbStream(nlbEvent('10.0.0.63', '50063', 'x'.repeat(7 * 1024)), opts);
+    // A 2 KiB chunk (itself under the pre-check cap) crosses the 8 KiB
+    // accumulation cap: the transport buffer overflows.
+    const res = await handleNlbStream(nlbEvent('10.0.0.63', '50063', 'x'.repeat(2 * 1024)), opts);
+    expect(res.statusCode).toBe(200);
+    expect(res.isBase64Encoded).toBe(true);
+    const decoded = Buffer.from(res.body, 'base64').toString('utf8');
+    expect(decoded).toContain('ERROR :Closing link: input buffer overflow');
+    // The Connections row is dropped entirely.
+    const row = await f.client.send(
+      new GetCommand({
+        TableName: f.tables.Connections,
+        Key: { connectionId: deriveNlbConnectionId('10.0.0.63', '50063') },
+      }),
+    );
+    expect(row.Item).toBeUndefined();
+  });
+
+  it('closes on overflow accumulated across invocations and never persists an oversized buffer', async () => {
+    const f = requireFx();
+    const opts = {
+      dynamo: f.client,
+      tables: f.tables,
+      serverConfig: SERVER_CONFIG,
+      motd: MOTD,
+      messages: bindMessageStore(SERVER_CONFIG),
+      managementApi: f.mgmt,
+    };
+    await handleNlbStream(nlbEvent('10.0.0.64', '50064', ''), opts);
+    // 8 KiB unterminated: at the cap — buffered and persisted, NOT closed.
+    await handleNlbStream(nlbEvent('10.0.0.64', '50064', 'a'.repeat(8192)), opts);
+    const midRow = await f.client.send(
+      new GetCommand({
+        TableName: f.tables.Connections,
+        Key: { connectionId: deriveNlbConnectionId('10.0.0.64', '50064') },
+      }),
+    );
+    expect(midRow.Item?.transportBuffer).toBe('a'.repeat(8192));
+
+    // One more byte-past-cap chunk: the restored buffer + chunk overflow.
+    const res = await handleNlbStream(nlbEvent('10.0.0.64', '50064', 'b'.repeat(1024)), opts);
+    const decoded = Buffer.from(res.body, 'base64').toString('utf8');
+    expect(decoded).toContain('ERROR :Closing link: input buffer overflow');
+    // The row is deleted — the oversized buffer never reaches a write.
+    const row = await f.client.send(
+      new GetCommand({
+        TableName: f.tables.Connections,
+        Key: { connectionId: deriveNlbConnectionId('10.0.0.64', '50064') },
+      }),
+    );
+    expect(row.Item).toBeUndefined();
+  });
+});
+
+// ---------------------------------------------------------------------------
+// Per-IP new-flow rate limit — stub-client unit tests (no DynamoDB Local
+// required, mirroring the connect.test.ts stub convention).
+// ---------------------------------------------------------------------------
+
+import { GetCommand as GetCmd, QueryCommand, UpdateCommand as UpdCmd } from '@aws-sdk/lib-dynamodb';
+import type { Clock } from '@serverless-ircd/irc-core';
+/** Fixed epoch base so window membership is deterministic. */
+const RATE_T0 = 1_700_000_000_000;
+
+/** A minimal existing-flow row (shape the NLB handler reads back). */
+function existingRow(connId: string): Record<string, unknown> {
+  return {
+    connectionId: connId,
+    registration: 'registered',
+    nick: 'existing',
+    capNegotiating: false,
+    secure: true,
+    caps: [],
+    userModes: { tls: true },
+    lastSeen: RATE_T0,
+    connectedSince: RATE_T0 - 5000,
+    idleSince: RATE_T0,
+    version: 1,
+    revision: 3,
+  };
+}
+
+interface NlbStub {
+  dynamo: DynamoDBDocumentClient;
+  /** QueryCommand inputs issued (rate-gate consultations). */
+  queries: Array<Record<string, unknown>>;
+  /** Whether the persist UpdateCommand ran (row written). */
+  persisted: boolean;
+}
+
+/**
+ * Stub client: `getRow` decides new-flow vs existing-flow; `gsiPages`
+ * scripts the sourceIp-GSI Query pages. Records queries + persists.
+ */
+function makeNlbStub(
+  getRow: Record<string, unknown> | undefined,
+  gsiPages: Array<Array<{ connectedSince: number }>> = [[]],
+): NlbStub {
+  const queries: Array<Record<string, unknown>> = [];
+  const flags = { persisted: false };
+  let queryCall = 0;
+  const dynamo = {
+    send: async (cmd: unknown) => {
+      if (cmd instanceof GetCmd) return getRow === undefined ? {} : { Item: getRow };
+      if (cmd instanceof QueryCommand) {
+        queries.push(cmd.input as Record<string, unknown>);
+        const page = gsiPages[queryCall] ?? [];
+        const last = queryCall;
+        queryCall += 1;
+        return {
+          Items: page,
+          LastEvaluatedKey: last + 1 < gsiPages.length ? { connectionId: `p-${last}` } : undefined,
+        };
+      }
+      if (cmd instanceof UpdCmd) {
+        flags.persisted = true;
+        return {};
+      }
+      throw new Error(`unexpected command: ${String(cmd)}`);
+    },
+  } as unknown as DynamoDBDocumentClient;
+  return {
+    dynamo,
+    get queries() {
+      return queries;
+    },
+    get persisted() {
+      return flags.persisted;
+    },
+  };
+}
+
+/** Fixed clock at `t` so the sliding-window cutoff is deterministic. */
+function clockAt(t: number): Clock {
+  return { now: () => t };
+}
+
+describe('NLB stream handler — per-IP new-flow rate limit', () => {
+  const RATE_CONFIG = makeTestServerConfig({
+    perIpConnectionRate: { max: 2, windowMs: 60_000 },
+  });
+
+  function opts(
+    stub: NlbStub,
+    clock: Clock = clockAt(RATE_T0),
+    config: typeof RATE_CONFIG = RATE_CONFIG,
+  ) {
+    return {
+      dynamo: stub.dynamo,
+      tables: {
+        Connections: 'Connections',
+        ChannelMeta: 'ChannelMeta',
+        ChannelMembers: 'ChannelMembers',
+        Nicks: 'Nicks',
+      } as TablesConfig,
+      serverConfig: config,
+      motd: MOTD,
+      managementApi: null,
+      clock,
+    };
+  }
+
+  it('admits a new flow under the budget and persists its row', async () => {
+    const stub = makeNlbStub(undefined, [[{ connectedSince: RATE_T0 - 1000 }]]);
+    const res = await handleNlbStream(nlbEvent('10.0.1.1', '60001', 'PING :tok\r\n'), opts(stub));
+    expect(res.statusCode).toBe(200);
+    expect(stub.persisted).toBe(true);
+  });
+
+  it('rejects the over-budget new flow with 429 and writes no row', async () => {
+    const stub = makeNlbStub(undefined, [
+      [{ connectedSince: RATE_T0 - 1000 }, { connectedSince: RATE_T0 - 2000 }],
+    ]);
+    const res = await handleNlbStream(nlbEvent('10.0.1.2', '60002', 'PING :tok\r\n'), opts(stub));
+    expect(res.statusCode).toBe(429);
+    expect(res.body).toBe('');
+    expect(res.isBase64Encoded).toBe(false);
+    expect(stub.persisted).toBe(false);
+    expect(stub.queries.length).toBe(1);
+    expect(stub.queries[0]?.IndexName).toBe('sourceIp-connectedSince');
+    expect(stub.queries[0]?.ExpressionAttributeValues).toMatchObject({ ':ip': '10.0.1.2' });
+  });
+
+  it('does not consult the rate budget for an existing flow (rate gates establishment only)', async () => {
+    // Rate-limiting EVERY chunk invocation would starve an active IRC
+    // session (each client line is a fresh invocation); the budget
+    // applies to connection establishment — the first chunk of a flow.
+    const connId = deriveNlbConnectionId('10.0.1.3', '60003');
+    const stub = makeNlbStub(existingRow(connId), [
+      [{ connectedSince: RATE_T0 - 1000 }, { connectedSince: RATE_T0 - 2000 }],
+    ]);
+    const res = await handleNlbStream(nlbEvent('10.0.1.3', '60003', 'PING :tok\r\n'), opts(stub));
+    expect(res.statusCode).toBe(200);
+    expect(stub.queries.length).toBe(0);
+    expect(stub.persisted).toBe(true);
+  });
+
+  it('window-reset: admissions older than the window stop consuming budget', async () => {
+    // Budget of 1; one prior admission stamped at RATE_T0.
+    const budgetOne = makeTestServerConfig({
+      perIpConnectionRate: { max: 1, windowMs: 60_000 },
+    });
+    const row = [{ connectedSince: RATE_T0 }];
+    // Inside the window (clock at RATE_T0+30 s): the admission counts →
+    // the new flow is over budget and rejected.
+    const early = makeNlbStub(undefined, [row]);
+    const earlyRes = await handleNlbStream(
+      nlbEvent('10.0.1.4', '60004', 'PING :tok\r\n'),
+      opts(early, clockAt(RATE_T0 + 30_000), budgetOne),
+    );
+    expect(earlyRes.statusCode).toBe(429);
+
+    // Past the window (clock at RATE_T0+61 s): the cutoff advanced past
+    // the admission → the budget recovered and the flow is admitted.
+    const later = makeNlbStub(undefined, [row]);
+    const laterRes = await handleNlbStream(
+      nlbEvent('10.0.1.4', '60005', 'PING :tok\r\n'),
+      opts(later, clockAt(RATE_T0 + 61_000), budgetOne),
+    );
+    expect(laterRes.statusCode).toBe(200);
+  });
+
+  it('paginates the GSI query when the per-IP partition spans pages', async () => {
+    const stub = makeNlbStub(undefined, [
+      [{ connectedSince: RATE_T0 - 1000 }, { connectedSince: RATE_T0 - 2000 }],
+      [{ connectedSince: RATE_T0 - 3000 }],
+    ]);
+    const res = await handleNlbStream(nlbEvent('10.0.1.5', '60006', 'PING :tok\r\n'), opts(stub));
+    expect(res.statusCode).toBe(429);
+    expect(stub.queries.length).toBe(2);
+  });
+
+  it('never consults the rate gate for a flow with missing flow headers', async () => {
+    // Malformed/missing flow headers are rejected by header validation
+    // (400) before any rate-gate consideration — and the stub's Get
+    // (existing-flow probe) and Query (rate gate) both stay silent.
+    const stub = makeNlbStub(undefined, [[]]);
+    const event: NlbStreamEvent = {
+      requestContext: { elb: { targetGroupArn: 'arn:aws:elasticloadbalancing:::tg' } },
+      version: '2.0',
+      body: Buffer.from('PING :tok\r\n', 'utf8').toString('base64'),
+      isBase64Encoded: true,
+      headers: { 'x-forwarded-port': '60007' },
+    };
+    const res = await handleNlbStream(event, opts(stub));
+    expect(res.statusCode).toBe(400);
+    expect(stub.queries.length).toBe(0);
+    expect(stub.persisted).toBe(false);
+  });
+});
+
+describe.skipIf(!available)('NLB per-IP new-flow rate limit (real GSI)', () => {
+  beforeEach(async () => {
+    fx = await makeFixture();
+  });
+  afterEach(async () => {
+    if (fx) await fx.cleanup();
+    fx = null;
+  });
+
+  /** Rate-pinned config: 2 new flows per 60 s window. */
+  const cfg = makeTestServerConfig({ perIpConnectionRate: { max: 2, windowMs: 60_000 } });
+
+  function flowOpts(clock: Clock) {
+    const f = requireFx();
+    return {
+      dynamo: f.client,
+      tables: f.tables,
+      serverConfig: cfg,
+      motd: MOTD,
+      messages: bindMessageStore(cfg),
+      managementApi: f.mgmt,
+      clock,
+    };
+  }
+
+  it('rejects the over-budget new flow with 429 (no row written)', async () => {
+    const t0 = { now: () => RATE_T0 };
+    const first = await handleNlbStream(nlbEvent('10.0.3.1', '61001', 'PING :a\r\n'), flowOpts(t0));
+    expect(first.statusCode).toBe(200);
+    const second = await handleNlbStream(
+      nlbEvent('10.0.3.1', '61002', 'PING :b\r\n'),
+      flowOpts(t0),
+    );
+    expect(second.statusCode).toBe(200);
+    const third = await handleNlbStream(nlbEvent('10.0.3.1', '61003', 'PING :c\r\n'), flowOpts(t0));
+    expect(third.statusCode).toBe(429);
+    expect(third.body).toBe('');
+    // No row for the rejected flow.
+    const row = await requireFx().client.send(
+      new GetCommand({
+        TableName: requireFx().tables.Connections,
+        Key: { connectionId: deriveNlbConnectionId('10.0.3.1', '61003') },
+      }),
+    );
+    expect(row.Item).toBeUndefined();
+  });
+
+  it('window-reset: the budget recovers once establishments age out', async () => {
+    const t0 = { now: () => RATE_T0 };
+    expect(
+      (await handleNlbStream(nlbEvent('10.0.3.2', '61011', 'PING :a\r\n'), flowOpts(t0)))
+        .statusCode,
+    ).toBe(200);
+    expect(
+      (await handleNlbStream(nlbEvent('10.0.3.2', '61012', 'PING :b\r\n'), flowOpts(t0)))
+        .statusCode,
+    ).toBe(200);
+    expect(
+      (await handleNlbStream(nlbEvent('10.0.3.2', '61013', 'PING :c\r\n'), flowOpts(t0)))
+        .statusCode,
+    ).toBe(429);
+    // Past the window both establishments decayed → the budget recovered.
+    const later = { now: () => RATE_T0 + 61_000 };
+    const after = await handleNlbStream(
+      nlbEvent('10.0.3.2', '61014', 'PING :d\r\n'),
+      flowOpts(later),
+    );
+    expect(after.statusCode).toBe(200);
+  });
+});