<commit_before>'use strict';
const fs = require('fs');
const thrift = require('thrift');
const parquet_thrift = require('../gen-nodejs/parquet_types')
const parquet_shredder = require('./shred')
const parquet_util = require('./util')
const parquet_codec_plain = require('./encodings/plain')
const parquet_codec_rle = require('./encodings/rle')

/**
 * Parquet File Magic String
 */
const PARQUET_MAGIC = 'PAR1';

/**
 * Parquet File Format Version
 */
const PARQUET_VERSION = 1;

/**
 * Default Page and Row Group sizes
 */
const PARQUET_DEFAULT_PAGE_SIZE = 8192;
const PARQUET_DEFAULT_ROW_GROUP_SIZE = 4096;

/**
 * Repetition and Definition Level Encoding
 */
const PARQUET_RDLVL_TYPE = 'INT32';
const PARQUET_RDLVL_ENCODING = 'RLE';

/**
 * Write a parquet file to an output stream. The ParquetWriter will perform
 * buffering/batching for performance, so close() must be called after all rows
 * are written.
 */
class ParquetWriter {

  /**
   * Create a new buffered parquet writer that writes to the specified file
   */
  static openFile(schema, path, opts, callback) {
    let envelopeWriter = ParquetEnvelopeWriter.openFile(schema, path, opts);
    envelopeWriter.writeHeader((err) => {
      if (err) {
        envelopeWriter.close();
        callback(err, null);
      } else {
        callback(null, new ParquetWriter(schema, envelopeWriter));
      }
    });
  }

  /**
   * Create a new buffered parquet writer for a given envelope writer
   */
  constructor(schema, envelopeWriter) {
    this.schema = schema;
    this.envelopeWriter = envelopeWriter;
    this.rowBuffer = {};
    this.rowGroupSize = PARQUET_DEFAULT_ROW_GROUP_SIZE;
    this.closed = false;
  }

  /**
   * Append a single row to the parquet file. Rows are buffered in memory until
   * rowGroupSize rows are in the buffer or close() is called
   */
  appendRow(row, callback) {
    if (!callback) {
      callback = () => {};
    }

    if (this.closed) {
      throw 'writer was closed';
    }

    parquet_shredder.shredRecord(this.schema, row, this.rowBuffer);

    if (this.rowBuffer.rowCount >= this.rowGroupSize) {
      this.envelopeWriter.writeRowGroup(this.rowBuffer, callback);
      this.rowBuffer = {};
    } else {
      callback(null);
    }
  }

  /**
   * Finish writing the parquet file and commit the footer to disk. This method
   * MUST be called after you are finished adding rows. You must not call this
   * method twice on the same object or add any rows after the close() method has
   * been called
   */
  close(callback) {
    if (!callback) {
      callback = () => {};
    }

    if (this.closed) {
      throw 'writer was closed';
    }

    this.closed = true;
    let closeFinally = () => {
      this.envelopeWriter.writeFooter((err) => {
        if (err) {
          callback(err);
        } else {
          this.envelopeWriter.close(callback);
          this.envelopeWriter = null;
        }
      });
    };

    if (this.rowBuffer.rowCount > 0) {
      this.envelopeWriter.writeRowGroup(this.rowBuffer, (err) => {
        if (err) {
          callback(err);
        } else {
          closeFinally(err);
        }
      });
      this.rowBuffer = {};
    } else {
      closeFinally();
    }
  }

  /**
   * Set the parquet row group size. This values controls the maximum number
   * of rows that are buffered in memory at any given time as well as the number
   * of rows that are co-located on disk. A higher value is generally better for
   * read-time I/O performance at the tradeoff of write-time memory usage.
   */
  setRowGroupSize(cnt) {
    this.rowGroupSize = cnt;
  }

  /**
   * Set the parquet data page size. The data page size controls the maximum
   * number of column values that are written to disk as a consecutive array
   */
  setPageSize(cnt) {
    this.writer.setPageSize(cnt);
  }

}

/**
 * Create a parquet file from a schema and a number of row groups. This class
 * performs direct, unbuffered writes to the underlying output stream and is
 * intendend for advanced and internal users; the writeXXX methods must be
 * called in the correct order to produce a valid file.
 */
class ParquetEnvelopeWriter {

  /**
   * Create a new parquet envelope writer that writes to the specified file
   */
  static openFile(schema, path, opts) {
    let outputStream = fs.createWriteStream(path, opts);

    let writeFn = function(buf, callback) {
      outputStream.write(buf, callback);
    }

    let closeFn = function(callback) {
      outputStream.end(callback);
    }

    return new ParquetEnvelopeWriter(schema, writeFn, closeFn, 0);
  }

  constructor(schema, writeFn, closeFn, fileOffset) {
    this.schema = schema;
    this.write = writeFn,
    this.close = closeFn;
    this.offset = fileOffset;
    this.rowCount = 0;
    this.rowGroups = [];
    this.pageSize = PARQUET_DEFAULT_PAGE_SIZE;
  }

  writeSection(buf, callback) {
    this.offset += buf.length;
    this.write(buf, callback);
  }

  /**
   * Encode the parquet file header
   */
  writeHeader(callback) {
    this.writeSection(Buffer.from(PARQUET_MAGIC), callback);
  }

  /**
   * Encode a parquet row group. The records object should be created using the
   * shredRecord method
   */
  writeRowGroup(records, callback) {
    let rgroup = encodeRowGroup(this.schema, records.rowCount, records.columnData, {
      baseOffset: this.offset,
      pageSize: this.pageSize
    });

    this.rowCount += records.rowCount;
    this.rowGroups.push(rgroup.metadata);
    this.writeSection(rgroup.body, callback);
  }

  /**
   * Write the parquet file footer
   */
  writeFooter(callback) {
    if (this.rowCount == 0) {
      throw 'cannot write parquet file with zero rows';
    }

    if (this.schema.columns.length == 0) {
      throw 'cannot write parquet file with zero columns';
    }

    this.writeSection(
        encodeFooter(this.schema, this.rowCount, this.rowGroups),
        callback);
  };

  /**
   * Set the parquet data page size. The data page size controls the maximum
   * number of column values that are written to disk as a consecutive array
   */
  setPageSize(cnt) {
    this.pageSize = cnt;
  }

}

/**
 * Encode a parquet data page
 */
function encodeDataPage(column, valueCount, values, rlevels, dlevels) {
  /* encode values */
  let valuesBuf;
  switch (column.encoding) {

    case 'PLAIN':
      valuesBuf = parquet_codec_plain.encodeValues(column.primitiveType, values);
      break;

    case 'RLE':
      valuesBuf = parquet_codec_rle.encodeValues(column.primitiveType, values);
      break;

    default:
      throw 'invalid encoding: ' + encoding;

  }

  /* encode repetition and definition levels */
  let rLevelsBuf = parquet_codec_rle.encodeValues(
      PARQUET_RDLVL_TYPE,
      rlevels,
      { bitWidth: parquet_util.getBitWidth(column.rLevelMax) });

  let dLevelsBuf = parquet_codec_rle.encodeValues(
      PARQUET_RDLVL_TYPE,
      dlevels,
      { bitWidth: parquet_util.getBitWidth(column.dLevelMax) });

  /* build page header */
  let pageBody = Buffer.concat([rLevelsBuf, dLevelsBuf, valuesBuf]);
  let pageHeader = new parquet_thrift.PageHeader()
  pageHeader.num_values = valueCount;
  pageHeader.type =
      parquet_util.setThriftEnum(parquet_thrift.PageType, 'DATA_PAGE');

  pageHeader.encoding =
      parquet_util.setThriftEnum(parquet_thrift.Encoding, column.encoding);

  pageHeader.definition_level_encoding =
    parquet_util.setThriftEnum(parquet_thrift.Encoding, PARQUET_RDLVL_ENCODING);

  pageHeader.repetition_level_encoding =
      parquet_util.setThriftEnum(parquet_thrift.Encoding, PARQUET_RDLVL_ENCODING);

  pageHeader.uncompressed_page_size = pageBody.length;
  pageHeader.compressed_page_size = pageBody.length;
  pageHeader.data_page_header = new parquet_thrift.DataPageHeader();
  pageHeader.data_page_header.num_values = pageHeader.num_values;
  pageHeader.data_page_header.encoding = pageHeader.encoding;

  pageHeader.data_page_header.definition_level_encoding =
      pageHeader.definition_level_encoding;

  pageHeader.data_page_header.repetition_level_encoding =
      pageHeader.repetition_level_encoding;

  return Buffer.concat([parquet_util.serializeThrift(pageHeader), pageBody]);
}

