lib/tracecontext/tracestate.js

/*
 * Copyright Elasticsearch B.V. and other contributors where applicable.
 * Licensed under the BSD 2-Clause License; you may not use this file except in
 * compliance with the BSD 2-Clause License.
 */

'use strict';

// Indirect usage of the singleton `Agent` to log.
function getLogger() {
  return require('../..').logger;
}

/**
 * Class for Managing Tracestate
 *
 * Class that creates objects for managing trace state.
 * This class is capable of parsing both tracestate strings
 * and tracestate binary representations, allowing clients
 * to get and set values in a single list-member/namespace
 * while preserving values in the other namespaces.
 *
 * Capable of working with either the binary of string
 * formatted tracestate values.
 *
 * Usage:
 *   const tracestate = TraceState.fromStringFormatString(headerTracestate, 'es')
 *   tracestate.setValue('s',1)
 *   const newHeader = tracestate.toW3cString()
 */
class TraceState {
  constructor(sourceBuffer, listMemberNamespace = 'es', defaultValues = {}) {
    if (!this._validateVendorKey(listMemberNamespace)) {
      throw new Error('Vendor namespace failed validation.');
    }

    // buffer representation of the trace state string.
    // The initial value of this.buffer will keep the
    // values set in the listMemberNamespace list-member,
    // but as soon as an initial value is set (via setValue)
    // then the listMemberNamespace values will be removed
    // from this.buffer and stored in the this.values.  While
    // slightly more complicated, this allows us to maintain
    // the order of list-member keys in an un-mutated tracestate
    // string, per the W3C spec
    this.buffer = sourceBuffer;

    this.listMemberNamespace = listMemberNamespace;

    // values for our namespace, set via setValue to
    // ensure names conform
    this.values = {};
    for (const key in defaultValues) {
      const value = defaultValues[key];
      this.setValue(key, value);
    }
  }

  setValue(key, value) {
    const strKey = String(key);
    const strValue = String(value);
    if (!this._validateElasicKeyAndValue(strKey, strValue)) {
      getLogger().trace(
        'could not set tracestate key, invalid characters detected',
      );
      return false;
    }
    const isFirstSet = Object.keys(this.values).length === 0;
    const oldValue = this.values[strKey];
    this.values[strKey] = value;

    // per: https://github.com/elastic/apm/blob/d5b2c87326548befcfec6731713932a00e430b99/specs/agents/tracing-distributed-tracing.md
    // If adding another key/value pair to the es entry would exceed the limit
    // of 256 chars, that key/value pair MUST be ignored by agents.
    // The key/value and entry separators : and ; have to be considered as well.
    const serializedValue = this._serializeValues(this.values);
    if (serializedValue.length > 256 && typeof oldValue === 'undefined') {
      delete this.values[strKey];
      return false;
    }
    if (serializedValue.length > 256 && typeof oldValue !== 'undefined') {
      this.values[strKey] = oldValue;
      return false;
    }

    // the first time we set a value, extract the mutable values from the
    // buffer and set this.values appropriately
    if (isFirstSet && Object.keys(this.values).length === 1) {
      const [buffer, values] = TraceState._removeMemberNamespaceFromBuffer(
        this.buffer,
        this.listMemberNamespace,
      );
      this.buffer = buffer;
      this.values = values;
      values[strKey] = value;
    }

    return true;
  }

  getValue(keyToGet) {
    const allValues = this.toObject();
    const rawValue = allValues[this.listMemberNamespace];
    if (!rawValue) {
      return rawValue;
    }
    const values = TraceState._parseValues(rawValue);
    return values[keyToGet];
  }

  toHexString() {
    const newBuffer = Buffer.alloc(this.buffer.length);
    let newBufferOffset = 0;
    for (let i = 0; i < this.buffer.length; i++) {
      const byte = this.buffer[i];
      if (byte === 0) {
        const indexOfKeyLength = i + 1;
        const indexOfKey = i + 2;
        const lengthKey = this.buffer[indexOfKeyLength];

        const indexOfValueLength = indexOfKey + lengthKey;
        const indexOfValue = indexOfValueLength + 1;
        const lengthValue = this.buffer[indexOfValueLength];

        const key = this.buffer
          .slice(indexOfKey, indexOfKey + lengthKey)
          .toString();
        // bail out if this is our mutable namespace
        if (key === this.listMemberNamespace) {
          continue;
        }

        // if this is not our key copy from the `0` byte to the end of the value
        this.buffer.copy(
          newBuffer,
          newBufferOffset,
          i,
          indexOfValue + lengthValue,
        );
        newBufferOffset += indexOfValue + lengthValue;

        // skip ahead to first byte after end of value
        i = indexOfValue + lengthValue - 1;
        continue;
      }
    }

    // now serialize the internal representation
    const ourBytes = [];
    if (Object.keys(this.values).length > 0) {
      // the zero byte
      ourBytes.push(0);

      // the length of the vendor namespace
      ourBytes.push(this.listMemberNamespace.length);
      // the chars of the vendor namespace
      for (let i = 0; i < this.listMemberNamespace.length; i++) {
        ourBytes.push(this.listMemberNamespace.charCodeAt(i));
      }

      // add the length of the value
      const serializedValue = this._serializeValues(this.values);
      ourBytes.push(serializedValue.length);

      // add the bytes of the value
      for (let i = 0; i < serializedValue.length; i++) {
        ourBytes.push(serializedValue.charCodeAt(i));
      }
    }
    const ourBuffer = Buffer.from(ourBytes);
    return Buffer.concat(
      [newBuffer, ourBuffer],
      newBuffer.length + ourBuffer.length,
    ).toString('hex');
  }

