[squeep-mystery-box] / lib / mystery-box.js

'use strict';

const crypto = require('crypto');
const zlib = require('zlib');
const { promisify } = require('util');
const common = require('./common');
const allVersions = require('./version-parameters');
const { performance } = require('perf_hooks');

const _fileScope = common.fileScope(__filename);

const brotliCompressAsync = promisify(zlib.brotliCompress);
const brotliDecompressAsync = promisify(zlib.brotliDecompress);
const deflateRawAsync = promisify(zlib.deflateRaw);
const inflateRawAsync = promisify(zlib.inflateRaw);
const scryptAsync = promisify(crypto.scrypt);

/**
 * Only you will know what's inside your...
 *    \  |            |               | __ )
 *   |\/ | |   |  __| __|  _ \  __|   | __ \   _ \ \  /
 *   |   | |   |\__ \ |    __/ |  |   | |   | (   |`  <
 *  _|  _|\__, |____/\__|\___|_| \__, |____/ \___/ _/\_\
 *        ____/                  ____/
 *
 * Our very own way of converting a buffer or serializable object
 * to and from an opaque and web-safe representation, which we can
 * let anyone see without disclosing the contents, nor allowing it
 * to be modified without detection.
 * 
 * The result is a Base64URL-encoded byte-array consisting of:
 * - version, indicating format and encryption algorithm
 * - flags, indicating state of the contents
 * - iv/nonce, randomness for the algorithm
 * - salt, applied to a secret to derive a key
 * - tag, additional room for encryption to validate the previous fields
 * - payload, encrypted version of possibly compressed and serialized object
 *
 */


const availableFlags = {
  Brotli: (1<<0), // Slightly better compression, but slow
  Flate: (1<<1), // Slightly worse compression, but much faster especially with larger payloads.  Prefer this.
  FutureCompression: (1<<0)|(1<<1), // Something else, like zstd maybe.

  Unused2: (1<<2),
  Unused3: (1<<3),
  Unused4: (1<<4),
  Unused5: (1<<5),
  Unused6: (1<<6),

  BufferPayload: (1<<7), // Payload is a buffer, not an object
};
const compressionFlagsMask = (availableFlags.Flate | availableFlags.Brotli);
const compressionFlagsShift = 0;
const payloadFlagsMask = (availableFlags.BufferPayload);
const payloadFlagsShift = 7;


class MysteryBox {
  /**
   * @param {Console} logger
   * @param {Object} options
   * @param {String|String[]} options.encryptionSecret - if an array, will always encrypt with first secret, will attempt to decrypt with all; useful for rolling secrets
   * @param {Number=} options.defaultFlags
   */
  constructor(logger, options = {}) {
    this.logger = logger;
    this.options = options;

    this.secrets = common.ensureArray(options.encryptionSecret);
    if (!this.secrets.length) {
      throw new Error('missing encryption secret');
    }

    // Filter any unavailable algorithms
    const availableCiphers = crypto.getCiphers();
    const availableHashes = crypto.getHashes();
    // Add legacy scrypt to available hashes for filtering key derivations
    availableHashes.push(allVersions.KD.SCRYPT);
    this.versionParameters = Object.entries(allVersions).reduce((acc, [v, p]) => {
      const validCipher = availableCiphers.includes(p.algorithm);
      const validKeyDeriver = availableHashes.includes(p.keyDeriver);
      if (validCipher && validKeyDeriver) {
        acc[v] = p; // eslint-disable-line security/detect-object-injection
      }
      return acc;
    }, {});

    // Default to highest
    this.bestVersion = Number(Object.keys(this.versionParameters).sort().pop());
    if (Number.isNaN(this.bestVersion)) {
      throw new Error('no supported versions available');
    }

    this.Flags = availableFlags;
    this.defaultFlags = 'defaultFlags' in options ? options.defaultFlags : availableFlags.Flate;
    if (this.defaultFlags < 0 || this.defaultFlags > 255) {
      throw new RangeError('Invalid default flag value');
    }
  }


  /**
   * Parse the bits out of the flags.
   */
  static _decodeFlags(flags) {
    return {
      compression: (flags & compressionFlagsMask) >> compressionFlagsShift,
      payloadIsBuffer: Boolean((flags & payloadFlagsMask) >> payloadFlagsShift),
    };
  }