/**
 * Encode an array of values into a parquet column chunk
 */
function encodeColumnChunk(opts) {
  let metadata = new parquet_thrift.ColumnMetaData();
  metadata.type = parquet_util.setThriftEnum(parquet_thrift.Type, opts.column.primitiveType);
  metadata.path_in_schema = opts.column.path;
  metadata.num_values = opts.values.length;
  metadata.data_page_offset = opts.baseOffset;
  metadata.encodings = [];

  /* list encodings */
  let encodingsSet = {};
  encodingsSet[PARQUET_RDLVL_ENCODING] = true;
  encodingsSet[opts.column.encoding] = true;
  for (let k in encodingsSet) {
    metadata.encodings.push(parquet_util.setThriftEnum(parquet_thrift.Encoding, k));
  }

  /* build pages */
  let values = opts.values;
  let pages = [];

  for (let voff = 0; voff <= values.length; voff += opts.pageSize) {
    /* extract page data vectors */
    let pageValues = [];
    let pageRLevels = [];
    let pageDLevels = [];

    for (let i = 0; i < Math.min(opts.pageSize, values.length - voff); ++i) {
      if (values[voff + i].value !== null) {
        pageValues.push(values[voff + i].value);
      }

      pageRLevels.push(values[voff + i].rlevel);
      pageDLevels.push(values[voff + i].dlevel);
    }

    /* encode data page */
    pages.push(
        encodeDataPage(
            opts.column,
            values.length,
            pageValues,
            pageRLevels,
            pageDLevels));
  }

  /* build column chunk header */
  let pagesBuf = Buffer.concat(pages);
  metadata.total_uncompressed_size = pagesBuf.length;
  metadata.total_compressed_size = pagesBuf.length;
  metadata.codec =
      parquet_util.setThriftEnum(parquet_thrift.CompressionCodec, 'UNCOMPRESSED');

  let metadataOffset = opts.baseOffset + pagesBuf.length;
  let body = Buffer.concat([pagesBuf, parquet_util.serializeThrift(metadata)]);
  return { body, metadata, metadataOffset };
}

/**
 * Encode a list of column values into a parquet row group
 */
function encodeRowGroup(schema, rowCount, rows, opts) {
  let metadata = new parquet_thrift.RowGroup();
  metadata.num_rows = rowCount;
  metadata.columns = [];
  metadata.total_byte_size = 0;

  let body = Buffer.alloc(0);
  for (let colDef of schema.columns) {
    let cchunkData = encodeColumnChunk({
      column: colDef,
      values: rows[colDef.name],
      baseOffset: opts.baseOffset + body.length,
      pageSize: opts.pageSize,
      encoding: colDef.encoding
    });

    let cchunk = new parquet_thrift.ColumnChunk();
    cchunk.file_offset = cchunkData.metadataOffset;
    cchunk.meta_data = cchunkData.metadata;
    metadata.columns.push(cchunk);
    metadata.total_byte_size += cchunkData.body.length;

    body = Buffer.concat([body, cchunkData.body]);
  }

  return { body, metadata };
}

/**
 * Encode a parquet file metadata footer
 */
function encodeFooter(schema, rowCount, rowGroups) {
  let metadata = new parquet_thrift.FileMetaData()
  metadata.version = PARQUET_VERSION;
  metadata.created_by = 'parquet.js';
  metadata.num_rows = rowCount;
  metadata.row_groups = rowGroups;
  metadata.schema = [];

  {
    let schemaRoot = new parquet_thrift.SchemaElement();
    schemaRoot.name = 'root';
    schemaRoot.num_children = schema.columns.length;
    metadata.schema.push(schemaRoot);
  }

  for (let colDef of schema.columns) {
    let schemaElem = new parquet_thrift.SchemaElement();
    schemaElem.name = colDef.name;
    schemaElem.type = parquet_util.setThriftEnum(parquet_thrift.Type, colDef.primitiveType);
    schemaElem.repetition_type = parquet_util.setThriftEnum(
        parquet_thrift.FieldRepetitionType,
        colDef.repetition_type);

    if (colDef.convertedType) {
      schemaElem.converted_type =
          parquet_util.setThriftEnum(parquet_thrift.ConvertedType, colDef.convertedType);
    }

    metadata.schema.push(schemaElem);
  }

  let metadataEncoded = parquet_util.serializeThrift(metadata);
  let footerEncoded = new Buffer(metadataEncoded.length + 8);
  metadataEncoded.copy(footerEncoded);
  footerEncoded.writeUInt32LE(metadataEncoded.length, metadataEncoded.length);
  footerEncoded.write(PARQUET_MAGIC, metadataEncoded.length + 4);
  return footerEncoded;
}

module.exports = {
  ParquetEnvelopeWriter,
  ParquetWriter,
};

<commit_msg>use async file stream open<commit_after>'use strict';
const fs = require('fs');
const thrift = require('thrift');
const parquet_thrift = require('../gen-nodejs/parquet_types')
const parquet_shredder = require('./shred')
const parquet_util = require('./util')
const parquet_codec_plain = require('./encodings/plain')
const parquet_codec_rle = require('./encodings/rle')

/**
 * Parquet File Magic String
 */
const PARQUET_MAGIC = 'PAR1';

/**
 * Parquet File Format Version
 */
const PARQUET_VERSION = 1;

/**
 * Default Page and Row Group sizes
 */
const PARQUET_DEFAULT_PAGE_SIZE = 8192;
const PARQUET_DEFAULT_ROW_GROUP_SIZE = 4096;

/**
 * Repetition and Definition Level Encoding
 */
const PARQUET_RDLVL_TYPE = 'INT32';
const PARQUET_RDLVL_ENCODING = 'RLE';

/**
 * Write a parquet file to an output stream. The ParquetWriter will perform
 * buffering/batching for performance, so close() must be called after all rows
 * are written.
 */
class ParquetWriter {

  /**
   * Create a new buffered parquet writer that writes to the specified file
   */
  static openFile(schema, path, opts, callback) {
    ParquetEnvelopeWriter.openFile(schema, path, opts, (err, envelopeWriter) => {
      if (err) {
        return callback(err, null);
      }

      envelopeWriter.writeHeader((err) => {
        if (err) {
          envelopeWriter.close();
          callback(err, null);
        } else {
          callback(null, new ParquetWriter(schema, envelopeWriter));
        }
      });
    });
  }

  /**
   * Create a new buffered parquet writer for a given envelope writer
   */
  constructor(schema, envelopeWriter) {
    this.schema = schema;
    this.envelopeWriter = envelopeWriter;
    this.rowBuffer = {};
    this.rowGroupSize = PARQUET_DEFAULT_ROW_GROUP_SIZE;
    this.closed = false;
  }

  /**
   * Append a single row to the parquet file. Rows are buffered in memory until
   * rowGroupSize rows are in the buffer or close() is called
   */
  appendRow(row, callback) {
    if (!callback) {
      callback = () => {};
    }

    if (this.closed) {
      throw 'writer was closed';
    }

    parquet_shredder.shredRecord(this.schema, row, this.rowBuffer);

    if (this.rowBuffer.rowCount >= this.rowGroupSize) {
      this.envelopeWriter.writeRowGroup(this.rowBuffer, callback);
      this.rowBuffer = {};
    } else {
      callback(null);
    }
  }

  /**
   * Finish writing the parquet file and commit the footer to disk. This method
   * MUST be called after you are finished adding rows. You must not call this
   * method twice on the same object or add any rows after the close() method has
   * been called
   */
  close(callback) {
    if (!callback) {
      callback = () => {};
    }

    if (this.closed) {
      throw 'writer was closed';
    }

    this.closed = true;
    let closeFinally = () => {
      this.envelopeWriter.writeFooter((err) => {
        if (err) {
          callback(err);
        } else {
          this.envelopeWriter.close(callback);
          this.envelopeWriter = null;
        }
      });
    };

    if (this.rowBuffer.rowCount > 0) {
      this.envelopeWriter.writeRowGroup(this.rowBuffer, (err) => {
        if (err) {
          callback(err);
        } else {
          closeFinally(err);
        }
      });
      this.rowBuffer = {};
    } else {
      closeFinally();
    }
  }

