smpp-js/src/reassembly.ts

import type { Concat } from './concat.ts';
import type { PduObject } from './pdu.ts';
import type { SmppLog } from './log.ts';
import { ExpiringGroups } from './expiring-groups.ts';
import { decodeMessage } from './message.ts';
import { detach, retainedOctets } from './retained-pdu.ts';
import { messageOctets } from './message-body.ts';
import { paramNumber, paramText } from './defs/types.ts';
import { uuidv7 } from './uuid.ts';

/** A concatenated message given up on, whose segments the peer has already been answered for. */
export type LostGroup = {
	parts: number;
	reason: 'evicted' | 'expired' | 'linkGone';
	smsId: string;
	total: number;
};

export type ReassemblerOptions = {
	log: SmppLog;
	max: number;
	maxOctets?: number | undefined;
	/** Injected so the ids a group is answered with can be read back in a test. */
	newId?: (() => string) | undefined;
	/** Injected so expiry can be exercised without a wall clock. */
	now?: (() => number) | undefined;
	onLost: (lost: LostGroup) => void;
	timeout: number;
};

/** Why a segment was not kept: its header joins no message here, or the store had no room for it. */
export type Refusal = 'full' | 'unplaceable';

/** What a segment did to its group. */
export type Collected =
	| { kept: false; refusal: Refusal }
	| {
		kept: true;
		/** The id base the group's segments are answered with. */
		smsId: string;
		/** Every segment in order, on the one that completes the message. */
		whole?: PduObject[] | undefined;
	};

export const defaultMaxOctets = 64 * 1024 * 1024;

type Group = {
	parts: Map<number, PduObject>;
	smsId: string;
	total: number;
};

// NUL: the one octet a C-Octet String address cannot hold, so no sender can forge another's key.
function groupKey(pduObj: PduObject, concat: Concat): string {
	return [
		paramText(pduObj.params.source_addr),
		paramText(pduObj.params.destination_addr),
		concat.spelling,
		String(concat.reference),
	].join('\u0000');
}

/** The text of a message, joining its segments in the order they were reassembled. */
export function decodeSegments(pduObjs: PduObject[]): string {
	let message = '';

	for (const pduObj of pduObjs) {
		const part = messageOctets(pduObj);

		if (part === undefined) continue;

		message += decodeMessage(
			part,
			paramNumber(pduObj.params.data_coding, 0),
			paramNumber(pduObj.params.esm_class, 0),
		).message;
	}

	return message;
}

/** Holds the segments of incomplete multipart messages until they are whole, capped and expiring. */
export class Reassembler {
	private readonly groups: ExpiringGroups<Group>;
	private readonly log: SmppLog;
	private readonly max: number;
	private readonly maxOctets: number;
	private readonly newId: () => string;
	private readonly onLost: (lost: LostGroup) => void;

	constructor(options: ReassemblerOptions) {
		this.maxOctets = options.maxOctets ?? defaultMaxOctets;
		this.groups = new ExpiringGroups<Group>({
			max: options.max,
			maxWeight: this.maxOctets,
			now: options.now,
			onSweep: () => { this.sweep(); },
			timeout: options.timeout,
		});
		this.log = options.log;
		this.max = options.max;
		this.newId = options.newId ?? uuidv7;
		this.onLost = options.onLost;
	}

	get size(): number {
		return this.groups.size;
	}

	/** The group the segment joined, and always an answer for it: an unanswered one stalls a peer. */
	collect(pduObj: PduObject, concat: Concat): Collected {
		this.sweep();

		const key = groupKey(pduObj, concat);
		const existing = this.groups.get(key);

		if (!this.placeable(concat, existing)) return { kept: false, refusal: 'unplaceable' };

		const group = existing ?? this.open(key, concat.total);

		group.parts.set(concat.part, detach(pduObj));

		if (group.parts.size < group.total) {
			// Its own arrival overran the octet cap, so the peer keeps it rather than being told we did.
			if (!this.trim(key, group)) return { kept: false, refusal: 'full' };

			return { kept: true, smsId: group.smsId };
		}

		this.groups.delete(key);

		return {
			kept: true,
			smsId: group.smsId,
			whole: [...group.parts.entries()].sort(([a], [b]) => a - b).map(([, part]) => part),
		};
	}

	clear(): void {
		for (const [, group] of this.groups.takeAll()) {
			this.lost(group, 'linkGone');
		}
	}

	/** Drops every group past its deadline. Runs before each collect and on its own timer. */
	sweep(): void {
		for (const [, group] of this.groups.takeExpired()) {
			this.lost(group, 'expired');
		}
	}

	/** Whether a segment can join a group at all: its own numbering, and the group's total. */
	private placeable(concat: Concat, existing: Group | undefined): boolean {
		if (concat.part < 1 || concat.total < 1 || concat.part > concat.total) {
			this.log.warn('reassembler - dropping an impossibly numbered segment', {
				part: concat.part,
				total: concat.total,
			});

			return false;
		}

		// Parts 1/2 and 2/3 would otherwise complete the stored two-part group as a truncated message.
		if (existing && existing.total !== concat.total) {
			this.log.warn('reassembler - dropping a segment with an inconsistent total', {
				existingTotal: existing.total,
				part: concat.part,
				total: concat.total,
			});

			return false;
		}

		return true;
	}

	private open(key: string, total: number): Group {
		if (this.groups.full) this.dropOldest();

		const group: Group = { parts: new Map(), smsId: this.newId(), total };

		this.groups.set(key, group);

		return group;
	}

	/** Drops the oldest groups until the retained payload is back under the octet cap. False if the current one went. */
	private trim(current: string, group: Group): boolean {
		let octets = 0;

		for (const part of group.parts.values()) {
			octets += retainedOctets(part);
		}

		let survived = true;

		for (const [key, oldest] of this.groups.weigh(current, octets)) {
			if (key === current) survived = false;

			// The refused segment is in the group but stays with the peer, so it is none of the loss.
			const answered = key === current ? oldest.parts.size - 1 : oldest.parts.size;

			if (answered > 0) this.lost(oldest, 'evicted', answered);
		}

		return survived;
	}

	private dropOldest(): void {
		const oldest = this.groups.takeOldest();

		if (oldest) this.lost(oldest[1], 'evicted');
	}

	/** Its segments are answered, so the peer will not send them again: this is traffic gone. */
	private lost(group: Group, reason: LostGroup['reason'], parts = group.parts.size): void {
		const lost: LostGroup = { parts, reason, smsId: group.smsId, total: group.total };

		this.log.warn('reassembler - gave up a concatenated message', {
			...lost,
			max: this.max,
			maxOctets: this.maxOctets,
			octets: this.groups.weight,
		});
		this.onLost(lost);
	}
}