  /**
   * Generate key data.
   */
  static async _keyFromSecret(deriver, secret, salt, keyBytes) {
    switch (deriver) {
      case allVersions.KD.SHAKE256: {
        const hash = crypto.createHash(allVersions.KD.SHAKE256, { outputLength: keyBytes });
        hash.update(salt);
        hash.update(secret);
        return hash.digest();
      }

      case allVersions.KD.BLAKE2B512: {
        const hash = crypto.createHash(allVersions.KD.BLAKE2B512);
        hash.update(salt);
        hash.update(secret);
        const digest = hash.digest();
        // should assert that keyBytes <= 64
        // but not concerned about that for now
        // until we have a new algorithm with bigger key size
        return digest.subarray(0, keyBytes);
      }

      case allVersions.KD.SCRYPT:
        return scryptAsync(secret, salt, keyBytes);

      default:
        throw new RangeError('unsupported key deriver');
    }
  }


  /**
   * Return bits and bit mask for given number of encoded bytes.
   * @param {Number} numBytes
   * @returns {Object}
   */
  static _versionHeaderBits(numBytes) {
    // Round up to 8 bits in result, just to be proper.
    const resultBits = (((numBytes + 7) >> 3) << 3) >>> 0;
    return {
      headerValue: ((0xff << (resultBits - numBytes + 1)) & 0xff) >>> 0,
      headerMask: ((0xff << (resultBits - numBytes)) & 0xff) >>> 0,
    };
  }


  /**
   * Parse a byte into the total number of bytes in the packed number,
   * returning that and the new value for the first byte, to update in the
   * buffer before parsing as an unsigned integer.
   * Number of packed bytes is indicated by location of first leading 0 bit.
   * Support for numbers larger than 127 is of dubious practicality, but here
   * it is anyhow.
   * @param {Number} firstByte
   * @returns {Object}
   */
  static _versionHeaderDecode(firstByte) {
    for (let numBytes = 1; numBytes <= 8; numBytes++) {
      const {
        headerValue,
        headerMask,
      } = MysteryBox._versionHeaderBits(numBytes);
      if (((firstByte & headerMask) >>> 0) === headerValue) {
        const restMask = (~headerMask & 0xff) >>> 0;
        return {
          numBytes,
          firstByte: (firstByte & restMask) >>> 0,
        };
      }
    }
    // Nine bytes would be an extravagence.
    throw new RangeError(`unsupported version header (0x${firstByte.toString(16)})`);
  }


  /**
   * Decode leading bytes of buffer as version identifier.
   * In the first byte, the position of the first unset bit indicates how
   * many total bytes comprise the version, in big-endian encoding.
   * Returns decoded version and number of bytes used to decode.
   * Only supports up to 6-byte numbers, and of those, only up to 4398046511103.
   * @param {Buffer} buf - N.B. will be mogrified
   * @returns {Object}
   */
  static _versionDecode(buf) {
    const headerByte = buf.readUInt8(0);
    const { numBytes: versionBytes, firstByte } = MysteryBox._versionHeaderDecode(headerByte);
    if (versionBytes === 1) {
      return {
        version: firstByte,
        versionBytes,
      };
    }

    if (versionBytes > 6) {
      throw new RangeError(`unsupported version (${versionBytes} bytes)`);
    }

    // Otherwise, update the masked first byte and parse the rest of the buffer.
    buf[0] = firstByte;
    return {
      version: buf.readUIntBE(0, versionBytes),
      versionBytes,
    };
  }


  /**
   * Encode a version identifier into a buffer of a variable number of bytes.
   * @param {Number} version
   * @returns {Object}
   */
  static _versionEncode(version) {
    let versionBytes = 0;

    if (version <= 0x7f) { // 0-127
      versionBytes = 1;
    } else if (version <= 0x3fff) { // 128-16383
      versionBytes = 2;
    } else if (version <= 0x1fffff) { // 16384-2097151
      versionBytes = 3;
    } else if (version <= 0x0fffffff) { // 2097152-268435455
      versionBytes = 4;
    } else if (version <= 0x07ffffffff) { // 268435456-34359738367
      versionBytes = 5;
    } else if (version <= 0x03ffffffffff) { // 34359738368-4398046511103
      versionBytes = 6;
    } else {
      throw new RangeError('version too large to encode');
    }

    const buffer = Buffer.alloc(versionBytes);
    buffer.writeUIntBE(version, 0, versionBytes);
    const headerByte = ((0xff << (8 - versionBytes + 1)) & 0xff) >>> 0;
    buffer[0] = (buffer[0] | headerByte) >>> 0;

    return {
      buffer,
      versionBytes,
    };
  }