  /**
   * Set the parquet row group size. This values controls the maximum number
   * of rows that are buffered in memory at any given time as well as the number
   * of rows that are co-located on disk. A higher value is generally better for
   * read-time I/O performance at the tradeoff of write-time memory usage.
   */
  setRowGroupSize(cnt) {
    this.rowGroupSize = cnt;
  }

  /**
   * Set the parquet data page size. The data page size controls the maximum
   * number of column values that are written to disk as a consecutive array
   */
  setPageSize(cnt) {
    this.writer.setPageSize(cnt);
  }

}

/**
 * Create a parquet file from a schema and a number of row groups. This class
 * performs direct, unbuffered writes to the underlying output stream and is
 * intendend for advanced and internal users; the writeXXX methods must be
 * called in the correct order to produce a valid file.
 */
class ParquetEnvelopeWriter {

  /**
   * Create a new parquet envelope writer that writes to the specified file
   */
  static openFile(schema, path, opts, callback) {
    let outputStream = fs.createWriteStream(path, opts);

    outputStream.on('open', function(fd) {
      let writeFn = function(buf, callback) {
        outputStream.write(buf, callback);
      }

      let closeFn = function(callback) {
        outputStream.end(callback);
      }

      callback(null, new ParquetEnvelopeWriter(schema, writeFn, closeFn, 0));
    });

    outputStream.on('error', function(err) {
      callback(err, null);
    });
  }

  constructor(schema, writeFn, closeFn, fileOffset) {
    this.schema = schema;
    this.write = writeFn,
    this.close = closeFn;
    this.offset = fileOffset;
    this.rowCount = 0;
    this.rowGroups = [];
    this.pageSize = PARQUET_DEFAULT_PAGE_SIZE;
  }

  writeSection(buf, callback) {
    this.offset += buf.length;
    this.write(buf, callback);
  }

  /**
   * Encode the parquet file header
   */
  writeHeader(callback) {
    this.writeSection(Buffer.from(PARQUET_MAGIC), callback);
  }

  /**
   * Encode a parquet row group. The records object should be created using the
   * shredRecord method
   */
  writeRowGroup(records, callback) {
    let rgroup = encodeRowGroup(this.schema, records.rowCount, records.columnData, {
      baseOffset: this.offset,
      pageSize: this.pageSize
    });

    this.rowCount += records.rowCount;
    this.rowGroups.push(rgroup.metadata);
    this.writeSection(rgroup.body, callback);
  }

  /**
   * Write the parquet file footer
   */
  writeFooter(callback) {
    if (this.rowCount == 0) {
      throw 'cannot write parquet file with zero rows';
    }

    if (this.schema.columns.length == 0) {
      throw 'cannot write parquet file with zero columns';
    }

    this.writeSection(
        encodeFooter(this.schema, this.rowCount, this.rowGroups),
        callback);
  };

  /**
   * Set the parquet data page size. The data page size controls the maximum
   * number of column values that are written to disk as a consecutive array
   */
  setPageSize(cnt) {
    this.pageSize = cnt;
  }

}

/**
 * Encode a parquet data page
 */
function encodeDataPage(column, valueCount, values, rlevels, dlevels) {
  /* encode values */
  let valuesBuf;
  switch (column.encoding) {

    case 'PLAIN':
      valuesBuf = parquet_codec_plain.encodeValues(column.primitiveType, values);
      break;

    case 'RLE':
      valuesBuf = parquet_codec_rle.encodeValues(column.primitiveType, values);
      break;

    default:
      throw 'invalid encoding: ' + encoding;

  }

  /* encode repetition and definition levels */
  let rLevelsBuf = parquet_codec_rle.encodeValues(
      PARQUET_RDLVL_TYPE,
      rlevels,
      { bitWidth: parquet_util.getBitWidth(column.rLevelMax) });

  let dLevelsBuf = parquet_codec_rle.encodeValues(
      PARQUET_RDLVL_TYPE,
      dlevels,
      { bitWidth: parquet_util.getBitWidth(column.dLevelMax) });

  /* build page header */
  let pageBody = Buffer.concat([rLevelsBuf, dLevelsBuf, valuesBuf]);
  let pageHeader = new parquet_thrift.PageHeader()
  pageHeader.num_values = valueCount;
  pageHeader.type =
      parquet_util.setThriftEnum(parquet_thrift.PageType, 'DATA_PAGE');

  pageHeader.encoding =
      parquet_util.setThriftEnum(parquet_thrift.Encoding, column.encoding);

  pageHeader.definition_level_encoding =
    parquet_util.setThriftEnum(parquet_thrift.Encoding, PARQUET_RDLVL_ENCODING);

  pageHeader.repetition_level_encoding =
      parquet_util.setThriftEnum(parquet_thrift.Encoding, PARQUET_RDLVL_ENCODING);

  pageHeader.uncompressed_page_size = pageBody.length;
  pageHeader.compressed_page_size = pageBody.length;
  pageHeader.data_page_header = new parquet_thrift.DataPageHeader();
  pageHeader.data_page_header.num_values = pageHeader.num_values;
  pageHeader.data_page_header.encoding = pageHeader.encoding;

  pageHeader.data_page_header.definition_level_encoding =
      pageHeader.definition_level_encoding;

  pageHeader.data_page_header.repetition_level_encoding =
      pageHeader.repetition_level_encoding;

  return Buffer.concat([parquet_util.serializeThrift(pageHeader), pageBody]);
}

/**
 * Encode an array of values into a parquet column chunk
 */
function encodeColumnChunk(opts) {
  let metadata = new parquet_thrift.ColumnMetaData();
  metadata.type = parquet_util.setThriftEnum(parquet_thrift.Type, opts.column.primitiveType);
  metadata.path_in_schema = opts.column.path;
  metadata.num_values = opts.values.length;
  metadata.data_page_offset = opts.baseOffset;
  metadata.encodings = [];

  /* list encodings */
  let encodingsSet = {};
  encodingsSet[PARQUET_RDLVL_ENCODING] = true;
  encodingsSet[opts.column.encoding] = true;
  for (let k in encodingsSet) {
    metadata.encodings.push(parquet_util.setThriftEnum(parquet_thrift.Encoding, k));
  }

  /* build pages */
  let values = opts.values;
  let pages = [];

  for (let voff = 0; voff <= values.length; voff += opts.pageSize) {
    /* extract page data vectors */
    let pageValues = [];
    let pageRLevels = [];
    let pageDLevels = [];

    for (let i = 0; i < Math.min(opts.pageSize, values.length - voff); ++i) {
      if (values[voff + i].value !== null) {
        pageValues.push(values[voff + i].value);
      }

      pageRLevels.push(values[voff + i].rlevel);
      pageDLevels.push(values[voff + i].dlevel);
    }

    /* encode data page */
    pages.push(
        encodeDataPage(
            opts.column,
            values.length,
            pageValues,
            pageRLevels,
            pageDLevels));
  }

  /* build column chunk header */
  let pagesBuf = Buffer.concat(pages);
  metadata.total_uncompressed_size = pagesBuf.length;
  metadata.total_compressed_size = pagesBuf.length;
  metadata.codec =
      parquet_util.setThriftEnum(parquet_thrift.CompressionCodec, 'UNCOMPRESSED');

  let metadataOffset = opts.baseOffset + pagesBuf.length;
  let body = Buffer.concat([pagesBuf, parquet_util.serializeThrift(metadata)]);
  return { body, metadata, metadataOffset };
}

/**
 * Encode a list of column values into a parquet row group
 */
