logos-messaging-js/packages/message-encryption/src/symmetric.ts

import { Decoder as DecoderV0 } from "@waku/core/lib/message/version_0";
import type {
  EncoderOptions as BaseEncoderOptions,
  IDecoder,
  IEncoder,
  IEncryptedMessage,
  IMessage,
  IMetaSetter,
  IProtoMessage,
  IRoutingInfo,
  PubsubTopic
} from "@waku/interfaces";
import { WakuMessage } from "@waku/proto";
import { Logger } from "@waku/utils";

import { generateSymmetricKey } from "./crypto/utils.js";
import { DecodedMessage } from "./decoded_message.js";
import {
  decryptSymmetric,
  encryptSymmetric,
  postCipher,
  preCipher
} from "./encryption.js";
import { OneMillion, Version } from "./misc.js";

export {
  decryptSymmetric,
  encryptSymmetric,
  postCipher,
  preCipher,
  generateSymmetricKey
};

const log = new Logger("message-encryption:symmetric");

class Encoder implements IEncoder {
  public constructor(
    public contentTopic: string,
    public routingInfo: IRoutingInfo,
    private symKey: Uint8Array,
    private sigPrivKey?: Uint8Array,
    public ephemeral: boolean = false,
    public metaSetter?: IMetaSetter
  ) {
    if (!contentTopic || contentTopic === "") {
      throw new Error("Content topic must be specified");
    }
  }

  public get pubsubTopic(): PubsubTopic {
    return this.routingInfo.pubsubTopic;
  }

  public async toWire(message: IMessage): Promise<Uint8Array | undefined> {
    const protoMessage = await this.toProtoObj(message);
    if (!protoMessage) return;

    return WakuMessage.encode(protoMessage);
  }

  public async toProtoObj(
    message: IMessage
  ): Promise<IProtoMessage | undefined> {
    const timestamp = message.timestamp ?? new Date();
    const preparedPayload = await preCipher(message.payload, this.sigPrivKey);

    const payload = await encryptSymmetric(preparedPayload, this.symKey);

    const protoMessage = {
      payload,
      version: Version,
      contentTopic: this.contentTopic,
      timestamp: BigInt(timestamp.valueOf()) * OneMillion,
      meta: undefined,
      rateLimitProof: message.rateLimitProof,
      ephemeral: this.ephemeral
    };

    if (this.metaSetter) {
      const meta = this.metaSetter(protoMessage);
      return { ...protoMessage, meta };
    }

    return protoMessage;
  }
}

export interface EncoderOptions extends BaseEncoderOptions {
  /** The symmetric key to encrypt the payload with. */
  symKey: Uint8Array;
  /** An optional private key to be used to sign the payload before encryption. */
  sigPrivKey?: Uint8Array;
}

/**
 * Creates an encoder that encrypts messages using symmetric encryption for the
 * given key, as defined in [26/WAKU2-PAYLOAD](https://rfc.vac.dev/spec/26/).
 *
 * An encoder is used to encode messages in the [`14/WAKU2-MESSAGE](https://rfc.vac.dev/spec/14/)
 * format to be sent over the Waku network. The resulting encoder can then be
 * pass to { @link @waku/interfaces!ISender.send } to automatically encrypt
 * and encode outgoing messages.
 *
 * The payload can optionally be signed with the given private key as defined
 * in [26/WAKU2-PAYLOAD](https://rfc.vac.dev/spec/26/).
 */
export function createEncoder({
  contentTopic,
  routingInfo,
  symKey,
  sigPrivKey,
  ephemeral = false,
  metaSetter
}: EncoderOptions): Encoder {
  return new Encoder(
    contentTopic,
    routingInfo,
    symKey,
    sigPrivKey,
    ephemeral,
    metaSetter
  );
}

class Decoder extends DecoderV0 implements IDecoder<IEncryptedMessage> {
  public constructor(
    contentTopic: string,
    routingInfo: IRoutingInfo,
    private symKey: Uint8Array
  ) {
    super(contentTopic, routingInfo);
  }

