Refactor decryption & key fetching logic

lib/hls-reader.js

@@ -1,23 +1,21 @@
 "use strict";
 
-var Url = require('url'),
-    Util = require('util'),
-    Crypto = require('crypto');
+var Util = require('util');
 
 var StreamProcess = require('streamprocess'),
-    oncemore = require('oncemore'),
-    UriStream = require('uristream');
+    oncemore = require('oncemore');
 
 var Readable = require('readable-stream/readable'),
     Passthrough = require('readable-stream/passthrough');
 
 var tssmooth = require('./tssmooth');
+var SegmentDecrypt = require('./segment-decrypt');
+
 
 var internals = {
-  keyCache: {},
+  NOOP: function(){},
 };
 
-var NOOP = function(){};
 
 // 'pipe' stream to a Readable
 function pump(src, dst, done) {
@@ -28,7 +26,7 @@ function pump(src, dst, done) {
   });
   oncemore(src).once('end', 'error', function(err) {
     // TODO: flush source buffer on error?
-    dst._read = NOOP;
+    dst._read = internals.NOOP;
     done(err);
   });
   dst._read = function() {
@@ -112,7 +110,11 @@ function HlsReader(segmentReader, options) {
 }
 Util.inherits(HlsReader, Readable);
 
-HlsReader.prototype._read = NOOP;
+HlsReader.prototype._read = internals.NOOP;
+
+HlsReader.prototype.destroy = function () {
+  
+};
 
 // the hook is used to prebuffer
 HlsReader.prototype.hook = function hook() {
@@ -144,56 +146,9 @@ HlsReader.prototype.hook = function hook() {
 };
 
 HlsReader.prototype.decrypt = function (stream, keyAttrs, next) {
-  if (!keyAttrs) return next(null, stream);
-
-  if (keyAttrs.enumeratedString('method') !== 'AES-128' ||
-      !keyAttrs.quotedString('uri') || !keyAttrs.hexadecimalInteger('iv')) {
-
-    // TODO: hard error when key is not recognized?
-    return next(new Error('unknown encryption parameters'));
-  }
-
-  return this.fetchKey(keyAttrs.quotedString('uri'), function(err, key) {
-    if (err)
-      return next(new Error('key fetch failed: ' + (err.stack || err)));
-
-    var iv = keyAttrs.hexadecimalInteger('iv');
-    try {
-      var decrypt = Crypto.createDecipheriv('aes-128-cbc', key, iv);
-    } catch (ex) {
-      return next(new Error('crypto setup failed: ' (ex.stack || ex)));
-    }
-
-    // forward stream errors
-    stream.on('error', function(err) {
-      decrypt.emit('error', err);
-    });
-
-    return next(null, stream.pipe(decrypt));
-  });
+  return SegmentDecrypt.decrypt(stream, keyAttrs, { base: this.reader.baseUrl, key: this.key, cookie: this.cookie }, next);
 };
 
-HlsReader.prototype.fetchKey = function (keyUri, next) {
-  if (this.key) return next(null, this.key);
-
-  var uri = Url.resolve(this.reader.url, keyUri);
-  var entry = internals.keyCache[uri];
-  if (entry && entry.length) return next(null, internals.keyCache[uri]);
-
-  var key = new Buffer(0);
-  var headers = {};
-  if (this.cookie)
-    headers.Cookie = this.cookie;
-
-  oncemore(UriStream(uri, { headers: headers, whitelist: ['http', 'https', 'data'], timeout: 10 * 1000 }))
-    .on('data', function(chunk) {
-      key = Buffer.concat([key, chunk]);
-    })
-    .once('error', 'end', function(err) {
-      internals.keyCache[uri] = key;
-      return next(err, key);
-    });
-};
 
 var hlsreader = module.exports = function hlsreader(segmentReader, options) {
   return new HlsReader(segmentReader, options);

lib/segment-decrypt.js

@@ -0,0 +1,116 @@
+"use strict";
+
+var Url = require('url');
+var Crypto = require('crypto');
+
+var Oncemore = require('oncemore');
+var UriStream = require('uristream');
+
+
+var internals = {
+  allowedProtocols: ['http', 'https', 'data'],
+  fetchTimeout: 10 * 1000,
+
+  keyCache: {},
+};
+
+
+internals.KeyFetcher = function (uri, cookie) {
+  this.uri = uri;
+  this.cookie = cookie;
+  this.key = null;
+
+  this._gets = [];
+};
+
+internals.KeyFetcher.prototype.fetch = function (next) {
+  var key = new Buffer(0);
+  var headers = {};
+  if (this.cookie) {
+    headers.Cookie = this.cookie;
+  }
+
+  Oncemore(UriStream(this.uri, { headers: headers, whitelist: internals.allowedProtocols, timeout: internals.fetchTimeout }))
+    .on('data', function(chunk) {
+      key = Buffer.concat([key, chunk]);
+    })
+    .once('error', 'end', function(err) {
+      return next(err, key);
+    });
+};
+
+internals.KeyFetcher.prototype.get = function (next) {
+  var self = this;
+
+  if (this.key && this.key.length) {
+    return next(null, this.key);
+  }
+
+  var complete = function (err, key) {
+    if (!err) {
+      self.key = key;
+    }
+
+    var gets = self._gets;
+    self._gets = [];
+
+    for (var idx = 0; idx < gets.length; idx++) {
+      var _next = gets[idx];
+      process.nextTick(function() {
+        _next(err, key);
+      });
+    }
+  };
+
+  if (this._gets.length === 0) {
+    this.fetch(complete);
+  }
+
+  return this._gets.push(next);
+};
+
+
+internals.fetchKey = function (keyUri, options, next) {
+  if (options.key) return next(null, options.key);
+
+  var uri = Url.resolve(options.base, keyUri);
+  var fetcher = internals.keyCache[uri];
+  if (!fetcher) {
+    fetcher = internals.keyCache[uri] = new internals.KeyFetcher(uri, options.cookie);
+  }
+  return fetcher.get(next);
+};
+
+
+exports.decrypt = function (stream, keyAttrs, options, next) {
+  var method = keyAttrs && keyAttrs.enumeratedString('method');
+  if (!method || method === 'NONE') {
+    return next(null, stream);
+  }
+
+  if (method !== 'AES-128' || !keyAttrs.quotedString('uri') || !keyAttrs.hexadecimalInteger('iv')) {
+
+    // TODO: hard error when key is not recognized?
+    return next(new Error('unknown encryption parameters'), stream);
+  }
+
+  return internals.fetchKey(keyAttrs.quotedString('uri'), options, function(err, key) {
+    if (err) {
+      return next(new Error('key fetch failed: ' + (err.stack || err)));
+    }
+
+    var iv = keyAttrs.hexadecimalInteger('iv');
+    try {
+      var decrypt = Crypto.createDecipheriv('aes-128-cbc', key, iv);
+    } catch (ex) {
+      return next(new Error('crypto setup failed: ' (ex.stack || ex)));
+    }
+
+    // forward stream errors
+    stream.on('error', function(err) {
+      decrypt.emit('error', err);
+    });
+
+    return next(null, stream.pipe(decrypt));
+  });
+};