function encodeRowGroup(schema, rowCount, rows, opts) {
  let metadata = new parquet_thrift.RowGroup();
  metadata.num_rows = rowCount;
  metadata.columns = [];
  metadata.total_byte_size = 0;

  let body = Buffer.alloc(0);
  for (let colDef of schema.columns) {
    let cchunkData = encodeColumnChunk({
      column: colDef,
      values: rows[colDef.name],
      baseOffset: opts.baseOffset + body.length,
      pageSize: opts.pageSize,
      encoding: colDef.encoding
    });

    let cchunk = new parquet_thrift.ColumnChunk();
    cchunk.file_offset = cchunkData.metadataOffset;
    cchunk.meta_data = cchunkData.metadata;
    metadata.columns.push(cchunk);
    metadata.total_byte_size += cchunkData.body.length;

    body = Buffer.concat([body, cchunkData.body]);
  }

  return { body, metadata };
}

/**
 * Encode a parquet file metadata footer
 */
function encodeFooter(schema, rowCount, rowGroups) {
  let metadata = new parquet_thrift.FileMetaData()
  metadata.version = PARQUET_VERSION;
  metadata.created_by = 'parquet.js';
  metadata.num_rows = rowCount;
  metadata.row_groups = rowGroups;
  metadata.schema = [];

  {
    let schemaRoot = new parquet_thrift.SchemaElement();
    schemaRoot.name = 'root';
    schemaRoot.num_children = schema.columns.length;
    metadata.schema.push(schemaRoot);
  }

  for (let colDef of schema.columns) {
    let schemaElem = new parquet_thrift.SchemaElement();
    schemaElem.name = colDef.name;
    schemaElem.type = parquet_util.setThriftEnum(parquet_thrift.Type, colDef.primitiveType);
    schemaElem.repetition_type = parquet_util.setThriftEnum(
        parquet_thrift.FieldRepetitionType,
        colDef.repetition_type);

    if (colDef.convertedType) {
      schemaElem.converted_type =
          parquet_util.setThriftEnum(parquet_thrift.ConvertedType, colDef.convertedType);
    }

    metadata.schema.push(schemaElem);
  }

  let metadataEncoded = parquet_util.serializeThrift(metadata);
  let footerEncoded = new Buffer(metadataEncoded.length + 8);
  metadataEncoded.copy(footerEncoded);
  footerEncoded.writeUInt32LE(metadataEncoded.length, metadataEncoded.length);
  footerEncoded.write(PARQUET_MAGIC, metadataEncoded.length + 4);
  return footerEncoded;
}

module.exports = {
  ParquetEnvelopeWriter,
  ParquetWriter,
};

<|endoftext|>


<commit_before>
/**
 * Module dependencies.
 */

var assert = require('assert');
var inherits = require('util').inherits;
var debug = require('debug')('wave:writer');
var Transform = require('stream').Transform;
var Parser = require('stream-parser');

// for node v0.8.x support, remove after v0.12.x
if (!Transform) Transform = require('readable-stream/transform');

/**
 * Module exports.
 */

module.exports = Writer;

/**
 * RIFF Chunk IDs in Buffers.
 *
 * @api private
 */

var RIFF = new Buffer('RIFF');
var WAVE = new Buffer('WAVE');
var fmt  = new Buffer('fmt ');
var data = new Buffer('data');

/**
 * The max size of the "data" chunk of a WAVE file. This is the max unsigned
 * 32-bit int value, minus 100 bytes (overkill, 44 would be safe) for the header.
 *
 * @api private
 */

var MAX_WAV = 4294967295 - 100;

/**
 * The `Writer` class accepts raw audio data written to it (only PCM audio data is
 * currently supported), and outputs a WAV file with a valid WAVE header at the
 * beginning specifying the formatting information of the audio stream.
 *
 * Note that there's an interesting problem, because the WAVE header also
 * specifies the total byte length of the audio data in the file, and there's no
 * way that we can know this ahead of time. Therefore the WAVE header will contain
 * a byte-length if `0` initially, which most WAVE decoders will know means to
 * just read until `EOF`.
 *
 * Optionally, if you are in a situation where you can seek back to the beginning
 * of the destination of the WAVE file (like writing to a regular file, for
 * example), then you may listen for the `"header"` event which will be emitted
 * _after_ all the data has been written, and you can go back and rewrite the new
 * header with proper audio byte length into the beginning of the destination
 * (though if your destination _is_ a regular file, you should use the the
 * `FileWriter` class instead).
 *
 * @param {Object} opts optional options object
 * @api public
 */

function Writer (opts) {
  if (!(this instanceof Writer)) {
    return new Writer(opts);
  }
  Transform.call(this, opts);

  // TODO: allow/properly handle other WAVE audio formats
  this.endianness = 'LE';
  this.format = 1; // raw PCM
  this.channels = 2;
  this.sampleRate = 44100;
  this.bitDepth = 16;

  if (opts) {
    if (null != opts.format) this.format = opts.format;
    if (null != opts.channels) this.channels = opts.channels;
    if (null != opts.sampleRate) this.sampleRate = opts.sampleRate;
    if (null != opts.bitDepth) this.bitDepth = opts.bitDepth;
  }

  // the first time _transform() is called, _writeHeader() will be run
  //this.immediate(this._writeHeader);
  this._bytes(0, this._writeHeader);
}
inherits(Writer, Transform);

/**
 * Mixin `Parser`.
 */

Parser(Writer.prototype);

/**
 * Writes the WAVE header.
 *
 * @api private
 */

Writer.prototype._writeHeader = function (chunk, write) {
  debug('_writeHeader()');
  var headerLength = 44; // TODO: 44 is only for format 1 (PCM), any other
                         // format will have a variable size...
  var dataLength = this.dataLength;
  if (null == dataLength) {
    debug('using default "dataLength" of %d', MAX_WAV);
    dataLength = MAX_WAV;
  }
  var fileSize = dataLength + headerLength;
  var header = new Buffer(headerLength);
  var offset = 0;

  // write the "RIFF" identifier
  RIFF.copy(header, offset);
  offset += RIFF.length;

  // write the file size minus the identifier and this 32-bit int
  header['writeUInt32' + this.endianness](fileSize - 8, offset);
  offset += 4;

  // write the "WAVE" identifier
  WAVE.copy(header, offset);
  offset += WAVE.length;

  // write the "fmt " sub-chunk identifier
  fmt.copy(header, offset);
  offset += fmt.length;

  // write the size of the "fmt " chunk
  // XXX: value of 16 is hard-coded for raw PCM format. other formats have
  // different size.
  header['writeUInt32' + this.endianness](16, offset);
  offset += 4;

  // write the audio format code
  header['writeUInt16' + this.endianness](this.format, offset);
  offset += 2;

  // write the number of channels
  header['writeUInt16' + this.endianness](this.channels, offset);
  offset += 2;

  // write the sample rate
  header['writeUInt32' + this.endianness](this.sampleRate, offset);
  offset += 4;

  // write the byte rate
  var byteRate = this.byteRate;
  if (null == byteRate) {
    byteRate = this.sampleRate * this.channels * this.bitDepth / 8;
  }
  header['writeUInt32' + this.endianness](byteRate, offset);
  offset += 4;

  // write the block align
  var blockAlign = this.blockAlign;
  if (null == blockAlign) {
    blockAlign = this.channels * this.bitDepth / 8;
  }
  header['writeUInt16' + this.endianness](blockAlign, offset);
  offset += 2;

  // write the bits per sample
  header['writeUInt16' + this.endianness](this.bitDepth, offset);
  offset += 2;

  // write the "data" sub-chunk ID
  data.copy(header, offset);
  offset += data.length;

  // write the remaining length of the rest of the data
  header['writeUInt32' + this.endianness](dataLength, offset);
  offset += 4;

  // save the "header" Buffer for the end, we emit the "header" event at the end
  // with the "size" values properly filled out. if this stream is being piped to
  // a file (or anything else seekable), then this correct header should be placed
  // at the very beginning of the file.
  this._header = header;
  this.headerLength = headerLength;

  // flush the header and after that pass-through "dataLength" bytes
  write(header);
  this._passthrough(dataLength, this._onEnd);
};

/**
 * Called for the "end" event of this Writer instance.
 *
 * @api private
 */

Writer.prototype._onEnd = function (write) {
  debug('_onEnd()');
};

/**
 * Emits a "header" event after the readable side of the stream has finished.
 *
 * @api private
 */

Writer.prototype._flush = function (write, done) {
  debug('_flush()');
  done();
  this.dataLength = this.bytesProcessed;
  process.nextTick(this._emitHeader.bind(this));
};

