From c5cda3043d8ade5cc98a08cd8ea1a2bc321b82f3 Mon Sep 17 00:00:00 2001
From: dignifiedquire <dignifiedquire@gmail.com>
Date: Thu, 17 Aug 2017 00:44:21 +0200
Subject: [PATCH] reduce buffer creation for ctr mode

---
 aes.js         |  8 ++++++--
 bench/index.js |  2 +-
 modes/ctr.js   | 18 +++++++++++++++---
 3 files changed, 22 insertions(+), 6 deletions(-)

diff --git a/aes.js b/aes.js
index 4d60109..ca32ab7 100644
--- a/aes.js
+++ b/aes.js
@@ -187,9 +187,13 @@ AES.prototype._reset = function () {
   this._invKeySchedule = invKeySchedule
 }
 
-AES.prototype.encryptBlock = function (M) {
+AES.prototype.encryptBlockRaw = function (M) {
   M = asUInt32Array(M)
-  var out = cryptBlock(M, this._keySchedule, G.SUB_MIX, G.SBOX, this._nRounds)
+  return cryptBlock(M, this._keySchedule, G.SUB_MIX, G.SBOX, this._nRounds)
+}
+
+AES.prototype.encryptBlock = function (M) {
+  var out = this.encryptBlockRaw(M)
   var buf = Buffer.allocUnsafe(16)
   buf.writeUInt32BE(out[0], 0)
   buf.writeUInt32BE(out[1], 4)
diff --git a/bench/index.js b/bench/index.js
index 172d3cc..4888fa4 100644
--- a/bench/index.js
+++ b/bench/index.js
@@ -5,7 +5,7 @@ let key = Buffer.alloc(16, 0xff)
 let iv = Buffer.alloc(16, 0x01)
 
 function test (mod, message) {
-  let cipher = mod.createCipheriv('aes-128-cbc', key, iv)
+  let cipher = mod.createCipheriv('aes-128-ctr', key, iv)
   let b = cipher.update(message)
   return Buffer.concat([b, cipher.final()])
 }
diff --git a/modes/ctr.js b/modes/ctr.js
index 0ef2278..ca282fb 100644
--- a/modes/ctr.js
+++ b/modes/ctr.js
@@ -16,14 +16,26 @@ function incr32 (iv) {
 }
 
 function getBlock (self) {
-  var out = self._cipher.encryptBlock(self._prev)
+  var out = self._cipher.encryptBlockRaw(self._prev)
   incr32(self._prev)
   return out
 }
 
+var blockSize = 16
 exports.encrypt = function (self, chunk) {
-  while (self._cache.length < chunk.length) {
-    self._cache = Buffer.concat([self._cache, getBlock(self)])
+  var chunkNum = Math.ceil(chunk.length / blockSize)
+  var start = self._cache.length
+  self._cache = Buffer.concat([
+    self._cache,
+    Buffer.allocUnsafe(chunkNum * blockSize)
+  ])
+  for (var i = 0; i < chunkNum; i++) {
+    var out = getBlock(self)
+    var offset = start + i * blockSize
+    self._cache.writeUInt32BE(out[0], offset + 0)
+    self._cache.writeUInt32BE(out[1], offset + 4)
+    self._cache.writeUInt32BE(out[2], offset + 8)
+    self._cache.writeUInt32BE(out[3], offset + 12)
   }
   var pad = self._cache.slice(0, chunk.length)
   self._cache = self._cache.slice(chunk.length)