whistle/lib/inspectors/data.js

var EventEmitter = require('events').EventEmitter;
var zlib = require('../util/zlib');
var util = require('../util');
var socketMgr = require('../socket-mgr');
var config = require('../config');

var MAX_BODY_SIZE = 360 * 1024;
var MAX_SIZE = (config.strict ? 256 : 1024) * 1024;
var MAX_REQ_BODY_SIZE = (config.strict ? 256 : 2048) * 1024;
var MAX_RES_BODY_SIZE = (config.strict ? 256 : 2048) * 1024;
var BIG_DATA_SIZE = 1024 * 1024 * 16;
var LOCALHOST = '127.0.0.1';
var MAX_PAYLOAD = util.MAX_FRAME_PAYLOAD;
var SSE_SEP = Buffer.from('\n\n');
var FRAME_SEP_KEY = 'x-whistle-custom-frame-separator';

function getZipType(options) {
  return options.headers && options.headers['content-encoding'];
}

function unzipBody(options, body, callback) {
  return zlib.unzip(getZipType(options), util.formatBuffer(body), callback);
}

function checkType(res) {
  if (!config.strict) {
    return true;
  }
  var type = res.headers['content-type'];
  if (!type) {
    return true;
  }
  type = util.getContentType(type);
  return type && type !== 'CSS' && type !== 'IMG';
}

function checkBodySize(data, useBigData) {
  if (
    config.strict &&
    data.body &&
    data.body.length > (useBigData ? BIG_DATA_SIZE : MAX_BODY_SIZE)
  ) {
    data.body = '';
  }
}

function getEventName(proxy) {
  if (util.listenerCount(proxy, '_request')) {
    return '_request';
  } else if (util.listenerCount(proxy, 'httpRequest')) {
    return 'httpRequest';
  }
}