/**
 * Emits the "header" event. This can safely be ignored, or if you are writing
 * this WAVE file to somewhere that is seekable (i.e. the filesystem), then you
 * should write this "header" buffer at the beginning of the file to get the
 * correct file size values in the file. This isn't too important since most audio
 * players look at the file size rather than those byte values in the header, but
 * it's good to when when possible.
 *
 * @api private
 */

Writer.prototype._emitHeader = function () {
  debug('_emitHeader()');
  var dataLength = this.dataLength;
  var headerLength = this.headerLength;
  var header = this._header;

  // write the file length at the beginning of the header
  header['writeUInt32' + this.endianness](dataLength + headerLength - 8, 4);

  // write the data length at the end of the header
  header['writeUInt32' + this.endianness](dataLength, headerLength - 4);

  this.emit('header', header);
};
<commit_msg>bypassed stream-parser on Writer to avoid assertion error<commit_after>
/**
 * Module dependencies.
 */

var inherits = require('util').inherits;
var debug = require('debug')('wave:writer');
var Transform = require('stream').Transform;

// for node v0.8.x support, remove after v0.12.x
if (!Transform) Transform = require('readable-stream/transform');

/**
 * Module exports.
 */

module.exports = Writer;

/**
 * RIFF Chunk IDs in Buffers.
 *
 * @api private
 */

var RIFF = new Buffer('RIFF');
var WAVE = new Buffer('WAVE');
var fmt  = new Buffer('fmt ');
var data = new Buffer('data');

/**
 * The max size of the "data" chunk of a WAVE file. This is the max unsigned
 * 32-bit int value, minus 100 bytes (overkill, 44 would be safe) for the header.
 *
 * @api private
 */

var MAX_WAV = 4294967295 - 100;

/**
 * The `Writer` class accepts raw audio data written to it (only PCM audio data is
 * currently supported), and outputs a WAV file with a valid WAVE header at the
 * beginning specifying the formatting information of the audio stream.
 *
 * Note that there's an interesting problem, because the WAVE header also
 * specifies the total byte length of the audio data in the file, and there's no
 * way that we can know this ahead of time. Therefore the WAVE header will contain
 * a byte-length if `0` initially, which most WAVE decoders will know means to
 * just read until `EOF`.
 *
 * Optionally, if you are in a situation where you can seek back to the beginning
 * of the destination of the WAVE file (like writing to a regular file, for
 * example), then you may listen for the `"header"` event which will be emitted
 * _after_ all the data has been written, and you can go back and rewrite the new
 * header with proper audio byte length into the beginning of the destination
 * (though if your destination _is_ a regular file, you should use the the
 * `FileWriter` class instead).
 *
 * @param {Object} opts optional options object
 * @api public
 */

function Writer (opts) {
  if (!(this instanceof Writer)) {
    return new Writer(opts);
  }
  Transform.call(this, opts);

  // TODO: allow/properly handle other WAVE audio formats
  this.endianness = 'LE';
  this.format = 1; // raw PCM
  this.channels = 2;
  this.sampleRate = 44100;
  this.bitDepth = 16;
  this.bytesProcessed = 0

  if (opts) {
    if (null != opts.format) this.format = opts.format;
    if (null != opts.channels) this.channels = opts.channels;
    if (null != opts.sampleRate) this.sampleRate = opts.sampleRate;
    if (null != opts.bitDepth) this.bitDepth = opts.bitDepth;
  }

  this._writeHeader()
}
inherits(Writer, Transform);

/**
 * Writes the WAVE header.
 *
 * @api private
 */

Writer.prototype._writeHeader = function () {
  debug('_writeHeader()');
  var headerLength = 44; // TODO: 44 is only for format 1 (PCM), any other
                         // format will have a variable size...
  var dataLength = this.dataLength;
  if (null == dataLength) {
    debug('using default "dataLength" of %d', MAX_WAV);
    dataLength = MAX_WAV;
  }
  var fileSize = dataLength + headerLength;
  var header = new Buffer(headerLength);
  var offset = 0;

  // write the "RIFF" identifier
  RIFF.copy(header, offset);
  offset += RIFF.length;

  // write the file size minus the identifier and this 32-bit int
  header['writeUInt32' + this.endianness](fileSize - 8, offset);
  offset += 4;

  // write the "WAVE" identifier
  WAVE.copy(header, offset);
  offset += WAVE.length;

  // write the "fmt " sub-chunk identifier
  fmt.copy(header, offset);
  offset += fmt.length;

  // write the size of the "fmt " chunk
  // XXX: value of 16 is hard-coded for raw PCM format. other formats have
  // different size.
  header['writeUInt32' + this.endianness](16, offset);
  offset += 4;

  // write the audio format code
  header['writeUInt16' + this.endianness](this.format, offset);
  offset += 2;

  // write the number of channels
  header['writeUInt16' + this.endianness](this.channels, offset);
  offset += 2;

  // write the sample rate
  header['writeUInt32' + this.endianness](this.sampleRate, offset);
  offset += 4;

  // write the byte rate
  var byteRate = this.byteRate;
  if (null == byteRate) {
    byteRate = this.sampleRate * this.channels * this.bitDepth / 8;
  }
  header['writeUInt32' + this.endianness](byteRate, offset);
  offset += 4;

  // write the block align
  var blockAlign = this.blockAlign;
  if (null == blockAlign) {
    blockAlign = this.channels * this.bitDepth / 8;
  }
  header['writeUInt16' + this.endianness](blockAlign, offset);
  offset += 2;

  // write the bits per sample
  header['writeUInt16' + this.endianness](this.bitDepth, offset);
  offset += 2;

  // write the "data" sub-chunk ID
  data.copy(header, offset);
  offset += data.length;

  // write the remaining length of the rest of the data
  header['writeUInt32' + this.endianness](dataLength, offset);
  offset += 4;

  // save the "header" Buffer for the end, we emit the "header" event at the end
  // with the "size" values properly filled out. if this stream is being piped to
  // a file (or anything else seekable), then this correct header should be placed
  // at the very beginning of the file.
  this._header = header;
  this.headerLength = headerLength;

  this.push(header);
};

/**
 * Called for the "end" event of this Writer instance.
 *
 * @api private
 */

Writer.prototype._onEnd = function (write) {
  debug('_onEnd()');
};

/**
 * Transform incoming data. We don't do anything special, just pass it through.
 *
 * @api private
 */

Writer.prototype._transform = function (chunk, enc, done) {
  this.push(chunk)
  this.bytesProcessed += chunk.length
  done()
};

/**
 * Emits a "header" event after the readable side of the stream has finished.
 *
 * @api private
 */

Writer.prototype._flush = function (done) {
  debug('_flush()');
  done();
  this.dataLength = this.bytesProcessed;
  process.nextTick(this._emitHeader.bind(this));
};

/**
 * Emits the "header" event. This can safely be ignored, or if you are writing
 * this WAVE file to somewhere that is seekable (i.e. the filesystem), then you
 * should write this "header" buffer at the beginning of the file to get the
 * correct file size values in the file. This isn't too important since most audio
 * players look at the file size rather than those byte values in the header, but
 * it's good to when when possible.
 *
 * @api private
 */

Writer.prototype._emitHeader = function () {
  debug('_emitHeader()');
  var dataLength = this.dataLength;
  var headerLength = this.headerLength;
  var header = this._header;

  // write the file length at the beginning of the header
  header['writeUInt32' + this.endianness](dataLength + headerLength - 8, 4);

  // write the data length at the end of the header
  header['writeUInt32' + this.endianness](dataLength, headerLength - 4);

  this.emit('header', header);
};
<|endoftext|>


<commit_before>var http = require('http');
var querystring = require('querystring');
var crypto = require('crypto');

var base64 = require('./extern/base64');
var sprintf = require('./extern/sprintf').sprintf;

var OTP = require('./otp').OTP;
var constants = require('./constants');
var exceptions = require('./exceptions');
var utils = require('./utils');

/**
 *
 * @constructor
 * @param {Number} client_id Client ID
 * @param {String} key API key
 * @param {Boolean} use_https True to use HTTPS (defaults to true)
 *
 * Note: You can get client_id and API key at https://upgrade.yubico.com/getapikey/
 */