  /**
   * Returns JSON reprenstation of tracestate key/value pairs
   *
   * Does not parse the mutable list namespace
   */
  toObject() {
    const map = this.toMap();
    const obj = {};
    for (const key of map.keys()) {
      obj[key] = map.get(key);
    }
    return obj;
  }

  toMap() {
    const map = new Map();

    // first, serialize values from the internal representation. This means
    // The W3C spec dictates that mutated values need to be on
    // the left of the new trace string
    if (Object.keys(this.values).length) {
      map.set(this.listMemberNamespace, this._serializeValues(this.values));
    }
    for (let i = 0; i < this.buffer.length; i++) {
      const byte = this.buffer[i];
      if (byte === 0) {
        const indexOfKeyLength = i + 1;
        const indexOfKey = i + 2;
        const lengthKey = this.buffer[indexOfKeyLength];

        const indexOfValueLength = indexOfKey + lengthKey;
        const indexOfValue = indexOfValueLength + 1;
        const lengthValue = this.buffer[indexOfValueLength];

        const key = this.buffer
          .slice(indexOfKey, indexOfKey + lengthKey)
          .toString();
        const value = this.buffer
          .slice(indexOfValue, indexOfValue + lengthValue)
          .toString();

        map.set(key, value);

        // skip ahead
        i = indexOfValue + lengthValue - 1;
        continue;
      }
    }

    return map;
  }

  toString() {
    return this.toW3cString();
  }

  toW3cString() {
    const json = this.toObject();
    const chars = [];
    for (const key in json) {
      const value = json[key];
      if (!value) {
        continue;
      }
      chars.push(key);
      chars.push('=');
      chars.push(value);
      chars.push(',');
    }
    chars.pop(); // remove final comma
    return chars.join('');
  }

  _serializeValues(keyValues) {
    const chars = [];
    for (const key in keyValues) {
      const value = keyValues[key];
      chars.push(`${key}:${value}`);
      chars.push(';');
    }
    chars.pop(); // last semi-colon
    return chars.join('');
  }

  _validateVendorKey(key) {
    if (key.length > 256 || key.length < 1) {
      return false;
    }
    const re = /^[abcdefghijklmnopqrstuvwxyz0123456789_\-*/]*$/;
    if (!key.match(re)) {
      return false;
    }
    return true;
  }

  _validateElasicKeyAndValue(key, value) {
    // 0x20` to `0x7E WITHOUT `,` or `=` or `;` or `;`
    const re =
      /^[ \][!"#$%&'()*+\-./0123456789<>?@ABCDEFGHIJKLMNOPQRSTUVWXYZ^_abcdefghijklmnopqrstuvwxyz{|}~]*$/;

    if (!key.match(re) || !value.match(re)) {
      return false;
    }

    if (key.length > 256 || value.length > 256) {
      return false;
    }

    return true;
  }

  static fromBinaryFormatHexString(string, listMemberNamespace = 'es') {
    return new TraceState(Buffer.from(string, 'hex'), listMemberNamespace);
  }

  static fromStringFormatString(string = '', listMemberNamespace = 'es') {
    // converts string format to byte format
    const bytes = [];

    const parts = string.split(',');
    for (let part of parts) {
      part = part.trim(); // optional whitespace (OWS)
      if (!part) {
        continue;
      }
      const [listMember, value] = part.split('=');
      if (!listMember || !value) {
        continue;
      }
      bytes.push(0);
      bytes.push(listMember.length);
      for (let i = 0; i < listMember.length; i++) {
        bytes.push(listMember.charCodeAt(i));
      }
      bytes.push(value.length);
      for (let i = 0; i < value.length; i++) {
        bytes.push(value.charCodeAt(i));
      }
    }

    return new TraceState(Buffer.from(bytes), listMemberNamespace);
  }

  static _parseValues(rawValues) {
    const parsedValues = {};
    const parts = rawValues.split(';');
    for (const keyValue of parts) {
      if (!keyValue) {
        continue;
      }
      const [key, value] = keyValue.split(':');
      if (!key || !value) {
        continue;
      }
      parsedValues[key] = value;
    }
    return parsedValues;
  }

  static _removeMemberNamespaceFromBuffer(buffer, listMemberNamespace) {
    const newBuffer = Buffer.alloc(buffer.length);
    let newBufferOffset = 0;
    const values = {};
    for (let i = 0; i < buffer.length; i++) {
      const byte = buffer[i];
      if (byte === 0) {
        const indexOfKeyLength = i + 1;
        const indexOfKey = i + 2;
        const lengthKey = buffer[indexOfKeyLength];

        const indexOfValueLength = indexOfKey + lengthKey;
        const indexOfValue = indexOfValueLength + 1;
        const lengthValue = buffer[indexOfValueLength];

        const key = buffer.slice(indexOfKey, indexOfKey + lengthKey).toString();

        // if this is our mutable namespace extract
        // and set the value in values, otherwise
        // copy into new buffer
        if (key === listMemberNamespace) {
          const rawValues = buffer
            .slice(indexOfValue, indexOfValue + lengthValue)
            .toString();
          const parsedValues = TraceState._parseValues(rawValues);
          for (const key in parsedValues) {
            values[key] = parsedValues[key];
          }
          continue;
        } else {
          buffer.copy(
            newBuffer,
            newBufferOffset,
            i,
            indexOfValue + lengthValue,
          );
          newBufferOffset += indexOfValue + lengthValue - i;
        }

        // skip ahead to first byte after end of value
        i = indexOfValue + lengthValue - 1;
        continue;
      }
    }

    // trim off extra 0 bytes
    const trimmedBuffer = newBuffer.slice(0, newBufferOffset);

    return [trimmedBuffer, values];
  }
}

module.exports = TraceState;