refactor aes.

lib/bcoin/aes.js

@@ -11,23 +11,26 @@
 /* jshint latedef: false */
 
 var assert = require('assert');
+var U32Array = typeof Uint32Array === 'function' ? Uint32Array : Array;
+var AES = exports;
 
 /**
- * An AES object for encrypting and decrypting blocks.
- * @exports AES
+ * An AES key object for encrypting
+ * and decrypting blocks.
+ * @exports AESKey
  * @constructor
  * @param {Buffer} key
  * @param {Number} bits
  * @param {Buffer?} iv
  */
 
-function AES(key, bits, iv) {
-  if (!(this instanceof AES))
-    return new AES(key, bits, iv);
+function AESKey(key, bits) {
+  if (!(this instanceof AESKey))
+    return new AESKey(key, bits);
 
+  this.rounds = null;
   this.userKey = key;
   this.bits = bits;
-  this.iv = iv;
 
   if (this.bits === 128)
     this.rounds = 10;
@@ -46,19 +49,14 @@ function AES(key, bits, iv) {
  * Destroy the object and zero the keys.
  */
 
-AES.prototype.destroy = function destroy() {
+AESKey.prototype.destroy = function destroy() {
   var i;
 
   assert(this.userKey, 'Already destroyed.');
 
-  this.userKey.fill(0);
+  // User should zero this.
   this.userKey = null;
 
-  if (this.iv) {
-    this.iv.fill(0);
-    this.iv = null;
-  }
-
   if (this.decryptKey) {
     for (i = 0; i < this.decryptKey.length; i++)
       this.decryptKey[i] = 0;
@@ -74,19 +72,19 @@ AES.prototype.destroy = function destroy() {
 
 /**
  * Convert the user key into an encryption key.
- * @returns {Object} key
+ * @returns {Uint32Array} key
  */
 
-AES.prototype.getEncryptKey = function getEncryptKey() {
+AESKey.prototype.getEncryptKey = function getEncryptKey() {
   var i = 0;
   var key, kp, tmp;
 
-  assert(this.userKey, 'Cannot use AES once it is destroyed.');
+  assert(this.userKey, 'Cannot use key once it is destroyed.');
 
   if (this.encryptKey)
     return this.encryptKey;
 
-  key = [];
+  key = new U32Array(60);
   kp = 0;
 
   key[kp + 0] = readU32(this.userKey, 0);
@@ -183,21 +181,25 @@ AES.prototype.getEncryptKey = function getEncryptKey() {
 
 /**
  * Convert the user key into a decryption key.
- * @returns {Object} key
+ * @returns {Uint32Array} key
  */
 
-AES.prototype.getDecryptKey = function getDecryptKey() {
-  var i, j, kp, key, tmp;
+AESKey.prototype.getDecryptKey = function getDecryptKey() {
+  var i, j, kp, enc, key, tmp;
 
-  assert(this.userKey, 'Cannot use AES once it is destroyed.');
+  assert(this.userKey, 'Cannot use key once it is destroyed.');
 
   if (this.decryptKey)
     return this.decryptKey;
 
   // First, start with an encryption schedule.
-  key = this.getEncryptKey().slice();
+  enc = this.getEncryptKey();
+  key = new U32Array(60);
   kp = 0;
 
+  for (i = 0; i < enc.length; i++)
+    key[i] = enc[i];
+
   this.decryptKey = key;
 
   // Invert the order of the round keys.
@@ -250,10 +252,10 @@ AES.prototype.getDecryptKey = function getDecryptKey() {
  * @returns {Buffer}
  */
 
-AES.prototype.encryptBlock = function encryptBlock(input) {
+AESKey.prototype.encryptBlock = function encryptBlock(input) {
   var output, kp, key, r, s0, s1, s2, s3, t0, t1, t2, t3;
 
-  assert(this.userKey, 'Cannot use AES once it is destroyed.');
+  assert(this.userKey, 'Cannot use key once it is destroyed.');
 
   key = this.getEncryptKey();
   kp = 0;
@@ -355,10 +357,10 @@ AES.prototype.encryptBlock = function encryptBlock(input) {
  * @returns {Buffer}
  */
 
-AES.prototype.decryptBlock = function decryptBlock(input) {
+AESKey.prototype.decryptBlock = function decryptBlock(input) {
   var output, kp, key, r, s0, s1, s2, s3, t0, t1, t2, t3;
 
-  assert(this.userKey, 'Cannot use AES once it is destroyed.');
+  assert(this.userKey, 'Cannot use AESKey once it is destroyed.');
 
   key = this.getDecryptKey();
   kp = 0;
@@ -455,40 +457,75 @@ AES.prototype.decryptBlock = function decryptBlock(input) {
 };
 
 /**
- * Encrypt data with aes 256.
- * @param {Buffer} data
+ * AES cipher.
+ * @exports AESCipher
+ * @constructor
  * @param {Buffer} key
  * @param {Buffer} iv
- * @param {Boolean?} chain - XOR chaining (cbc).
+ * @param {Number} bits
+ * @param {String} mode
+ */
+
+function AESCipher(key, iv, bits, mode) {
+  if (!(this instanceof AESCipher))
+    return new AESCipher(key, iv, mode);
+
+  assert(mode === 'ecb' || mode === 'cbc', 'Unknown mode.');
+
+  this.key = new AESKey(key, bits);
+  this.mode = mode;
+  this.prev = iv;
+  this.waiting = null;
+}
+
+/**
+ * Encrypt blocks of data.
+ * @param {Buffer} data
  * @returns {Buffer}
  */
 
-AES.encrypt = function encrypt(data, key, iv, chain) {
-  var aes = new AES(key, 256, iv);
-  var trailing = data.length % 16;
-  var len = data.length - trailing;
+AESCipher.prototype.update = function update(data) {
   var blocks = [];
-  var i, prev, block, left, pad;
+  var i, len, trailing, block;
 
-  // Setup initialization vector.
-  prev = aes.iv;
+  if (this.waiting) {
+    data = Buffer.concat([this.waiting, data]);
+    this.waiting = null;
+  }
+
+  trailing = data.length % 16;
+  len = data.length - trailing;
 
   // Encrypt all blocks except for the last.
   for (i = 0; i < len; i += 16) {
     block = data.slice(i, i + 16);
-    if (chain)
-      block = xor(block, prev);
-    prev = aes.encryptBlock(block);
-    blocks.push(prev);
+    if (this.mode === 'cbc')
+      block = xor(block, this.prev);
+    this.prev = this.key.encryptBlock(block);
+    blocks.push(this.prev);
   }
 
+  if (trailing > 0)
+    this.waiting = data.slice(len);
+
+  return Buffer.concat(blocks);
+};
+
+/**
+ * Finalize the cipher.
+ * @returns {Buffer}
+ */
+
+AESCipher.prototype.final = function final() {
+  var i, block, left, pad;
+
   // Handle padding on the last block.
-  if (trailing === 0) {
+  if (!this.waiting) {
     block = new Buffer(16);
     block.fill(16);
   } else {
-    block = data.slice(len, len + trailing);
-    left = 16 - trailing;
+    block = this.waiting;
+    left = 16 - block.length;
     pad = new Buffer(left);
     pad.fill(left);
     block = Buffer.concat([block, pad]);
@@ -496,70 +533,150 @@ AES.encrypt = function encrypt(data, key, iv, chain) {
 
   // Encrypt the last block,
   // as well as the padding.
-  if (chain)
-    block = xor(block, prev);
+  if (this.mode === 'cbc')
+    block = xor(block, this.prev);
 
-  block = aes.encryptBlock(block);
-  blocks.push(block);
+  block = this.key.encryptBlock(block);
 
-  aes.destroy();
+  this.key.destroy();
+
+  return block;
+};
+
+/**
+ * AES decipher.
+ * @exports AESDecipher
+ * @constructor
+ * @param {Buffer} key
+ * @param {Buffer} iv
+ * @param {Number} bits
+ * @param {String} mode
+ */
+
+function AESDecipher(key, iv, bits, mode) {
+  if (!(this instanceof AESDecipher))
+    return new AESDecipher(key, iv, mode);
+
+  assert(mode === 'ecb' || mode === 'cbc', 'Unknown mode.');
+
+  this.key = new AESKey(key, bits);
+  this.mode = mode;
+  this.prev = iv;
+  this.waiting = null;
+  this.lastBlock = null;
+}
+
+/**
+ * Decrypt blocks of data.
+ * @param {Buffer} data
+ */
+
+AESDecipher.prototype.update = function update(data) {
+  var blocks = [];
+  var i, chunk, block, len, trailing;
+
+  if (this.waiting) {
+    data = Buffer.concat([this.waiting, data]);
+    this.waiting = null;
+  }
+
+  trailing = data.length % 16;
+  len = data.length - trailing;
+
+  // Decrypt all blocks.
+  for (i = 0; i < len; i += 16) {
+    chunk = this.prev;
+    this.prev = data.slice(i, i + 16);
+    block = this.key.decryptBlock(this.prev);
+    if (this.mode === 'cbc')
+      block = xor(block, chunk);
+    blocks.push(block);
+  }
+
+  if (trailing > 0)
+    this.waiting = data.slice(len);
+
+  if (this.lastBlock) {
+    blocks.unshift(this.lastBlock);
+    this.lastBlock = null;
+  }
+
+  // Keep a reference to the last
+  // block for the padding check.
+  this.lastBlock = blocks.pop();
 
   return Buffer.concat(blocks);
 };
 
+/**
+ * Finalize the decipher.
+ * @returns {Buffer}
+ */
+
+AESDecipher.prototype.final = function final() {
+  var i, b, n, block;
+
+  assert(!this.waiting, 'Bad decrypt (trailing bytes).');
+  assert(this.lastBlock, 'Bad decrypt (no data).');
+
+  // Check padding on the last block.
+  block = this.lastBlock;
+  b = 16;
+  n = block[b - 1];
+
+  if (n === 0 || n > b)
+    throw new Error('Bad decrypt (padding).');
+
+  for (i = 0; i < n; i++) {
+    if (block[--b] !== n)
+      throw new Error('Bad decrypt (padding).');
+  }
+
+  // Slice off the padding unless
+  // the entire block was padding.
+  if (n === 16)
+    return new Buffer(0);
+
+  block = block.slice(0, -n);
+
+  this.key.destroy();
+
+  return block;
+};
+
+/**
+ * Encrypt data with aes 256.
+ * @param {Buffer} data
+ * @param {Buffer} key
+ * @param {Buffer} iv
+ * @param {String} mode
+ * @returns {Buffer}
+ */
+
+AES.encrypt = function encrypt(data, key, iv, bits, mode) {
+  var cipher = new AESCipher(key, iv, bits, mode);
+  return Buffer.concat([
+    cipher.update(data),
+    cipher.final()
+  ]);
+};
+
 /**
  * Decrypt data with aes 256.
  * @param {Buffer} data
  * @param {Buffer} key
- * @param {Buffer} iv
- * @param {Boolean?} chain - XOR chaining (cbc).
+ * @param {Buffer|null} iv
+ * @param {Number} bits
+ * @param {String} mode
  * @returns {Buffer}
  */
 
-AES.decrypt = function decrypt(data, key, iv, chain) {
-  var aes = new AES(key, 256, iv);
-  var blocks = [];
-  var i, b, n, prev, chunk, block;
-
-  assert(data.length > 0);
-  assert(data.length % 16 === 0);
-
-  // Setup initialization vector.
-  prev = aes.iv;
-
-  // Decrypt all blocks.
-  for (i = 0; i < data.length; i += 16) {
-    chunk = prev;
-    prev = data.slice(i, i + 16);
-    block = aes.decryptBlock(prev);
-    if (chain)
-      block = xor(block, chunk);
-    blocks.push(block);
-  }
-
-  // Check padding on the last block.
-  block = blocks.pop();
-  b = 16;
-  n = block[b - 1];
-
-  if (n === 0 || n > b)
-    throw new Error('Bad decrypt');
-
-  for (i = 0; i < n; i++) {
-    if (block[--b] !== n)
-      throw new Error('Bad decrypt');
-  }
-
-  // Slice off the padding unless
-  // the entire block was padding.
-  if (n < 16) {
-    block = block.slice(0, -n);
-    blocks.push(block);
-  }
-
-  aes.destroy();
-
-  return Buffer.concat(blocks);
+AES.decrypt = function decrypt(data, key, iv, bits, mode) {
+  var decipher = new AESDecipher(key, iv, bits, mode);
+  return Buffer.concat([
+    decipher.update(data),
+    decipher.final()
+  ]);
 };
 
 /**
@@ -576,7 +693,7 @@ AES.ecb = {};
  */
 
 AES.ecb.encrypt = function encrypt(data, key) {
-  return AES.encrypt(data, key, null, false);
+  return AES.encrypt(data, key, null, 256, 'ecb');
 };
 
 /**
@@ -587,7 +704,7 @@ AES.ecb.encrypt = function encrypt(data, key) {
  */
 
 AES.ecb.decrypt = function decrypt(data, key) {
-  return AES.decrypt(data, key, null, false);
+  return AES.decrypt(data, key, null, 256, 'ecb');
 };
 
 /**
@@ -605,7 +722,7 @@ AES.cbc = {};
  */
 
 AES.cbc.encrypt = function encrypt(data, key, iv) {
-  return AES.encrypt(data, key, iv, true);
+  return AES.encrypt(data, key, iv, 256, 'cbc');
 };
 
 /**
@@ -617,9 +734,17 @@ AES.cbc.encrypt = function encrypt(data, key, iv) {
  */
 
 AES.cbc.decrypt = function decrypt(data, key, iv) {
-  return AES.decrypt(data, key, iv, true);
+  return AES.decrypt(data, key, iv, 256, 'cbc');
 };
 
+/*
+ * Expose
+ */
+
+AES.Key = AESKey;
+AES.Cipher = AESCipher;
+AES.Decipher = AESDecipher;
+
 /*
  * Helpers
  */
@@ -1225,9 +1350,3 @@ var RCON = [
   0x10000000, 0x20000000, 0x40000000, 0x80000000,
   0x1B000000, 0x36000000
 ];
-
-/*
- * Expose
- */
-
-module.exports = AES;