function emitDataEvents(req, res, proxy) {
  var now = Date.now();
  var eventName = getEventName(proxy);
  eventName && proxy.emit(eventName, req.fullUrl);
  if (!util.showPluginReq(req) || util.isHide(req)) {
    return;
  }

  var reqSepInfo = parseFrameSep(req.headers);
  var cseSep = reqSepInfo && reqSepInfo.sep;
  var resSepInfo;

  function emitFrame(chunk, fromServer) {
    proxy.emit('frame', {
      reqId: req.reqId,
      frameId: util.getFrameId(),
      isClient: !fromServer,
      length: chunk.length,
      bin: chunk
    });
  }

  function parseFrameSep(headers) {
    var sep = headers[FRAME_SEP_KEY];
    if (!util.isString(sep)) {
      return;
    }
    var result = {};
    delete headers[FRAME_SEP_KEY];
    try {
      sep = decodeURIComponent(sep);
    } catch (e) {}
    if (sep[0] === '/') {
      result.keepSep = true;
      sep = sep.slice(1);
    }
    if (sep) {
      result.sep = Buffer.from(sep);
      result.sepLen = sep.length;
    }
    return result;
  }

  function parseFrame(buf, isRes, end) {
    var len = buf && buf.length;
    if (!len) {
      return null;
    }
    var sep, sepLen, keepSep;
    var info = isRes ? resSepInfo : reqSepInfo;
    if (info) {
      sep = info.sep;
      sepLen = info.sepLen;
      keepSep = info.keepSep;
    } else {
      sep = SSE_SEP;
      sepLen = 2;
    }
    var index = buf.indexOf(sep);
    while (index !== -1) {
      if (keepSep) {
        index += sepLen;
      }
      index && emitFrame(buf.slice(0, index), isRes);
      buf = buf.slice(keepSep ? index : index + sepLen);
      index = buf.indexOf(sep);
    }
    len = buf.length;
    if (!len) {
      return null;
    }
    if (end) {
      emitFrame(buf, isRes);
      return null;
    }
    if (len > MAX_PAYLOAD) {
      emitFrame(buf.slice(0, -1), isRes);
      buf = buf.slice(-1);
    }
    return buf;
  }
  var getTTFB = function () {
    if (req._ttfb >= now) {
      return req._ttfb - now;
    }
  };
  var _res = {};
  var reqEmitter = new EventEmitter();
  var reqData = {
    method: util.toUpperCase(req.method) || 'GET',
    httpVersion: req.httpVersion || '1.1',
    ip: req.clientIp,
    port: req.clientPort,
    isWhistleHttps: req.isWhistleHttps,
    rawHeaderNames: req.rawHeaderNames,
    headers: req.headers
  };
  var resData = {};
  var reqBody = false;
  var resBody = false;
  var cleared;
  var data = {
    useH2: req.isH2,
    h2Id: req._alpn,
    fc: req.fromComposer ? 1 : undefined,
    isPR: req.isPluginReq ? 1 : undefined,
    testId: req._testId,
    _clientId: req._clientId,
    startTime: now,
    ttfb: getTTFB(),
    id: req.reqId,
    sniPlugin: req.sniPlugin,
    url: req.fullUrl,
    req: reqData,
    res: resData,
    rules: req.rules,
    fwdHost: req._fwdHost,
    pipe: req._pipeRule,
    rulesHeaders: req.rulesHeaders,
    abort: function (clear) {
      if (clear === true) {
        data = reqData = resData = reqBody = resBody = false;
        cleared = true;
      } else {
        var err = new Error('Aborted');
        err.code = 'ERR_WHISTLE_ABORTED';
        req.emit('error', err);
      }
    }
  };
  proxy.emit('request', reqEmitter, data);

  var requestTime;
  var endTime;
  var updateEvent;
  var useFrames;
  var enable = req.enable;
  var disable = req.disable;
  var useBigData =
    (enable.bigData || enable.largeData) &&
    !disable.bigData &&
    !disable.largeData;
  var setRequestTime = function (time) {
    data.requestTime = time;
    data.connectTime = req._connectTime;
    reqData.size = reqData.size || 0;
  };
  var setEndTime = function () {
    if (data.requestTime || !requestTime) {
      data.endTime = endTime || Date.now();
      setRequestTime(data.requestTime || data.endTime);
    } else if (endTime == null) {
      endTime = endTime || Date.now();
    }
    resData.size = resData.size || 0;
  };
  var updateVersion = function () {
    if (req.useH2 != null) {
      data.useH2 = req.useH2;
    }
    data.httpsTime = req.httpsTime;
    data.useHttp = req.useHttp;
    if (data.useH2 && !req.useHttp) {
      reqData.httpVersion = '2.0';
    } else {
      reqData.httpVersion = '1.1';
    }
  };
  var setUnzipSize = function (body, obj) {
    var len = body ? body.length : -1;
    if (len >= 0 && len !== obj.size) {
      obj.unzipSize = len;
    }
  };
  req.setServerPort = function (serverPort) {
    req.serverPort = serverPort;
    setReqStatus(LOCALHOST);
    resData.port = serverPort;
  };
  req.setClientId = function (clientId) {
    data.clientId = clientId;
  };
  req.setTTFB = function() {
    data.ttfb = getTTFB();
  };
  res.setCurTrailers = function (trailers, rawTrailerNames) {
    resData.trailers = trailers;
    resData.rawTrailerNames = rawTrailerNames;
  };
  var reqDone;
  var resDone;
  var captureReqStream;
  var maxReqSize;
  var reqInfo;
  var setReqBody = function (end) {
    reqBody = util.formatBuffer(reqBody);
    if (captureReqStream && cseSep) {
      reqBody = parseFrame(reqBody, false, end);
    } else {
      reqData.body = reqBody;
    }
  };
  var updateReqInfo = function() {
    if (reqInfo && !captureReqStream) {
      captureReqStream = !disable.captureStream && (resDone || enable.captureStream) && !getZipType(reqInfo);
      maxReqSize = useBigData ? BIG_DATA_SIZE : (captureReqStream ? MAX_REQ_BODY_SIZE : MAX_SIZE);
      if (captureReqStream) {
        data.isCse = !!cseSep;
        reqBody && !reqData.body && setReqBody();
      }
    }
  };

  var handleReqBody = function (stream, info) {
    if (reqDone) {
      return;
    }
    reqDone = true;
    reqInfo = info || (req._needGunzip ? { method: req.method, headers: '' } : req);
    if (!cleared && util.hasRequestBody(reqInfo)) {
      reqBody = null;
    }
    reqData.size = 0;
    updateReqInfo();
    var write = stream.write;
    var end = stream.end;
    stream.write = function (chunk) {
      if (chunk) {
        reqData.size += chunk.length;
        if (reqBody || reqBody === null) {
          if (useFrames) {
            reqBody = '';
          } else {
            reqBody = util.joinBuffer(reqBody, chunk);
            if (captureReqStream) {
              setReqBody();
              if (reqBody && reqBody.length > maxReqSize) {
                reqBody = false;
              }
            } else if (reqData.size > maxReqSize) {
              reqBody = false;
            }
          }
        }
      }
      updateEvent && proxy.emit(updateEvent, req.fullUrl, false);
      return write.apply(this, arguments);
    };
    stream.end = function (chunk) {
      requestTime = Date.now();
      if (useFrames) {
        reqBody = '';
      } else if (chunk) {
        reqData.size += chunk.length;
        reqBody = util.joinBuffer(reqBody, chunk);
      }
      unzipBody(reqInfo, reqBody, function (err, body) {
        setRequestTime(requestTime);
        if (endTime) {
          data.endTime = endTime;
        }
        reqBody = err ? util.getErrorStack(err) : body;
        setReqBody(true);
        !data.isCse && setUnzipSize(body, reqData);
        checkBodySize(reqData, useBigData);
      });
      updateEvent && proxy.emit(updateEvent, req.fullUrl, false, true);
      return end.apply(this, arguments);
    };
  };
  var handleResBody = function (stream, info) {
    if (resDone) {
      return;
    }
    resDone = true;
    info = info || (res._needGunzip ? { statusCode: res.statusCode, headers: '' } : res);
    var isSse;
    var captureStream;
    if (!disable.captureStream && !getZipType(info)) {
      var headers = info.headers;
      captureStream = requestTime == null || enable.captureStream;
      if (headers) {
        isSse = util.trimType(headers['content-type']) === 'text/event-stream';
        captureStream = captureStream || isSse;
        resSepInfo = parseFrameSep(headers);
        if (resSepInfo && resSepInfo.sep) {
          isSse = captureStream;
        }
      }
    }
    var maxResSize = useBigData ? BIG_DATA_SIZE : (captureStream ? MAX_RES_BODY_SIZE : MAX_SIZE);
    var setResBody = function (end) {
      resBody = util.formatBuffer(resBody);
      if (isSse) {
        resBody = parseFrame(resBody, true, end);
      } else {
        resData.body = resBody;
      }
    };
    data.isSse = isSse;
    captureStream && updateReqInfo();
    if (!cleared && util.hasBody(info, req) && checkType(info)) {
      if (info.headers['content-length'] > maxResSize && !captureStream) {
        resBody = false;
      } else {
        resBody = null;
      }
    }
    resData.size = 0;
    var write = stream.write;
    var end = stream.end;
    stream.write = function (chunk) {
      if (chunk) {
        resData.size += chunk.length;
        if (resBody || resBody === null) {
          if (useFrames) {
            resBody = '';
          } else {
            resBody = util.joinBuffer(resBody, chunk);
            if (captureStream) {
              setResBody();
              if (resBody && resBody.length > maxResSize) {
                resBody = false;
              }
            } else if (resData.size > maxResSize) {
              resBody = false;
            }
          }
        }
      }
      updateEvent && proxy.emit(updateEvent, req.fullUrl, true);
      return write.apply(this, arguments);
    };
    stream.end = function (chunk) {
      endTime = Date.now();
      delete data.abort;
      if (useFrames) {
        resBody = '';
      } else if (chunk) {
        resData.size += chunk.length;
        resBody = util.joinBuffer(resBody, chunk);
      }
      resData.hasGzipError = unzipBody(info, resBody, function (err, body) {
        if (!data.resError) {
          resBody = err ? util.getErrorStack(err) : body;
          setResBody(true);
        }
        !isSse && setUnzipSize(body, resData);
        checkBodySize(resData, useBigData);
        setEndTime();
        reqEmitter.emit('end', data);
      });
      updateEvent && proxy.emit(updateEvent, req.fullUrl, true, true);
      return end.apply(this, arguments);
    };
  };

  var hasReqPipe = req._pipePluginPorts.reqWritePort;
  var hasResPipe = req._pipePluginPorts.resWritePort;

  if (hasReqPipe) {
    req.on('bodyStreamReady', handleReqBody);
  }

  if (hasResPipe) {
    res.on('bodyStreamReady', handleResBody);
  }

  req.once('dest', function (_req) {
    util.onWriteEnd(_req, function () {
      if (!requestTime) {
        setRequestTime(Date.now());
      }
    });
    setReqStatus();
    reqEmitter.emit('send', data);
    !hasReqPipe && handleReqBody(_req, req);
  });
  res.once('src', function (r) {
    _res = r;
    data.pipe = req._pipeRule;
    resData.rawHeaderNames = res.rawHeaderNames;
    if (!data.endTime) {
      setResStatus();
      reqEmitter.emit('response', data);
    }
    !hasResPipe && handleResBody(res, _res);
  });
  req.once('error', handleError);
  res.once('error', handleError);
  res.once('close', handleError);
  util.onWriteEnd(res, setEndTime);

  function handleError(err) {
    req._hasError = true;
    util.clearDelayTimer(req);
    if (!data || endTime || data.endTime) {
      return;
    }
    // 触发 endTime
    !endTime && setEndTime();
    if (data.responseTime && !err) {
      return;
    }
    delete data.abort;
    // 不能乱优化，跟 var msg = err && err.message; if (msg !== 'Aborted' && msg !== 'aborted') 不一样
    if (err && err.message !== 'Aborted' && err.message !== 'aborted') {
      data.resError = true;
      if (resData.body) {
        data.errMsg = util.getErrorStack(err);
      } else {
        resData.body = util.getErrorStack(err);
      }
      util.emitError(reqEmitter, data);
      setResStatus(502);
    } else {
      data.reqError = true;
      if (!reqData.body) {
        reqData.body = 'aborted';
      } else if (!resData.body) {
        resData.body = 'aborted';
      } else {
        data.errMsg = 'aborted';
      }
      if (req.__resHeaders) {
        _res.headers = req.__resHeaders;
      }
      reqEmitter.emit('abort', data);
      setResStatus('aborted' + (req.__statusCode ? ' (' + req.__statusCode + ')' : ''));
    }
  }

  function setReqStatus(defaultHost) {
    data.dnsTime = (req.dnsTime || 0) + now;
    data.realUrl = data.url === req.realUrl ? undefined : req.realUrl;
    updateVersion();
    resData.ip = req.hostIp || defaultHost;
    resData.phost = req._phost && req._phost.host;
  }

  req.initCustomParser = function () {
    if (req.customParser) {
      data.useFrames = false;
      socketMgr.setPending(req);
      req.disableCustomParser = function () {
        data.useFrames = null;
        req.customParser = null;
        req.enableCustomParser = null;
        req.disableCustomParser = null;
        socketMgr.removePending(req);
        delete req.headers['x-whistle-frame-parser'];
      };
      req.enableCustomParser = function (svrRes) {
        useFrames = true;
        reqData.body = '';
        resData.body = '';
        reqData.base64 = '';
        resData.base64 = '';
        data.useFrames = true;
        socketMgr.setContext(
          req,
          svrRes,
          eventName,
          { data: [] },
          { data: [] }
        );
        socketMgr.removePending(req);
        delete req.headers['x-whistle-frame-parser'];
      };
    }
  };

  function setResStatus(defaultCode) {
    if (data.responseTime) {
      return;
    }
    setReqStatus(LOCALHOST);
    resData.statusCode = _res.statusCode || defaultCode || 502;
    resData.statusMessage = _res.statusMessage;
    data.responseTime = Date.now();
    if (!requestTime && !data.requestTime) {
      updateEvent = eventName;
    }
    resData.headers = _res.headers;
  }
}

module.exports = function (req, res, next) {
  emitDataEvents(req, res, this);
  var reqDelay = util.getMatcherValue(req.rules.reqDelay);
  util.setDelayTimer(req, util.delay(reqDelay, function () {
    if (util.needAbortReq(req)) {
      return res.destroy();
    }
    next();
  }));
};