  /**
   * Put contents into a mysterious box.
   * @param {Object|Buffer} contents
   * @param {Number=} version
   * @param {Number=} flags
   * @returns {String}
   */
  async pack(contents, version = this.bestVersion, flags = this.defaultFlags) {
    const _scope = _fileScope('pack');
    const timingsMs = {
      start: performance.now(),
      preCompress: 0,
      postCompress: 0,
      preCrypt: 0,
      postCrypt: 0,
      end: 0,
    };
    const stats = {
      version,
      flags: undefined,
      serializedBytes: undefined, // original contents size in bytes
      compressedBytes: undefined, // compressed contents size in bytes
    };

    if (!(version in this.versionParameters)) {
      throw new RangeError(`MysteryBox format version ${version} not supported`);
    }
    // eslint-disable-next-line security/detect-object-injection
    const v = this.versionParameters[version];

    const { compression, payloadIsBuffer } = MysteryBox._decodeFlags(flags);

    if (Buffer.isBuffer(contents)) {
      // Ensure payloadIsBuffer flag is set when contents are indeed a Buffer
      flags |= this.Flags.BufferPayload;
    } else {
      if (payloadIsBuffer) {
        // Flag is set, but contents are not a Buffer? Try to coerce contents.
        contents = Buffer.from(contents);
      } else {
        // Otherwise attempt to serialize the object
        contents = JSON.stringify(contents);
      }
    }
    stats.serializedBytes = Buffer.byteLength(contents);

    const { buffer: versionBuffer, versionBytes } = MysteryBox._versionEncode(v.version);
    if (versionBytes !== v.versionBytes) {
      throw new Error('internal inconsistency, mismatched version byte length');
    }

    const [iv, salt] = await Promise.all([
      v.ivBytes,
      v.saltBytes,
    ].map((b) => common.randomBytesAsync(b)));

    timingsMs.preCompress = performance.now();
    let compressedContents;
    switch (compression) {
      case 0: // No compression requested
        compressedContents = contents;
        break;
      case this.Flags.Brotli:
        compressedContents = await brotliCompressAsync(contents);
        break;
      case this.Flags.Flate:
        compressedContents = await deflateRawAsync(contents);
        break;
    }
    stats.compressedBytes = Buffer.byteLength(compressedContents);
    // const compressionRatio = stats.compressedBytes / stats.serializedBytes;
    timingsMs.postCompress = timingsMs.preCrypt = performance.now();

    let payload;
    if (stats.compressedBytes >= stats.serializedBytes) {
      // If compression is not beneficial enough, or detrimental, do not use
      flags = flags & ~compressionFlagsMask;
      payload = contents;
      stats.compressedBytes = undefined;
    } else {
      payload = compressedContents;
    }

    const flagsBuffer = Buffer.alloc(v.flagsBytes);
    flagsBuffer.writeUInt8(flags, 0);
    stats.flags = this._prettyFlags(flags);

    // Authenticate all this data
    const aadBuffer = Buffer.concat([versionBuffer, flagsBuffer, iv, salt]);

    // Always encrypt with first secret
    const secret = this.secrets[0];
    const key = await MysteryBox._keyFromSecret(v.keyDeriver, secret, salt, v.keyBytes);
    const cipher = crypto.createCipheriv(v.algorithm, key, iv, v.algOptions);
    cipher.setAAD(aadBuffer);
    const encrypted = cipher.update(payload);
    const final = cipher.final();
    const tag = cipher.getAuthTag();

    const result = Buffer.concat([versionBuffer, flagsBuffer, iv, salt, tag, encrypted, final]).toString('base64url');
    timingsMs.end = timingsMs.postCrypt = performance.now();

    this.logger.debug(_scope, 'statistics', { ...stats, ...MysteryBox._timingsLog(timingsMs) });

    return result;
  }