function Yubico(client_id, key, use_https) {
    this._client_id = client_id;
    this._key = (key !== null) ? base64.decode(key) : null;
    this._use_https = use_https || true;
}

/**
 * Verify the OTP token.
 *
 * @param {Number} timeout Connection timeout in seconds (defaults to 15
 *                         seconds)
 * @param {Function} callback Callback which is called with a possible error
 *                            as the first argument and true as the second
 *                            one if the provided OTP is valid.
 */
Yubico.prototype.verify = function(otp, timeout, callback) {
    var timeout_, callback_, nonce;

    if (typeof timeout === 'function') {
        timeout_ = constants.TIMEOUT * 1000;
        callback_ = timeout;
    }
    else {
        timeout_ = timeout * 1000;
        callback_ = callback;
    }

    this._verify_token(otp, timeout_, function(err, success) {
        if (err) {
            callback_(err, false);
            return;
        }

        callback_(null, true);
    });
};

/**
 * Verify multiple OTP tokens.
 *
 * @param {Array} otp_list Array of OTP tokens.
 * @param {Number} max_time_window How many seconds can pass between the first
 *                 and the last OTP generation.
 * @param {Number} timeout Connection timeout in seconds (defaults to 15
 *                         seconds)
 * @param {Function} callback Callback which is called with a possible error
 *                            as the first argument and true as the second
 *                            one if all the provided OTPs are valid.
 */
Yubico.prototype.verify_multi = function(otp_list, max_time_window, timeout, callback) {
    var self = this;
    var max_time_window_, timeout_, callback_;
    var i, otp_list_length, otp, otp_object, device_id, otps = [], device_ids = [];
    var callback_called = false, result_count = 0;

    if (typeof max_time_window === 'function') {
        max_time_window_ = constants.MAX_TIME_WINDOW;
        timeout_ = (constants.TIMEOUT * 1000);
        callback_ = max_time_window;
    }
    else {
        max_time_window_ = max_time_window || constants.MAX_TIME_WINDOW;
        timeout_ = (timeout * 1000) || (constants.TIMEOUT * 1000);
        callback_ = callback;
    }

    var call_callback = function(err, success) {
        if (callback_called) {
            return;
        }

        callback_called = true;
        callback_(err, success);
    };

    var handle_got_result = function(err, result) {
        if (err) {
            call_callback(err, false);
            return;
        }

        result_count++;
        if (result_count === otp_list_length) {
            call_callback(null, true);
        }
        else {
            self._verify_token(otps.shift(), timeout_, handle_got_result);
        }
    };

    otp_list_length = otp_list.length;
    for (i = 0; i < otp_list_length; i++) {
        otp = otp_list[i];
        otp_object = new OTP(otp);
        device_id = otp_object.get_device_id();

        if (i > 0 && device_ids.indexOf(device_id) === -1) {
            callback_(new Error('OTPs contain different device IDs'), false);
            return;
        }

        device_ids.push(device_id);
        otps.push(otp);
    }

    this._verify_token(otps.shift(), timeout_, handle_got_result);
};

/**
 * Verify the provided token.
 *
 * @param {String} otp OTP token.
 * @param {Number} timeout Number of seconds to wait for sync responses; if
 *                         absent, let the server decides.
 * @param {Function} callback Callback which is called with a possible error
 *                            as the first argument and true as the second
 *                            one if the provided OTP is valid.
 */
Yubico.prototype._verify_token = function(otp, timeout, callback) {
    var self = this;
    var otp_, nonce, query_string;

    var i, client, request, host, port, path, error;
    var timeout_ids = [], client_objects = [], request_objects = [];
    var got_response = false, callback_called = false;
    var client_timed_out_count = 0;

    otp_ = new OTP(otp).get_otp();
    nonce = base64.encode(utils.randstr(40));
    nonce = nonce.substr(0, 30);
    query_string = this._generate_query_string(otp_, nonce, true, 75, (timeout / 1000));

    var call_callback = function(err, success) {
        if (callback_called) {
            return;
        }

        callback_called = true;
        callback(err, success);
    };

    var handle_got_response = function() {
        var i, client_object, request_object, timeout_id;
        got_response = true;

        for (i = 0; i < client_objects.length; i++) {
            client_object = client_objects[i];
            client_object.removeAllListeners('error');

            if (client_object.writable) {
                client_object.destroy();
            }
        }

        for (i = 0; i < request_objects.length; i++) {
            request_object = request_objects[i];
            request_object.removeAllListeners('response');
        }

        for (i = 0; i < timeout_ids.length; i++) {
            timeout_id = timeout_ids[i];
            clearTimeout(timeout_id);
        }
    };

    var clear_connect_timeout = function(timeout_id) {
        clearTimeout(timeout_id);
    };

    var handle_response = function(response, client, timeout_id) {
        var data_buffer = [];

        if (got_response || callback_called) {
            return;
        }

        if (!client.verifyPeer()) {
            // Invalid SSL certificate
            error = exceptions.InvalidSSLCertificateError(client.getPeerCertificate());
            call_callback(error, false);
            return;
        }

        clear_connect_timeout(timeout_id);

        response.on('data', function(chunk) {
            data_buffer.push(chunk);
        });

        response.on('end', function() {
            var response, split, parsed_response, signature, parameters;
            var status, response_otp, generated_signature, error;
            response = data_buffer.join('');

            status = response.match('status=([a-zA-Z0-9_]+)');
            response_otp = response.match('otp=([a-zA-Z0-9]+)');

            if (!status) {
                handle_got_response();

                call_callback(new Error('Missing status attribute'), false);
                return;
            }

            status = status[1].trim().toLowerCase();
            response_otp = response_otp[1].trim();
            if (otp !== response_otp) {
                handle_got_response();

                call_callback(new Error('OTP in the response does not match the' +
                                        ' provided OTP'), false);
                return;
            }

            if (self._key) {
                // Verify the response signature
                parsed_response = self._parse_parameters_from_response(response);
                signature = parsed_response[0];
                parameters = parsed_response[1];
                generated_signature = self._generate_message_signature(parameters, self._key);

                if (signature !== generated_signature) {
                    handle_got_response();

                    error = new exceptions.SignatureVerificationError(generated_signature, signature);
                    call_callback(error, false);
                    return;
                }
            }

            if (status.toLowerCase() === 'ok') {
                handle_got_response();

                call_callback(null, true);
                return;
            }
            else if (status === 'no_such_client') {
                handle_got_response();

                error = new exceptions.InvaliClientError(self._client_id);
                call_callback(error);
                return;
            }
            else if (status === 'replayed_otp' || status === 'bad_otp' ||
                     status === 'bad_signature') {
                handle_got_response();

                error = new exceptions.StatusCodeError(status);
                call_callback(error);
                return;
            }
            else if (status === 'replayed_request') {
                return;
            }

            handle_got_response();
            error = new exceptions.StatusCodeError(status);
            call_callback(error, false);
        });
    };

    var handle_error = function(err) {
        client.removeAllListeners('response');
        call_callback(err, false);
    };

    var handle_timeout = function(client) {
        if (got_response || callback_called) {
            return;
        }

        if (client.writable) {
            client.destroy();
        }

        client_timed_out_count++;
        client.removeAllListeners('response');

        if (client_timed_out_count === client_objects.length) {
            call_callback(new exceptions.ConnectionTimeoutError(timeout));
        }
    };

    port = (this._use_https) ? 443 : 80;
    for (i = 0; i < constants.API_HOSTS.length ; i++) {
        host = constants.API_HOSTS[i];
        path = sprintf('%s?%s', constants.API_PATH, query_string);

        client = http.createClient(port, host, this._use_https);
        request = client.request('GET', path, {
            'host': host,
            'User-Agent': sprintf('NodeJS Yubico library v%s.%s.%s',
                                  constants.VERSION[0], constants.VERSION[1],
                                  constants.VERSION[2])
        });

        request.end();
        client.on('error', handle_error);

        (function(client) {
            var timeout_id;

            timeout_id = setTimeout(function() {
                handle_timeout(client);
            }, timeout);

            request.on('response', function(response) {
                handle_response(response, client, timeout_id);
            });

            timeout_ids.push(timeout_id);
        })(client);

        client_objects.push(client);
        request_objects.push(request);

    }
};

