feat: Support for binary data streaming

mladkau · mladkau · commit c29914edfab6 · 2019-10-08T13:41:37.000+01:00
diff --git a/lib/client/socket.js b/lib/client/socket.js
@@ -30,10 +30,12 @@ module.exports = ({ url, token, filters, config, identifyingMetadata }) => {
 
   logger.info({ url }, 'broker client is connecting to broker server');
 
-  const response = relay.response(filters, config);
+  const response = relay.response(filters, config, io);
+  const streamingResponse = relay.streamingResponse;
 
   // RS note: this bind doesn't feel right, it feels like a sloppy way of
   // getting the filters into the request function.
+  io.on('chunk', streamingResponse(token));
   io.on('request', response());
   io.on('error', ({ type, description }) => {
     if (type === 'TransportError') {
@@ -53,7 +55,7 @@ module.exports = ({ url, token, filters, config, identifyingMetadata }) => {
   io.on('close', () => {
     logger.warn({ url, token }, 'websocket connection to the broker server was closed');
   });
-  
+
   // only required if we're manually opening the connection
   // io.open();
 
diff --git a/lib/relay.js b/lib/relay.js
@@ -8,12 +8,41 @@ const replace = require('./replace-vars');
 const tryJSONParse = require('./try-json-parse');
 const logger = require('./log');
 const version = require('./version');
+const stream = require('stream');
 
 module.exports = {
   request: requestHandler,
   response: responseHandler,
+  streamingResponse: streamResponseHandler,
 };
 
+const binaryRequest = [
+  'application/vnd.docker.image.rootfs.diff.tar.gzip',
+];
+
+const streams = new Map();
+
+// Stream processing is only available on
+// the server side i.e. the client sends
+// binary chunks to the server.
+function streamResponseHandler(token) {
+  return (streamingID, chunk, finished) => {
+    var streamBuffer = streams.get(streamingID);
+
+    if (streamBuffer) {
+      if (chunk) {
+        streamBuffer.write(chunk);
+      }
+      if (finished) {
+        streamBuffer.end();
+        streams.delete(streamingID);
+      }
+    } else {
+      logger.warn({ streamingID, token }, 'discarding binary chunk');
+    }
+  };
+}
+
 // 1. Request coming in over HTTP conn (logged)
 // 2. Filter for rule match (log and block if no match)
 // 3. Relay over websocket conn (logged)
@@ -43,6 +72,29 @@ function requestHandler(filterRules) {
 
       req.url = result.url;
       logContext.ioUrl = result.url;
+
+      // check if this is a streaming request for binary data
+      if (req.headers && binaryRequest.indexOf(req.headers['accept']) !== -1) {
+        const streamingID = uuid();
+        const streamBuffer = new stream.PassThrough();
+
+        logContext.streamingID = streamingID;
+        logger.debug(logContext, 'sending binary request over websocket connection');
+
+        streams.set(streamingID, streamBuffer);
+        streamBuffer.pipe(res);
+
+        res.locals.io.send('request', {
+          url: req.url,
+          method: req.method,
+          body: req.body,
+          headers: req.headers,
+          streamingID,
+        });
+
+        return;
+      }
+
       logger.debug(logContext, 'sending request over websocket connection');
 
       // relay the http request over the websocket, handle websocket response
@@ -51,6 +103,7 @@ function requestHandler(filterRules) {
         method: req.method,
         body: req.body,
         headers: req.headers,
+        streamingID : '',
       }, ioResponse => {
         logContext.ioStatus = ioResponse.status;
         logContext.ioHeaders = ioResponse.headers;
@@ -84,15 +137,16 @@ function requestHandler(filterRules) {
 // 3. Relay over HTTP conn (logged)
 // 4. Get response over HTTP conn (logged)
 // 5. Send response over websocket conn
-function responseHandler(filterRules, config) {
+function responseHandler(filterRules, config, io) {
   const filters = Filters(filterRules);
 
-  return (brokerToken) => ({ url, headers = {}, method, body = null } = {}, emit) => {
+  return (brokerToken) => ({ url, headers = {}, method, body = null, streamingID = '' } = {}, emit) => {
     const logContext = {
       url,
       method,
       headers,
       requestId: headers['snyk-request-id'] || uuid(),
+      streamingID,
     };
 
     logger.debug(logContext, 'received request over webscoket connection');
@@ -197,15 +251,31 @@ function responseHandler(filterRules, config) {
 
       logger.debug(logContext, 'sending websocket request over HTTP connection');
 
-      request({
+      var req = {
         url: result.url,
         headers: headers,
         method,
         body,
         agentOptions: {
           ca: config.caCert, // Optional CA cert
         },
-      }, (error, response, body) => {
+      };
+
+      // check if this is a streaming request for binary data
+      if (streamingID) {
+        logger.debug(logContext, 'serving binary stream request');
+
+        req.encoding = null; // no encoding for binary data
+
+        request(req).on('data', chunk => {
+          io.send('chunk', streamingID, chunk, false);
+        }).on('end', () => {
+          io.send('chunk', streamingID, '', true);
+        });
+        return;
+      }
+
+      request(req, (error, response, body) => {
         if (error) {
           logContext.error = error;
           logger.error(logContext, 'error while sending websocket request over HTTP connection');
diff --git a/lib/server/socket.js b/lib/server/socket.js
@@ -8,7 +8,8 @@ module.exports = ({ server, filters, config }) => {
   io.plugin('emitter', Emitter);
 
   const connections = new Map();
-  const response = relay.response(filters, config);
+  const response = relay.response(filters, config, io);
+  const streamingResponse = relay.streamingResponse;
 
   io.on('error', error => logger.error({ error }, 'Primus/engine.io server error'));
 
@@ -57,9 +58,10 @@ module.exports = ({ server, filters, config }) => {
       clientPool.unshift({ socket, metadata: clientData.metadata });
       connections.set(token, clientPool);
 
+      socket.on('chunk', streamingResponse(token));
       socket.on('request', response(token));
     });
-    
+
     ['close', 'end', 'disconnect'].forEach(e => socket.on(e, () => close(e)));
     socket.on('error', (error) => {
       logger.warn({ error }, 'error on websocket connection');
