refactor(extension,streams): Abstract envelope handling

[grypez]

It would be nice to abstract envelope handling in the extension so that we can compartmentalize our hard-to-validate assumptions about the messages being sent and received between various interlocutors.

[grypez]

envelope.ts

import { isObject } from '@metamask/utils';

import type { CapTpMessage, WrappedIframeMessage } from './message.js';
import { isCapTpMessage, isWrappedIframeMessage } from './message.js';

export enum EnvelopeLabel {
  Command = 'message',
  CapTp = 'capTp',
}

type Envelope<Label extends EnvelopeLabel, Content> = {
  label: Label;
  content: Content;
};

type ContentOf<Env> = Env extends Envelope<EnvelopeLabel, infer Content>
  ? Content
  : never;
type LabelOf<Env> = Env extends Envelope<infer Label, unknown> ? Label : never;
type GenericEnvelope<Env> = Envelope<LabelOf<Env>, ContentOf<Env>>;

type LabeledWith<Label extends EnvelopeLabel> = {
  label: Label;
  [key: string]: unknown;
};

type ContainerOf<Content> = {
  content: Content;
  [key: string]: unknown;
};

const isLabeled = <Label extends EnvelopeLabel>(
  value: unknown,
  label?: Label,
): value is LabeledWith<Label> =>
  isObject(value) &&
  typeof value.label !== 'undefined' &&
  (label === undefined || value.label === label);

type EnvelopeKit<Env extends GenericEnvelope<Env>> = {
  label: Env['label'];
  sniff: (value: unknown) => value is LabeledWith<Env['label']>;
  check: (value: unknown) => value is Env;
  wrap: (content: Env['content']) => Env;
  unwrap: (envelope: Env) => Env['content'];
};

const makeEnvelopeKit = <Env extends GenericEnvelope<Env>>(
  label: Env['label'],
  isContent: (value: unknown) => value is Env['content'],
): EnvelopeKit<Env> => {
  const hasLabel = (value: unknown): value is LabeledWith<Env['label']> =>
    isLabeled(value, label);
  const hasContent = (value: unknown): value is ContainerOf<ContentOf<Env>> =>
    isObject(value) && (isContent ?? ((_: unknown) => true))(value.content);
  return {
    label,
    sniff: hasLabel,
    check: (value: unknown): value is Env =>
      hasLabel(value) && hasContent(value),
    wrap: (content: ContentOf<Env>) =>
      ({
        label,
        content,
      } as Env),
    unwrap: (envelope: Env): Env['content'] => {
      if (!hasLabel(envelope)) {
        throw new Error(
          // @ts-expect-error Exhaustiveness check
          `Expected envelope labelled "${label.value}" but got "${envelope.label}".`,
        );
      }
      return envelope.content;
    },
  };
};

type CommandEnvelope = Envelope<EnvelopeLabel.Command, WrappedIframeMessage>;
type CapTpEnvelope = Envelope<EnvelopeLabel.CapTp, CapTpMessage>;

export const commandEnveloper = makeEnvelopeKit<CommandEnvelope>(EnvelopeLabel.Command, isWrappedIframeMessage);
export const capTpEnveloper = makeEnvelopeKit<CapTpEnvelope>(EnvelopeLabel.CapTp, isCapTpMessage);

export type StreamEnvelope =
  | CommandEnvelope
  | CapTpEnvelope;

type StreamEnvelopeContent<Label extends LabelOf<StreamEnvelope>> =
Label extends EnvelopeLabel.Command ? CommandEnvelope['content'] :
Label extends EnvelopeLabel.CapTp ? CapTpEnvelope['content'] : never;

const envelopeKits: { [Label in EnvelopeLabel]: EnvelopeKit<StreamEnvelope> } =
  Object.fromEntries([
    [
      EnvelopeLabel.Command,
      commandEnveloper,
    ],
    [EnvelopeLabel.CapTp, capTpEnveloper],
  ]);

export const isStreamEnvelope = (value: unknown): value is StreamEnvelope =>
  isLabeled(value) &&
  Object.values(envelopeKits).some((kit) => kit.check(value));

type ContentHandler<Label extends LabelOf<StreamEnvelope>> = (
  content: StreamEnvelopeContent<Label>,
) => Promise<unknown>;

type ErrorHandler = (...args: unknown[]) => never;

type ContentHandlerBag = { [Label in EnvelopeLabel]: ContentHandler<Label> };

export const makeStreamEnvelopeHandler =
  (contentHandlers: ContentHandlerBag, errorHandler: ErrorHandler) =>
  async (value: unknown) => {
    const envelope = isStreamEnvelope(value)
      ? value
      : errorHandler(
          `Envelope handler received unexpected value ${JSON.stringify(value)}`,
        );
    const kit =
      envelopeKits[envelope.label] ??
      errorHandler(
        `Envelope handler received an envelope with unknown label "${envelope.label}"`,
      );
    const handler: ContentHandler<typeof envelope.label> =
      contentHandlers[envelope.label] ??
      errorHandler(
        `Envelope handler received an envelope with known but unexpected label "${envelope.label}"`,
      );
    return handler(kit.unwrap(envelope));
  };

[grypez]

envelope.test.ts

import { describe, it, expect } from 'vitest';
import { capTpEnveloper, commandEnveloper, EnvelopeLabel, isStreamEnvelope, makeStreamEnvelopeHandler } from './envelope.js';
import { Command, WrappedIframeMessage } from './message.js';

describe('envelope', () => {
  describe('isStreamEnvelope', () => {
    it('returns true for valid messages', () => {
      expect(
        isStreamEnvelope({
          label: EnvelopeLabel.Command,
          content: {
            id: '1',
            message: { type: 'evaluate', data: '1 + 1' },
          },
        }),
      ).toBe(true);
    });

    it.each`
      exportName            | enveloper           | content
      ${'commandEnveloper'} | ${commandEnveloper} | ${{ id: '1', message: { type: Command.Evaluate, data: '1 + 1' } }}
      ${'capTpEnveloper'}   | ${capTpEnveloper}   | ${{ id: '4', message: { type: 'CTP_CALL', epoch: 0 } }}
    `('returns true for valid contents wrapped by $exportName', ({ enveloper, content }) => {
      expect(
        isStreamEnvelope(enveloper.wrap(content))).toBe(true);
    });

    it('returns false for invalid messages', () => {
      const invalidMessages = [
        {},
        { id: '1' },
        { message: { type: 'evaluate' } },
        { id: '1', message: { type: 'evaluate' } },
        { id: '1', message: { type: 'evaluate', data: 1 } },
      ];

      invalidMessages.forEach((message) => {
        expect(isStreamEnvelope(message)).toBe(false);
      });
    });
  });

  describe('makeStreamEnvelopeHandler', () => {
    it('handles valid StreamEnvelopes', async () => {
      const handler = makeStreamEnvelopeHandler({
        message: async content => 'message',
        capTp: async content => 'capTp',
      }, problem => { throw new Error(String(problem)); });

      expect(await handler(commandEnveloper.wrap({ id: '1', message: { type: Command.Evaluate, data: '2 + 3' }}))).toStrictEqual('message');
    })
  })
});