nats-io / nats.js / 14248085491

import { decode, Empty, encode, TE } from "./encoders.ts";

import { CR_LF, CRLF, getResolveFn, newTransport } from "./transport.ts";

import { deferred, delay, extend, timeout } from "./util.ts";

import { DataBuffer } from "./databuffer.ts";

import { Servers } from "./servers.ts";

import { QueuedIteratorImpl } from "./queued_iterator.ts";

import { MuxSubscription } from "./muxsubscription.ts";

import { Heartbeat } from "./heartbeats.ts";

import { Kind, Parser } from "./parser.ts";

import { MsgImpl } from "./msg.ts";

import { Features, parseSemVer } from "./semver.ts";

import {

  DEFAULT_MAX_PING_OUT,

  DEFAULT_PING_INTERVAL,

  DEFAULT_RECONNECT_TIME_WAIT,

} from "./options.ts";

import { errors, InvalidArgumentError } from "./errors.ts";

const FLUSH_THRESHOLD = 1024 * 32;

export const INFO = /^INFO\s+([^\r\n]+)\r\n/i;

const PONG_CMD = encode("PONG\r\n");

const PING_CMD = encode("PING\r\n");

export class Connect {

  echo?: boolean;

  no_responders?: boolean;

  protocol: number;

  verbose?: boolean;

  pedantic?: boolean;

  jwt?: string;

  nkey?: string;

  sig?: string;

  user?: string;

  pass?: string;

  auth_token?: string;

  tls_required?: boolean;

  name?: string;

  lang: string;

  version: string;

  headers?: boolean;

  constructor(

    transport: { version: string; lang: string },

    opts: ConnectionOptions,

    nonce?: string,

  ) {

    this.protocol = 1;

    this.version = transport.version;

    this.lang = transport.lang;

    this.verbose = opts.verbose;

    this.pedantic = opts.pedantic;

    this.tls_required = opts.tls ? true : undefined;

    this.name = opts.name;

    const creds =

      (opts && typeof opts.authenticator === "function"

        : {}) || {};

    extend(this, creds);

  }

}

class SlowNotifier {

  slow: number;

  cb: (pending: number) => void;

  notified: boolean;

  constructor(slow: number, cb: (pending: number) => void) {

    this.slow = slow;

    this.cb = cb;

    this.notified = false;

  }

  maybeNotify(pending: number): void {

    if (pending <= this.slow) {

      this.notified = false;

    } else {

      if (!this.notified) {

        this.cb(pending);

        this.notified = true;

      }

    }

  }

}

