363 lines
10 KiB
TypeScript
363 lines
10 KiB
TypeScript
import * as bindings from "mtp/raw";
|
|
import { MTPRatchet } from "./ratchet.js";
|
|
import type { MTPSessionState } from "./session";
|
|
import { concatBytes, writeU64BE } from "./utils.js";
|
|
|
|
export const MTP_E2EE_VERSION = 1;
|
|
export const FLAG_INIT = 0x01;
|
|
export const FLAG_DEVICE_SECRET = 0x02;
|
|
export const FLAG_KEY_ROTATION = 0x04;
|
|
export const MAX_RATCHET_SKIP = 100;
|
|
const SUPPORTED_FLAGS = FLAG_INIT | FLAG_DEVICE_SECRET | FLAG_KEY_ROTATION;
|
|
const HEADER_FIXED_LEN = 1 + 1 + 8 + 8 + 4 + 2 + 4;
|
|
|
|
export interface ParsedEncryptedMessage {
|
|
version: 1;
|
|
flags: number;
|
|
senderClientId: bigint;
|
|
recipientClientId: bigint;
|
|
messageNumber: number;
|
|
kemCiphertext?: Uint8Array;
|
|
ciphertext: Uint8Array;
|
|
/** Compatibility alias for older SDK tests/callers. */
|
|
header?: EncryptedMessageHeader;
|
|
/** Compatibility alias for older SDK tests/callers. */
|
|
aeadPayload?: Uint8Array;
|
|
}
|
|
|
|
export interface EncryptedMessageHeader {
|
|
version: 1;
|
|
flags: number;
|
|
senderClientId: bigint;
|
|
recipientClientId: bigint;
|
|
messageNumber: number;
|
|
kemCiphertext?: Uint8Array;
|
|
}
|
|
export interface SerializedEncryptedMessage {
|
|
header: EncryptedMessageHeader;
|
|
aeadPayload: Uint8Array;
|
|
}
|
|
|
|
function readU64BE(bytes: Uint8Array, offset: number): bigint {
|
|
let value = 0n;
|
|
for (let i = 0; i < 8; i++) {
|
|
value = (value << 8n) | BigInt(bytes[offset + i]);
|
|
}
|
|
return value;
|
|
}
|
|
|
|
function writeU32BE(value: number): Uint8Array {
|
|
if (!Number.isSafeInteger(value) || value < 0 || value > 0xffff_ffff) {
|
|
throw new Error("u32 value out of range");
|
|
}
|
|
return new Uint8Array([
|
|
(value >>> 24) & 0xff,
|
|
(value >>> 16) & 0xff,
|
|
(value >>> 8) & 0xff,
|
|
value & 0xff,
|
|
]);
|
|
}
|
|
|
|
function assertSupported(message: ParsedEncryptedMessage): void {
|
|
if (message.version !== MTP_E2EE_VERSION) {
|
|
throw new Error(
|
|
`Unsupported encrypted message version: ${message.version}`,
|
|
);
|
|
}
|
|
if ((message.flags & ~SUPPORTED_FLAGS) !== 0) {
|
|
throw new Error(`Unsupported encrypted message flags: ${message.flags}`);
|
|
}
|
|
const isInit = (message.flags & FLAG_INIT) !== 0;
|
|
if (isInit && !message.kemCiphertext?.length) {
|
|
throw new Error("Init message must include KEM ciphertext");
|
|
}
|
|
if (!isInit && message.kemCiphertext?.length) {
|
|
throw new Error("Non-init message must not include KEM ciphertext");
|
|
}
|
|
if (!message.ciphertext.length) {
|
|
throw new Error("Encrypted message ciphertext must be non-empty");
|
|
}
|
|
}
|
|
|
|
export function serializeEncryptedMessage(
|
|
message: ParsedEncryptedMessage,
|
|
): Uint8Array;
|
|
export function serializeEncryptedMessage(
|
|
message: SerializedEncryptedMessage,
|
|
): Uint8Array;
|
|
export function serializeEncryptedMessage(
|
|
message: ParsedEncryptedMessage | SerializedEncryptedMessage,
|
|
): Uint8Array {
|
|
const normalized: ParsedEncryptedMessage =
|
|
"header" in message
|
|
? {
|
|
...message.header,
|
|
ciphertext: message.aeadPayload,
|
|
}
|
|
: message;
|
|
|
|
assertSupported(normalized);
|
|
const kemCiphertext = normalized.kemCiphertext ?? new Uint8Array(0);
|
|
if (kemCiphertext.length > 0xffff) {
|
|
throw new Error("KEM ciphertext too long");
|
|
}
|
|
|
|
return concatBytes([
|
|
new Uint8Array([normalized.version]),
|
|
new Uint8Array([normalized.flags]),
|
|
writeU64BE(normalized.senderClientId),
|
|
writeU64BE(normalized.recipientClientId),
|
|
writeU32BE(normalized.messageNumber),
|
|
new Uint8Array([
|
|
(kemCiphertext.length >>> 8) & 0xff,
|
|
kemCiphertext.length & 0xff,
|
|
]),
|
|
kemCiphertext,
|
|
writeU32BE(normalized.ciphertext.length),
|
|
normalized.ciphertext,
|
|
]);
|
|
}
|
|
|
|
export function parseEncryptedMessage(
|
|
bytes: Uint8Array,
|
|
): ParsedEncryptedMessage {
|
|
let offset = 0;
|
|
if (!(bytes instanceof Uint8Array)) {
|
|
throw new Error("Encrypted message must be bytes");
|
|
}
|
|
if (bytes.length < HEADER_FIXED_LEN) {
|
|
throw new Error("Encrypted message too short");
|
|
}
|
|
|
|
const version = bytes[offset++];
|
|
const flags = bytes[offset++];
|
|
const senderClientId = readU64BE(bytes, offset);
|
|
offset += 8;
|
|
const recipientClientId = readU64BE(bytes, offset);
|
|
offset += 8;
|
|
const messageNumber =
|
|
((bytes[offset] << 24) |
|
|
(bytes[offset + 1] << 16) |
|
|
(bytes[offset + 2] << 8) |
|
|
bytes[offset + 3]) >>>
|
|
0;
|
|
offset += 4;
|
|
const kemLen = (bytes[offset] << 8) | bytes[offset + 1];
|
|
offset += 2;
|
|
|
|
let kemCiphertext: Uint8Array | undefined;
|
|
if (kemLen > 0) {
|
|
if (bytes.length < offset + kemLen + 4) {
|
|
throw new Error("Encrypted message KEM ciphertext truncated");
|
|
}
|
|
kemCiphertext = bytes.slice(offset, offset + kemLen);
|
|
offset += kemLen;
|
|
}
|
|
|
|
if (bytes.length < offset + 4) {
|
|
throw new Error("Encrypted message missing ciphertext length");
|
|
}
|
|
const ciphertextLen =
|
|
((bytes[offset] << 24) |
|
|
(bytes[offset + 1] << 16) |
|
|
(bytes[offset + 2] << 8) |
|
|
bytes[offset + 3]) >>>
|
|
0;
|
|
offset += 4;
|
|
if (bytes.length < offset + ciphertextLen) {
|
|
throw new Error("Encrypted message ciphertext truncated");
|
|
}
|
|
const ciphertext = bytes.slice(offset, offset + ciphertextLen);
|
|
offset += ciphertextLen;
|
|
if (offset !== bytes.length) {
|
|
throw new Error("Encrypted message has trailing data");
|
|
}
|
|
|
|
const parsed: ParsedEncryptedMessage = {
|
|
version: version as 1,
|
|
flags,
|
|
senderClientId,
|
|
recipientClientId,
|
|
messageNumber,
|
|
kemCiphertext,
|
|
ciphertext,
|
|
};
|
|
parsed.header = {
|
|
version: parsed.version,
|
|
flags: parsed.flags,
|
|
senderClientId: parsed.senderClientId,
|
|
recipientClientId: parsed.recipientClientId,
|
|
messageNumber: parsed.messageNumber,
|
|
kemCiphertext: parsed.kemCiphertext,
|
|
};
|
|
parsed.aeadPayload = parsed.ciphertext;
|
|
assertSupported(parsed);
|
|
return parsed;
|
|
}
|
|
|
|
function buildAAD(header: EncryptedMessageHeader): Uint8Array {
|
|
return concatBytes([
|
|
new Uint8Array([header.version]),
|
|
new Uint8Array([header.flags]),
|
|
writeU64BE(header.senderClientId),
|
|
writeU64BE(header.recipientClientId),
|
|
writeU32BE(header.messageNumber),
|
|
]);
|
|
}
|
|
|
|
export function encryptedMessageAAD(
|
|
header: EncryptedMessageHeader,
|
|
extra?: Uint8Array,
|
|
): Uint8Array {
|
|
return extra?.length
|
|
? concatBytes([buildAAD(header), extra])
|
|
: buildAAD(header);
|
|
}
|
|
|
|
export async function encryptPayload(args: {
|
|
plaintext: Uint8Array;
|
|
session: MTPSessionState;
|
|
kemCiphertext?: Uint8Array;
|
|
aad?: Uint8Array;
|
|
}): Promise<{
|
|
payload: Uint8Array;
|
|
session: MTPSessionState;
|
|
}> {
|
|
const step = await MTPRatchet.stepSend(args.session.sendChainKey);
|
|
const header: EncryptedMessageHeader = {
|
|
version: 1,
|
|
flags: args.kemCiphertext ? FLAG_INIT : 0,
|
|
senderClientId: args.session.ownClientId,
|
|
recipientClientId: args.session.peerClientId,
|
|
messageNumber: args.session.sendCount,
|
|
kemCiphertext: args.kemCiphertext,
|
|
};
|
|
const aad = args.aad ?? encryptedMessageAAD(header);
|
|
const cipher = new bindings.WasmChaCha20Poly1305(step.key);
|
|
let ciphertext: Uint8Array;
|
|
try {
|
|
ciphertext = cipher.encrypt(args.plaintext, aad);
|
|
} finally {
|
|
cipher.free();
|
|
step.key.fill(0);
|
|
}
|
|
|
|
const payload = serializeEncryptedMessage({ ...header, ciphertext });
|
|
return {
|
|
payload,
|
|
session: {
|
|
...args.session,
|
|
sendChainKey: step.chainKey,
|
|
sendCount: args.session.sendCount + 1,
|
|
updatedAt: Date.now(),
|
|
},
|
|
};
|
|
}
|
|
|
|
export async function decryptPayload(args: {
|
|
payload: Uint8Array;
|
|
session: MTPSessionState;
|
|
expectedRecipientClientId?: bigint;
|
|
aad?: Uint8Array;
|
|
}): Promise<{
|
|
plaintext: Uint8Array;
|
|
session: MTPSessionState;
|
|
}> {
|
|
const parsed = parseEncryptedMessage(args.payload);
|
|
const expectedRecipientClientId =
|
|
args.expectedRecipientClientId ?? args.session.ownClientId;
|
|
if (parsed.recipientClientId !== expectedRecipientClientId) {
|
|
throw new Error("Encrypted message recipient mismatch");
|
|
}
|
|
if (parsed.senderClientId !== args.session.peerClientId) {
|
|
throw new Error("Encrypted message sender mismatch");
|
|
}
|
|
|
|
const existingSkippedMessageKeys = args.session.skippedMessageKeys ?? [];
|
|
const cachedKeyIndex = existingSkippedMessageKeys.findIndex(
|
|
(skipped) => skipped.messageNumber === parsed.messageNumber,
|
|
);
|
|
if (parsed.messageNumber < args.session.recvCount && cachedKeyIndex < 0) {
|
|
throw new Error("Encrypted message replay message number");
|
|
}
|
|
|
|
let chainKey = args.session.recvChainKey;
|
|
let messageKey: Uint8Array | undefined;
|
|
let skippedMessageKeys = existingSkippedMessageKeys.slice();
|
|
const newlyDerivedKeys: Uint8Array[] = [];
|
|
let nextRecvCount = args.session.recvCount;
|
|
|
|
if (cachedKeyIndex >= 0) {
|
|
// Work on a copy so an invalid ciphertext cannot consume the cached key.
|
|
messageKey = skippedMessageKeys[cachedKeyIndex].key.slice();
|
|
} else {
|
|
const gap = parsed.messageNumber - args.session.recvCount;
|
|
if (gap > MAX_RATCHET_SKIP) {
|
|
throw new Error(
|
|
`Encrypted message receive gap exceeds max skip (${MAX_RATCHET_SKIP})`,
|
|
);
|
|
}
|
|
|
|
const steps = gap + 1;
|
|
for (let i = 0; i < steps; i += 1) {
|
|
const step = await MTPRatchet.stepRecv(chainKey);
|
|
if (i === steps - 1) {
|
|
messageKey = step.key;
|
|
} else {
|
|
skippedMessageKeys.push({
|
|
messageNumber: args.session.recvCount + i,
|
|
key: step.key,
|
|
});
|
|
newlyDerivedKeys.push(step.key);
|
|
}
|
|
if (chainKey !== args.session.recvChainKey) chainKey.fill(0);
|
|
chainKey = step.chainKey;
|
|
}
|
|
nextRecvCount = parsed.messageNumber + 1;
|
|
}
|
|
if (!messageKey) {
|
|
throw new Error("Failed to derive receive message key");
|
|
}
|
|
|
|
const header: EncryptedMessageHeader = {
|
|
version: parsed.version,
|
|
flags: parsed.flags,
|
|
senderClientId: parsed.senderClientId,
|
|
recipientClientId: parsed.recipientClientId,
|
|
messageNumber: parsed.messageNumber,
|
|
kemCiphertext: parsed.kemCiphertext,
|
|
};
|
|
const aad = args.aad ?? encryptedMessageAAD(header);
|
|
const cipher = new bindings.WasmChaCha20Poly1305(messageKey);
|
|
let plaintext: Uint8Array;
|
|
try {
|
|
plaintext = cipher.decrypt(parsed.ciphertext, aad);
|
|
} catch (error) {
|
|
for (const key of newlyDerivedKeys) key.fill(0);
|
|
if (chainKey !== args.session.recvChainKey) chainKey.fill(0);
|
|
throw error;
|
|
} finally {
|
|
cipher.free();
|
|
messageKey.fill(0);
|
|
}
|
|
|
|
if (cachedKeyIndex >= 0) {
|
|
const [consumed] = skippedMessageKeys.splice(cachedKeyIndex, 1);
|
|
consumed.key.fill(0);
|
|
}
|
|
while (skippedMessageKeys.length > MAX_RATCHET_SKIP) {
|
|
const evicted = skippedMessageKeys.shift();
|
|
evicted?.key.fill(0);
|
|
}
|
|
|
|
return {
|
|
plaintext,
|
|
session: {
|
|
...args.session,
|
|
recvChainKey: chainKey,
|
|
recvCount: nextRecvCount,
|
|
skippedMessageKeys,
|
|
updatedAt: Date.now(),
|
|
},
|
|
};
|
|
}
|