  /**
   * Take contents out of a mysterious box.
   * @param {String} box - Base64URL encoded payload
   * @returns {Object}
   */
  async unpack(box) {
    const _scope = _fileScope('unpack');
    const timingsMs = {
      start: performance.now(),
      preCompress: 0,
      postCompress: 0,
      preCrypt: 0,
      postCrypt: 0,
      end: 0,
    };

    if (!box) {
      throw new RangeError('nothing to unpack');
    }

    const raw = Buffer.from(box, 'base64url');
    let offset = 0;

    const { version, versionBytes } = MysteryBox._versionDecode(raw);
    if (!(version in this.versionParameters)) {
      throw new RangeError('unsupported version');
    }
    // eslint-disable-next-line security/detect-object-injection
    const v = this.versionParameters[version];

    if (v.versionBytes !== versionBytes) {
      throw new Error('internal inconsistency, mismatched version byte length');
    }
    offset += v.versionBytes;

    const minBytes = v.versionBytes + v.flagsBytes + v.ivBytes + v.saltBytes + v.tagBytes;
    if (raw.length < minBytes) {
      throw new RangeError('not enough to unpack');
    }

    const flags = raw.subarray(offset, offset + v.flagsBytes).readUInt8(0);
    offset += v.flagsBytes;

    const { compression, payloadIsBuffer } = MysteryBox._decodeFlags(flags);

    const iv = raw.subarray(offset, offset + v.ivBytes);
    offset += v.ivBytes;

    const salt = raw.subarray(offset, offset + v.saltBytes);
    offset += v.saltBytes;

    const aad = raw.subarray(0, offset); // Everything up to here

    const tag = raw.subarray(offset, offset + v.tagBytes);
    offset += v.tagBytes;

    const encrypted = raw.subarray(offset);

    timingsMs.preCrypt = performance.now();

    let decrypted;
    let err;
    let success = false;
    for await (const secret of this.secrets) {
      const key = await MysteryBox._keyFromSecret(v.keyDeriver, secret, salt, v.keyBytes);
      const decipher = crypto.createDecipheriv(v.algorithm, key, iv, v.algOptions);
      decipher.setAAD(aad);
      decipher.setAuthTag(tag);

      try {
        decrypted = Buffer.concat([decipher.update(encrypted), decipher.final()]);
        success = true;
        break;
      } catch (e) {
        err = e;
        continue;
      }
    }
    if (!success) {
      throw err;
    }

    let payload;
    timingsMs.preCompress = timingsMs.postCrypt = performance.now();
    switch (compression) {
      case 0: // No compression
        payload = decrypted;
        break;
      case this.Flags.Brotli:
        payload = await brotliDecompressAsync(decrypted);
        break;
      case this.Flags.Flate:
        payload = await inflateRawAsync(decrypted);
        break;
    }
    timingsMs.end = timingsMs.postCompress = performance.now();

    if (!payloadIsBuffer) {
      payload = JSON.parse(payload.toString('utf8'));
    }

    this.logger.debug(_scope, 'statistics', { version, flags: this._prettyFlags(flags), ...MysteryBox._timingsLog(timingsMs) });

    return payload;
  }


  /**
   * Pretty-print flag values
   * @param {Number} flags
   * @returns {String}
   */
  _prettyFlags(flags) {
    const flagNames = Object.entries(this.Flags).reduce((acc, cur) => {
      const [flagName, flagValue] = cur;
      if ((flags & flagValue) === flagValue) {
        acc.push(flagName);
      }
      return acc;
    }, []);
    return `0x${flags.toString(16).padStart(2, '0')} [${flagNames.join(',')}]`;
  }


  /**
   * Everyone loves numbers.
   * @param {Object} timingsMs
   * @returns 
   */
  static _timingsLog({ start, preCompress, postCompress, preCrypt, postCrypt, end }) {
    return {
      totalMs: end - start,
      compressMs: postCompress - preCompress,
      cryptMs: postCrypt - preCrypt,
    };
  }

}

// Expose for stubbing in tests
MysteryBox._test = { crypto };

module.exports = MysteryBox;