export class SubscriptionImpl extends QueuedIteratorImpl<Msg>

  sid!: number;

  queue?: string;

  draining: boolean;

  max?: number;

  subject: string;

  drained?: Promise<void>;

  protocol: ProtocolHandler;

  timer?: Timeout<void>;

  info?: unknown;

  cleanupFn?: (sub: Subscription, info?: unknown) => void;

  closed: Deferred<void | Error>;

  requestSubject?: string;

  slow?: SlowNotifier;

  constructor(

    protocol: ProtocolHandler,

    subject: string,

    opts: SubscriptionOptions = {},

  ) {

    super();

    extend(this, opts);

    this.protocol = protocol;

    this.subject = subject;

    this.draining = false;

    this.noIterator = typeof opts.callback === "function";

    this.closed = deferred<void | Error>();

    const asyncTraces = !(protocol.options?.noAsyncTraces || false);

    if (opts.timeout) {

      this.timer = timeout<void>(opts.timeout, asyncTraces);

          this.stop(err);

          if (this.noIterator) {

            this.callback(err, {} as Msg);

          }

        });

    }

    if (!this.noIterator) {

      this.iterClosed.then((err: void | Error) => {

        this.closed.resolve(err);

        this.unsubscribe();

      });

    }

  }

  setSlowNotificationFn(slow: number, fn?: (pending: number) => void): void {

    this.slow = undefined;

    if (fn) {

      this.slow = new SlowNotifier(slow, fn);

    }

  }

  callback(err: Error | null, msg: Msg) {

    this.cancelTimeout();

    err ? this.stop(err) : this.push(msg);

      this.slow.maybeNotify(this.getPending());

    }

  }

  close(err?: Error): void {

    if (!this.isClosed()) {

      this.cancelTimeout();

      const fn = () => {

        this.stop();

        this.closed.resolve(err);

      };

      if (this.noIterator) {

        fn();

      } else {

        this.push(fn);

      }

    }

  }

  unsubscribe(max?: number): void {

    this.protocol.unsubscribe(this, max);

  }

  cancelTimeout(): void {

    if (this.timer) {

      this.timer.cancel();

      this.timer = undefined;

    }

  }

  drain(): Promise<void> {

    if (this.protocol.isClosed()) {

      return Promise.reject(new errors.ClosedConnectionError());

    }

    if (this.isClosed()) {

      return Promise.reject(

        new errors.InvalidOperationError("subscription is already closed"),

    }

    if (!this.drained) {

      this.draining = true;

      this.protocol.unsub(this);

      this.drained = this.protocol.flush(deferred<void>())

        .then(() => {

          this.protocol.subscriptions.cancel(this);

    }

    return this.drained;

  }

  isDraining(): boolean {

    return this.draining;

  }

  isClosed(): boolean {

    return this.done;

  }

  getSubject(): string {

    return this.subject;

  }

  getID(): number {

    return this.sid;

  }

}

export class Subscriptions {

  mux: SubscriptionImpl | null;

  subs: Map<number, SubscriptionImpl>;

  sidCounter: number;

  constructor() {

    this.sidCounter = 0;

    this.mux = null;

    this.subs = new Map<number, SubscriptionImpl>();

  }

  size(): number {

    return this.subs.size;

  }

  add(s: SubscriptionImpl): SubscriptionImpl {

    this.sidCounter++;

    s.sid = this.sidCounter;

    this.subs.set(s.sid, s as SubscriptionImpl);

    return s;

  }

  setMux(s: SubscriptionImpl | null): SubscriptionImpl | null {

    this.mux = s;

    return s;

  }

  getMux(): SubscriptionImpl | null {

    return this.mux;

  }

  get(sid: number): SubscriptionImpl | undefined {

    return this.subs.get(sid);

  }

  resub(s: SubscriptionImpl): SubscriptionImpl {

    this.sidCounter++;

    this.subs.delete(s.sid);

    s.sid = this.sidCounter;

    this.subs.set(s.sid, s);

    return s;

  }

  all(): (SubscriptionImpl)[] {

    return Array.from(this.subs.values());

  }

  cancel(s: SubscriptionImpl): void {

    if (s) {

      s.close();

      this.subs.delete(s.sid);

    }

  }

  handleError(err: PermissionViolationError): boolean {

    const subs = this.all();

    let sub;

    if (err.operation === "subscription") {

      sub = subs.find((s) => {

        return s.subject === err.subject && s.queue === err.queue;

      });

      sub = subs.find((s) => {

        return s.requestSubject === err.subject;

      });

    }

    if (sub) {

      sub.callback(err, {} as Msg);

      sub.close(err);

      this.subs.delete(sub.sid);

      return sub !== this.mux;

    return false;

  }

  close() {

    this.subs.forEach((sub) => {

      sub.close();

    });

  }

}

export class ProtocolHandler implements Dispatcher<ParserEvent> {

  connected: boolean;

  connectedOnce: boolean;

  infoReceived: boolean;

  info?: ServerInfo;

  muxSubscriptions: MuxSubscription;

  options: ConnectionOptions;

  outbound: DataBuffer;

  pongs: Array<Deferred<void>>;

  subscriptions: Subscriptions;

  transport!: Transport;

  noMorePublishing: boolean;

  connectError?: (err?: Error) => void;

  publisher: Publisher;

  _closed: boolean;

  closed: Deferred<Error | void>;

  listeners: QueuedIteratorImpl<Status>[];

  heartbeats: Heartbeat;

  parser: Parser;

  outMsgs: number;

  inMsgs: number;

  outBytes: number;

  inBytes: number;

  pendingLimit: number;

  lastError?: Error;

  abortReconnect: boolean;

  whyClosed: string;

  servers: Servers;

  server!: ServerImpl;

  features: Features;

  connectPromise: Promise<void> | null;

  dialDelay: Delay | null;

  raceTimer?: Timeout<void>;

  constructor(options: ConnectionOptions, publisher: Publisher) {

    this._closed = false;

    this.connected = false;

    this.connectedOnce = false;

    this.infoReceived = false;

    this.noMorePublishing = false;

    this.abortReconnect = false;

    this.listeners = [];

    this.pendingLimit = FLUSH_THRESHOLD;

    this.outMsgs = 0;

    this.inMsgs = 0;

    this.outBytes = 0;

    this.inBytes = 0;

    this.options = options;

    this.publisher = publisher;

    this.subscriptions = new Subscriptions();

    this.muxSubscriptions = new MuxSubscription();

    this.outbound = new DataBuffer();

    this.pongs = [];

    this.whyClosed = "";

    this.pendingLimit = options.pendingLimit || this.pendingLimit;

    this.features = new Features({ major: 0, minor: 0, micro: 0 });

    this.connectPromise = null;

    this.dialDelay = null;

    this.servers = new Servers(servers, {

      randomize: !options.noRandomize,

    });

    this.closed = deferred<Error | void>();

    this.parser = new Parser(this);

  }

  resetOutbound(): void {

    this.outbound.reset();

    const pongs = this.pongs;

    this.pongs = [];

    const err = new errors.RequestError("connection disconnected");

    err.stack = "";

    pongs.forEach((p) => {

      p.reject(err);

    });

    this.parser = new Parser(this);

    this.infoReceived = false;

  }

  dispatchStatus(status: Status): void {

    this.listeners.forEach((q) => {

      q.push(status);

    });

  }

  private prepare(): Deferred<void> {

      this.transport.discard();

    }

    this.info = undefined;

    this.resetOutbound();

    const pong = deferred<void>();

    pong.catch(() => {

    });

    this.pongs.unshift(pong);

    this.connectError = (err?: Error) => {

      pong.reject(err);

    };

    this.transport = newTransport();

    this.transport.closed()

      .then(async (_err?) => {

        this.connected = false;

        if (!this.isClosed()) {

          await this.disconnected(this.transport.closeError || this.lastError);

          return;

        }

      });

    return pong;

  }

  public disconnect(): void {

    this.dispatchStatus({ type: "staleConnection" });

    this.transport.disconnect();

  }

  public reconnect(): Promise<void> {

    if (this.connected) {

      this.dispatchStatus({

        type: "forceReconnect",

      });

      this.transport.disconnect();

    }

    return Promise.resolve();

  }

  async disconnected(err?: Error): Promise<void> {

    this.dispatchStatus(

      {

        type: "disconnect",

        server: this.servers.getCurrentServer().toString(),

      },

    if (this.options.reconnect) {

      await this.dialLoop()

        .then(() => {

          this.dispatchStatus(

            {

              type: "reconnect",

              server: this.servers.getCurrentServer().toString(),

            },

            this.lastError = undefined;

          }

        })

        .catch((err) => {

          this.close(err).catch();

        });

      await this.close(err).catch();

    }

  }

  async dial(srv: Server): Promise<void> {

    const pong = this.prepare();

    try {

      this.raceTimer = timeout(this.options.timeout || 20000);

      const cp = this.transport.connect(srv, this.options);

      await Promise.race([cp, this.raceTimer]);

      (async () => {

        try {

          for await (const b of this.transport) {

            this.parser.parse(b);

          }

      })().then();

      pong.reject(err);

    }

    try {

      await Promise.race([this.raceTimer, pong]);

      this.raceTimer?.cancel();

      this.connected = true;

      this.connectError = undefined;

      this.sendSubscriptions();

      this.connectedOnce = true;

      this.server.didConnect = true;

      this.server.reconnects = 0;

      this.flushPending();

      this.heartbeats.start();

      this.raceTimer?.cancel();

      await this.transport.close(err as Error);

      throw err;

    }

  }

  async _doDial(srv: Server): Promise<void> {

    const { resolve } = this.options;

    const alts = await srv.resolve({

      fn: getResolveFn(),

      debug: this.options.debug,

      randomize: !this.options.noRandomize,

      resolve,

    });

    let lastErr: Error | null = null;

    for (const a of alts) {

      try {

        lastErr = null;

        this.dispatchStatus(

          { type: "reconnecting" },

        await this.dial(a);

        return;

        lastErr = err as Error;

      }

    throw lastErr;

  }

  dialLoop(): Promise<void> {

    if (this.connectPromise === null) {

      this.connectPromise = this.dodialLoop();

      this.connectPromise

        .then(() => {})

        .catch(() => {})

        .finally(() => {

          this.connectPromise = null;

        });

    }

    return this.connectPromise;

  }

  async dodialLoop(): Promise<void> {

    let lastError: Error | undefined;

    while (true) {

        this.servers.clear();

      }

      let maxWait = wait;

      const srv = this.selectServer();

        if (lastError) {

          throw lastError;

          throw new errors.ConnectionError("connection refused");

        }

      }

      const now = Date.now();

        srv.lastConnect = Date.now();

        try {

          await this._doDial(srv);

          break;

          lastError = err as Error;

            if (this.options.waitOnFirstConnect) {

              continue;

            }

            this.servers.removeCurrentServer();

          }

          srv.reconnects++;

            this.servers.removeCurrentServer();

          }

        }

        maxWait = Math.min(maxWait, srv.lastConnect + wait - now);

        this.dialDelay = delay(maxWait);

        await this.dialDelay;

      }

    }

  }

  public static async connect(

    options: ConnectionOptions,

    publisher: Publisher,

  ): Promise<ProtocolHandler> {

    const h = new ProtocolHandler(options, publisher);

    await h.dialLoop();

    return h;

  }

  static toError(s: string): Error {

    let err: Error | null = errors.PermissionViolationError.parse(s);

    if (err) {

      return err;

    err = errors.UserAuthenticationExpiredError.parse(s);

    if (err) {

      return err;

    }

    err = errors.AuthorizationError.parse(s);

    if (err) {

      return err;

    }

    return new errors.ProtocolError(s);

  }

  processMsg(msg: MsgArg, data: Uint8Array) {

    this.inMsgs++;

    this.inBytes += data.length;

    const sub = this.subscriptions.get(msg.sid) as SubscriptionImpl;

      return;

    }

    sub.received += 1;

    if (sub.callback) {

      sub.callback(null, new MsgImpl(msg, data, this));

    }

      sub.unsubscribe();

    }

  }

  processError(m: Uint8Array) {

    let s = decode(m);

    if (s.startsWith("'") && s.endsWith("'")) {

      s = s.slice(1, s.length - 1);

    }

    const err = ProtocolHandler.toError(s);

    switch (err.constructor) {

      case errors.PermissionViolationError: {

        const pe = err as PermissionViolationError;

        const mux = this.subscriptions.getMux();

        this.subscriptions.handleError(pe);

        this.muxSubscriptions.handleError(isMuxPermission, pe);

          this.subscriptions.setMux(null);

        }

      }

    }

    this.dispatchStatus({ type: "error", error: err });

    this.handleError(err);

  }

  handleError(err: Error) {

    if (

      err instanceof errors.UserAuthenticationExpiredError ||

      err instanceof errors.AuthorizationError

      this.handleAuthError(err);

    }

      this.lastError = err;

    }

  }

  handleAuthError(err: UserAuthenticationExpiredError | AuthorizationError) {

    if (

      (this.lastError instanceof errors.UserAuthenticationExpiredError ||

        this.lastError instanceof errors.AuthorizationError) &&

      this.options.ignoreAuthErrorAbort === false

    ) {

      this.abortReconnect = true;

    }

    if (this.connectError) {

      this.connectError(err);

    } else {

      this.disconnect();

    }

  }

  processPing() {

    this.transport.send(PONG_CMD);

  }

  processPong() {

    const cb = this.pongs.shift();

    if (cb) {

      cb.resolve();

    }

  }

  processInfo(m: Uint8Array) {

    const info = JSON.parse(decode(m));

    this.info = info;

      ? undefined

      : this.servers.update(info, this.transport.isEncrypted());

    if (!this.infoReceived) {

      this.features.update(parseSemVer(info.version));

      this.infoReceived = true;

        this.servers.updateTLSName();

      }

      const { version, lang } = this.transport;

      try {

        const c = new Connect(

          { version, lang },

          this.options,

          info.nonce,

        if (info.headers) {

          c.headers = true;

          c.no_responders = true;

        }

        const cs = JSON.stringify(c);

        this.transport.send(

          encode(`CONNECT ${cs}${CR_LF}`),

        this.transport.send(PING_CMD);

        this.close(err as Error).catch();

      }

    }

    if (updates) {

      const { added, deleted } = updates;

      this.dispatchStatus({ type: "update", added, deleted });

    }

      this.dispatchStatus(

        {

          type: "ldm",

          server: this.servers.getCurrentServer().toString(),

        },

    }

  }

  push(e: ParserEvent): void {