  public async fromProtoObj(
    pubsubTopic: string,
    protoMessage: IProtoMessage
  ): Promise<IEncryptedMessage | undefined> {
    const cipherPayload = protoMessage.payload;

    if (protoMessage.version !== Version) {
      log.error(
        "Failed to decrypt due to incorrect version, expected:",
        Version,
        ", actual:",
        protoMessage.version
      );
      return;
    }

    let payload;

    try {
      payload = await decryptSymmetric(cipherPayload, this.symKey);
    } catch (e) {
      log.error(
        `Failed to decrypt message using asymmetric decryption for contentTopic: ${this.contentTopic}`,
        e
      );
      return;
    }

    if (!payload) {
      log.error(
        `Failed to decrypt payload for contentTopic ${this.contentTopic}`
      );
      return;
    }

    const res = postCipher(payload);

    if (!res) {
      log.error(
        `Failed to decode payload for contentTopic ${this.contentTopic}`
      );
      return;
    }

    log.info("Message decrypted", protoMessage);
    return new DecodedMessage(
      pubsubTopic,
      protoMessage,
      res.payload,
      res.sig?.signature,
      res.sig?.publicKey
    );
  }
}

/**
 * Creates a decoder that decrypts messages using symmetric encryption, using
 * the given key as defined in [26/WAKU2-PAYLOAD](https://rfc.vac.dev/spec/26/).
 *
 * A decoder is used to decode messages from the [14/WAKU2-MESSAGE](https://rfc.vac.dev/spec/14/)
 * format when received from the Waku network. The resulting decoder can then be
 * pass to { @link @waku/interfaces!IReceiver.subscribe } to automatically decrypt and
 * decode incoming messages.
 *
 * @param contentTopic The resulting decoder will only decode messages with this content topic.
 * @param routingInfo Routing information, depends on the network config (static vs auto sharding)
 * @param symKey The symmetric key used to decrypt the message.
 */
export function createDecoder(
  contentTopic: string,
  routingInfo: IRoutingInfo,
  symKey: Uint8Array
): Decoder {
  return new Decoder(contentTopic, routingInfo, symKey);
}

/**
 * Result of decrypting a message with AES symmetric encryption.
 */
export interface SymmetricDecryptionResult {
  /** The decrypted payload */
  payload: Uint8Array;
  /** The signature if the message was signed */
  signature?: Uint8Array;
  /** The recovered public key if the message was signed */
  signaturePublicKey?: Uint8Array;
}

/**
 * AES symmetric encryption.
 *
 *
 * Follows [26/WAKU2-PAYLOAD](https://rfc.vac.dev/spec/26/) encryption standard.
 */
export class SymmetricEncryption {
  /**
   * Creates an AES Symmetric encryption instance.
   *
   * @param symKey - The symmetric key for encryption (32 bytes recommended)
   * @param sigPrivKey - Optional private key to sign messages before encryption
   */
  public constructor(
    private symKey: Uint8Array,
    private sigPrivKey?: Uint8Array
  ) {}

  /**
   * Encrypts a byte array payload.
   *
   * The encryption process:
   * 1. Optionally signs the payload with the private key
   * 2. Adds padding to obscure payload size
   * 3. Encrypts using AES-256-GCM
   *
   * @param payload - The data to encrypt
   * @returns The encrypted payload
   */
  public async encrypt(payload: Uint8Array): Promise<Uint8Array> {
    const preparedPayload = await preCipher(payload, this.sigPrivKey);
    return encryptSymmetric(preparedPayload, this.symKey);
  }
}

/**
 * AES symmetric decryption.
 *
 * Follows [26/WAKU2-PAYLOAD](https://rfc.vac.dev/spec/26/) encryption standard.
 */
export class SymmetricDecryption {
  /**
   * Creates an AES Symmetric decryption instance.
   *
   * @param symKey - The symmetric key for decryption (must match encryption key)
   */
  public constructor(private symKey: Uint8Array) {}

  /**
   * Decrypts an encrypted byte array payload.
   *
   * The decryption process:
   * 1. Decrypts using AES-256-GCM
   * 2. Removes padding
   * 3. Verifies and recovers signature if present
   *
   * @param encryptedPayload - The encrypted data (from [[SymmetricEncryption.encrypt]])
   * @returns Object containing the decrypted payload and signature info, or undefined if decryption fails
   */
  public async decrypt(
    encryptedPayload: Uint8Array
  ): Promise<SymmetricDecryptionResult | undefined> {
    try {
      const decryptedData = await decryptSymmetric(
        encryptedPayload,
        this.symKey
      );

      if (!decryptedData) {
        return undefined;
      }

      const result = postCipher(decryptedData);

      if (!result) {
        return undefined;
      }

      return {
        payload: result.payload,
        signature: result.sig?.signature,
        signaturePublicKey: result.sig?.publicKey
      };
    } catch (error) {
      log.error("Failed to decrypt payload", error);
      return undefined;
    }
  }
}