/**
 * Generate a query string which is sent to the validation servers.
 *
 * @param {String} otp Yubikey token.
 * @param {String} nonce  16 to 40 characters long string with random data.
 * @param {String} sl A value 0 to 100 indicating percentage of syncing required
 *                    by client, or strings "fast" or "secure" to use
 *                    server-configured values; if
 *                    absent, let the server decides.
 * @param {Number} timeout Number of seconds to wait for sync responses; if
 *                         absent, let the server decide.
 * @return {String} Generated query string.
 */
Yubico.prototype._generate_query_string = function(otp, nonce, timestamp, sl, timeout) {
    var data, query_string, hmac_signature;

    data = { 'id': this._client_id,
             'otp': otp,
             'nonce': nonce
    };

    if (timestamp) {
        data.timestamp = 1;
    }

    if (sl) {
        if (sl < 0 || sl > 100 && [ 'fast', 'secure' ].indexOf(sl) === -1) {
            throw new Error('sl parameter value must be between 0 and 100 or string "fast" or "secure"');
        }

        data.sl = sl;
    }

    if (timeout) {
        data.timeout = timeout;
    }

    query_string = querystring.encode(data);

    if (this._key) {
        hmac_signature = this._generate_message_signature(query_string, this._key);
        hmac_signature = hmac_signature.replace(/\+/g, '%2B');
        query_string = sprintf('%s&h=%s', query_string, hmac_signature);
    }

    return query_string;
};

/**
 * Parse parameters from the response.
 *
 * @param {String} response Response string
 * @return {Array} Array where the first member is parsed signature and the
 *                second one if the rest of the parsed response returned as
 *                a query string.
 */
Yubico.prototype._parse_parameters_from_response = function(response) {
    var i, split = [], pairs, pair, signature, query_string;

    pairs = response.split('\n');
    for (i = 0; i < pairs.length; i++) {
        pair = pairs[i].trim();

        if (pair === '') {
            continue;
        }

        split.push(pair);
    }

    signature = split[0].replace('h=', '');
    query_string = split.slice(1).join('&');

    return [ signature, query_string ];
};

/**
  * Convert a query string to an object.
 *
 * @param {String} query_string Query string (e.g. foo=bar&bar=baz)
 * @return {Object} Object where the key is a parameter name and the value is
 *                  is the parameter value.
 */
Yubico.prototype._query_string_to_object = function(query_string) {
    var i, pairs, pair, parameters = {};

    pairs = query_string.split('&');
    for (i = 0; i < pairs.length; i++) {
        pair = pairs[i].trim();

        if (pair === '') {
            continue;
        }

        pair = pair.split('=');
        parameters[pair[0]] = pair[1];
    }

    return parameters;
};

/**
 * Generate HMAC-SHA1 signature for the given query string using the
 * provided key.
 *
 * @param {String} query_string Query string
 * @param {String} key Cryptographic key.
 * @return {String} base64 encoded signature.
 */
Yubico.prototype._generate_message_signature = function(query_string, key) {
    var i, pairs, pair, pairs_sorted, pairs_string, hmac, signature;
    pairs = query_string.split('&');

    pairs_sorted = [];
    for (i = 0; i < pairs.length; i++) {
        pair = pairs[i];
        pair = pair.split('=');
        pairs_sorted.push(pair);
    }

    pairs_sorted.sort();

    pairs = [];
    for (i = 0; i < pairs_sorted.length; i++) {
        pair = pairs_sorted[i];
        pairs[i] = pair.join('=');
    }

    pairs_string = pairs.join('&');

    hmac = crypto.createHmac('sha1', key);
    hmac.update(pairs_string);
    signature = hmac.digest('base64');

    return signature;
};

exports.Yubico = Yubico;
<commit_msg>Add return_response argument to the _verify_token method.<commit_after>var http = require('http');
var querystring = require('querystring');
var crypto = require('crypto');

var base64 = require('./extern/base64');
var sprintf = require('./extern/sprintf').sprintf;

var OTP = require('./otp').OTP;
var constants = require('./constants');
var exceptions = require('./exceptions');
var utils = require('./utils');

/**
 *
 * @constructor
 * @param {Number} client_id Client ID
 * @param {String} key API key
 * @param {Boolean} use_https True to use HTTPS (defaults to true)
 *
 * Note: You can get client_id and API key at https://upgrade.yubico.com/getapikey/
 */
function Yubico(client_id, key, use_https) {
    this._client_id = client_id;
    this._key = (key !== null) ? base64.decode(key) : null;
    this._use_https = use_https || true;
}

/**
 * Verify the OTP token.
 *
 * @param {Number} timeout Connection timeout in seconds (defaults to 15
 *                         seconds)
 * @param {Function} callback Callback which is called with a possible error
 *                            as the first argument and true as the second
 *                            one if the provided OTP is valid.
 */
Yubico.prototype.verify = function(otp, timeout, callback) {
    var timeout_, callback_, nonce;

    if (typeof timeout === 'function') {
        timeout_ = constants.TIMEOUT * 1000;
        callback_ = timeout;
    }
    else {
        timeout_ = timeout * 1000;
        callback_ = callback;
    }

    this._verify_token(otp, timeout_, false, function(err, success) {
        if (err) {
            callback_(err, false);
            return;
        }

        callback_(null, true);
    });
};

/**
 * Verify multiple OTP tokens.
 *
 * @param {Array} otp_list Array of OTP tokens.
 * @param {Number} max_time_window How many seconds can pass between the first
 *                 and the last OTP generation.
 * @param {Number} timeout Connection timeout in seconds (defaults to 15
 *                         seconds)
 * @param {Function} callback Callback which is called with a possible error
 *                            as the first argument and true as the second
 *                            one if all the provided OTPs are valid.
 */
Yubico.prototype.verify_multi = function(otp_list, max_time_window, timeout, callback) {
    var self = this;
    var max_time_window_, timeout_, callback_;
    var i, otp_list_length, otp, otp_object, device_id, otps = [], device_ids = [];
    var callback_called = false, result_count = 0;

    if (typeof max_time_window === 'function') {
        max_time_window_ = constants.MAX_TIME_WINDOW;
        timeout_ = (constants.TIMEOUT * 1000);
        callback_ = max_time_window;
    }
    else {
        max_time_window_ = max_time_window || constants.MAX_TIME_WINDOW;
        timeout_ = (timeout * 1000) || (constants.TIMEOUT * 1000);
        callback_ = callback;
    }

    var call_callback = function(err, success) {
        if (callback_called) {
            return;
        }

        callback_called = true;
        callback_(err, success);
    };

    var handle_got_result = function(err, result) {
        if (err) {
            call_callback(err, false);
            return;
        }

        result_count++;
        if (result_count === otp_list_length) {
            call_callback(null, true);
        }
        else {
            self._verify_token(otps.shift(), timeout_, handle_got_result);
        }
    };

    otp_list_length = otp_list.length;
    for (i = 0; i < otp_list_length; i++) {
        otp = otp_list[i];
        otp_object = new OTP(otp);
        device_id = otp_object.get_device_id();

        if (i > 0 && device_ids.indexOf(device_id) === -1) {
            callback_(new Error('OTPs contain different device IDs'), false);
            return;
        }

        device_ids.push(device_id);
        otps.push(otp);
    }

    this._verify_token(otps.shift(), timeout_, handle_got_result);
};

/**
 * Verify the provided token.
 *
 * @param {String} otp OTP token.
 * @param {Number} timeout Number of seconds to wait for sync responses; if
 *                         absent, let the server decides.
 * @param {Boolean} return_response true to call callback with the response
 *                                  string as the last argument.
 * @param {Function} callback Callback which is called with a possible error
 *                            as the first argument, true as the second one
 *                            if the provided OTP is valid and the response
 *                            string as the third one if return_response equals
 *                            true.
 */
Yubico.prototype._verify_token = function(otp, timeout, return_response, callback) {
    var self = this;
    var otp_, nonce, query_string;

    var i, client, request, host, port, path, error;
    var timeout_ids = [], client_objects = [], request_objects = [];
    var got_response = false, callback_called = false;
    var client_timed_out_count = 0;

    otp_ = new OTP(otp).get_otp();
    nonce = base64.encode(utils.randstr(40));
    nonce = nonce.substr(0, 30);
    query_string = this._generate_query_string(otp_, nonce, true, 75, (timeout / 1000));

    var call_callback = function(err, success, response) {
        if (callback_called) {
            return;
        }

        if (!return_response) {
            response = undefined;
        }

        callback_called = true;
        callback(err, success, response);
    };

    var handle_got_response = function() {
        var i, client_object, request_object, timeout_id;
        got_response = true;

        for (i = 0; i < client_objects.length; i++) {
            client_object = client_objects[i];
            client_object.removeAllListeners('error');

            if (client_object.writable) {
                client_object.destroy();
            }
        }

        for (i = 0; i < request_objects.length; i++) {
            request_object = request_objects[i];
            request_object.removeAllListeners('response');
        }

        for (i = 0; i < timeout_ids.length; i++) {
            timeout_id = timeout_ids[i];
            clearTimeout(timeout_id);
        }
    };

    var clear_connect_timeout = function(timeout_id) {
        clearTimeout(timeout_id);
    };

    var handle_response = function(response, client, timeout_id) {
        var data_buffer = [];

        if (got_response || callback_called) {
            return;
        }

        if (!client.verifyPeer()) {
            // Invalid SSL certificate
            error = exceptions.InvalidSSLCertificateError(client.getPeerCertificate());
            call_callback(error, false, null);
            return;
        }

        clear_connect_timeout(timeout_id);

        response.on('data', function(chunk) {
            data_buffer.push(chunk);
        });

        response.on('end', function() {
            var response, split, parsed_response, signature, parameters;
            var status, response_otp, generated_signature, error;
            response = data_buffer.join('');

            status = response.match('status=([a-zA-Z0-9_]+)');
            response_otp = response.match('otp=([a-zA-Z0-9]+)');

            if (!status) {
                handle_got_response();

                call_callback(new Error('Missing status attribute'), false, response);
                return;
            }

            status = status[1].trim().toLowerCase();
            response_otp = response_otp[1].trim();
            if (otp !== response_otp) {
                handle_got_response();

                call_callback(new Error('OTP in the response does not match the' +
                                        ' provided OTP'), false, response);
                return;
            }

            if (self._key) {
                // Verify the response signature
                parsed_response = self._parse_parameters_from_response(response);
                signature = parsed_response[0];
                parameters = parsed_response[1];
                generated_signature = self._generate_message_signature(parameters, self._key);

                if (signature !== generated_signature) {
                    handle_got_response();

                    error = new exceptions.SignatureVerificationError(generated_signature, signature);
                    call_callback(error, false, response);
                    return;
                }
            }

            if (status.toLowerCase() === 'ok') {
                handle_got_response();

                call_callback(null, true, response);
                return;
            }
            else if (status === 'no_such_client') {
                handle_got_response();

                error = new exceptions.InvaliClientError(self._client_id);
                call_callback(error, false, response);
                return;
            }
            else if (status === 'replayed_otp' || status === 'bad_otp' ||
                     status === 'bad_signature') {
                handle_got_response();

                error = new exceptions.StatusCodeError(status);
                call_callback(error, false, response);
                return;
            }
            else if (status === 'replayed_request') {
                return;
            }

            handle_got_response();
            error = new exceptions.StatusCodeError(status);
            call_callback(error, false, response);
        });
    };

    var handle_error = function(err) {
        client.removeAllListeners('response');
        call_callback(err, false);
    };

    var handle_timeout = function(client) {
        if (got_response || callback_called) {
            return;
        }

        if (client.writable) {
            client.destroy();
        }

        client_timed_out_count++;
        client.removeAllListeners('response');

        if (client_timed_out_count === client_objects.length) {
            call_callback(new exceptions.ConnectionTimeoutError(timeout));
        }
    };

    port = (this._use_https) ? 443 : 80;
    for (i = 0; i < constants.API_HOSTS.length ; i++) {
        host = constants.API_HOSTS[i];
        path = sprintf('%s?%s', constants.API_PATH, query_string);

        client = http.createClient(port, host, this._use_https);
        request = client.request('GET', path, {
            'host': host,
            'User-Agent': sprintf('NodeJS Yubico library v%s.%s.%s',
                                  constants.VERSION[0], constants.VERSION[1],
                                  constants.VERSION[2])
        });

        request.end();
        client.on('error', handle_error);

        (function(client) {
            var timeout_id;

            timeout_id = setTimeout(function() {
                handle_timeout(client);
            }, timeout);

            request.on('response', function(response) {
                handle_response(response, client, timeout_id);
            });

            timeout_ids.push(timeout_id);
        })(client);

        client_objects.push(client);
        request_objects.push(request);
    }
};

/**
 * Generate a query string which is sent to the validation servers.
 *
 * @param {String} otp Yubikey token.
 * @param {String} nonce  16 to 40 characters long string with random data.
 * @param {String} sl A value 0 to 100 indicating percentage of syncing required
 *                    by client, or strings "fast" or "secure" to use
 *                    server-configured values; if
 *                    absent, let the server decides.
 * @param {Number} timeout Number of seconds to wait for sync responses; if
 *                         absent, let the server decide.
 * @return {String} Generated query string.
 */
Yubico.prototype._generate_query_string = function(otp, nonce, timestamp, sl, timeout) {
    var data, query_string, hmac_signature;

    data = { 'id': this._client_id,
             'otp': otp,
             'nonce': nonce
    };

    if (timestamp) {
        data.timestamp = 1;
    }

    if (sl) {
        if (sl < 0 || sl > 100 && [ 'fast', 'secure' ].indexOf(sl) === -1) {
            throw new Error('sl parameter value must be between 0 and 100 or string "fast" or "secure"');
        }

        data.sl = sl;
    }

    if (timeout) {
        data.timeout = timeout;
    }

    query_string = querystring.encode(data);

    if (this._key) {
        hmac_signature = this._generate_message_signature(query_string, this._key);
        hmac_signature = hmac_signature.replace(/\+/g, '%2B');
        query_string = sprintf('%s&h=%s', query_string, hmac_signature);
    }

    return query_string;
};

/**
 * Parse parameters from the response.
 *
 * @param {String} response Response string
 * @return {Array} Array where the first member is parsed signature and the
 *                second one if the rest of the parsed response returned as
 *                a query string.
 */
Yubico.prototype._parse_parameters_from_response = function(response) {
    var i, split = [], pairs, pair, signature, query_string;

    pairs = response.split('\n');
    for (i = 0; i < pairs.length; i++) {
        pair = pairs[i].trim();

        if (pair === '') {
            continue;
        }

        split.push(pair);
    }

    signature = split[0].replace('h=', '');
    query_string = split.slice(1).join('&');

    return [ signature, query_string ];
};

/**
  * Convert a query string to an object.
 *
 * @param {String} query_string Query string (e.g. foo=bar&bar=baz)
 * @return {Object} Object where the key is a parameter name and the value is
 *                  is the parameter value.
 */
Yubico.prototype._query_string_to_object = function(query_string) {
    var i, pairs, pair, parameters = {};

    pairs = query_string.split('&');
    for (i = 0; i < pairs.length; i++) {
        pair = pairs[i].trim();

        if (pair === '') {
            continue;
        }

        pair = pair.split('=');
        parameters[pair[0]] = pair[1];
    }

    return parameters;
};

/**
 * Generate HMAC-SHA1 signature for the given query string using the
 * provided key.
 *
 * @param {String} query_string Query string
 * @param {String} key Cryptographic key.
 * @return {String} base64 encoded signature.
 */
Yubico.prototype._generate_message_signature = function(query_string, key) {
    var i, pairs, pair, pairs_sorted, pairs_string, hmac, signature;
    pairs = query_string.split('&');

    pairs_sorted = [];
    for (i = 0; i < pairs.length; i++) {
        pair = pairs[i];
        pair = pair.split('=');
        pairs_sorted.push(pair);
    }

    pairs_sorted.sort();

    pairs = [];
    for (i = 0; i < pairs_sorted.length; i++) {
        pair = pairs_sorted[i];
        pairs[i] = pair.join('=');
    }

    pairs_string = pairs.join('&');

    hmac = crypto.createHmac('sha1', key);
    hmac.update(pairs_string);
    signature = hmac.digest('base64');

    return signature;
};

exports.Yubico = Yubico;
<|endoftext|>