whistle/lib/socket-mgr.js

var PassThrough = require('stream').PassThrough;
var pluginMgr = require('./plugins');
var wsParser = require('ws-parser');
var util = require('./util');
var config = require('./config');
var Buffer = require('safe-buffer').Buffer;

var pendingReqList = [];
var INTERVAL = 22 * 1000;
var proxy;
var index = 0;
var MAX_PAYLOAD = 1024 * 256;
var conns = {};
var PING = new Buffer.from('iQA=', 'base64');
var PONG = new Buffer.from('ioAn6ubf', 'base64');
var PAUSE_STATUS = 1;
var IGNORE_STATUS = 2;
var MAX_COMPOSE_FRAME_COUNT = 5;

function getFrameId() {
  ++index;
  if (index > 999) {
    index = 0;
  }
  if (index > 99) {
    return Date.now() + '-' + index;
  }
  if (index > 9) {
    return Date.now() + '-0' + index;
  }
  return Date.now() + '-00' + index;
}

exports = module.exports = function(p) {
  proxy = p;
};

function handleSocketEnd(req, res, callback) {
  util.onSocketEnd(req, callback);
  util.onSocketEnd(res, callback);
}

function handleClose(req, res, justTunnel) {
  handleSocketEnd(req, res, function(err) {
    var ctx = conns[req.reqId];
    ctx && ctx.clearup();
    var closed = req._hasClosed;
    req._hasClosed = true;
    // 确保两个连接都关掉才行
    if (closed) {
      req.emit('_closed');
    }
    if (closed || req.customParser || justTunnel) {
      return;
    }
    proxy.emit('frame', {
      reqId:  req.reqId,
      frameId: getFrameId(),
      closed: true,
      code: req._errorCode || res._errorCode,
      err: err && err.message
    });
  });
}

function getStatus(ctx, status, name) {
  status = parseInt(status, 10);
  name = name || 'receiveStatus';
  var oldStatus = ctx[name] || 0;
  status = ctx[name] = (status > 0 || status < 3) ? status : 0;
  return status !== oldStatus ? status : -1;
}

function setConnStatus(ctx, status, statusObj, name) {
  statusObj.pause = statusObj.ignore = undefined;
  status = getStatus(ctx, status, name);
  if (status === 1) {
    statusObj.pause = true;
    return;
  }
  if (status === 2) {
    statusObj.ignore = true;
    statusObj.chunk = null;
    statusObj.ignoring = !!statusObj.callback;
  }
  statusObj.emitData && statusObj.emitData();
  if (statusObj.callback) {
    statusObj.addToReceiver && statusObj.addToReceiver();
    statusObj.callback(null, statusObj.chunk);
    statusObj.addToReceiver = null;
    statusObj.callback = null;
    statusObj.chunk = null;
  }
}

function initStatus(ctx, enable) {
  if (enable.pauseSend) {
    ctx.setSendStatus(PAUSE_STATUS);
  } else if (enable.ignoreSend) {
    ctx.setSendStatus(IGNORE_STATUS);
  }
  if (enable.pauseReceive) {
    ctx.setReceiveStatus(PAUSE_STATUS);
  } else if (enable.ignoreReceive) {
    ctx.setReceiveStatus(IGNORE_STATUS);
  }
}

function removePending(reqId) {
  var index = pendingReqList.indexOf(reqId);
  if (index !== -1) {
    pendingReqList.splice(index, 1);
  }
}

function pipeStream(src, target, useSrc) {
  if (!src || !target) {
    return src || target;
  }
  src.pipe(target);
  return useSrc ? src : target;
}

function isHide(req) {
  return config.rulesMode || (req.filter.hide && !req.disable.hide);
}

function emitDataToProxy(req, chunk, fromClient, ignore) {
  if (isHide(req) || req._hasClosed) {
    return;
  }
  proxy.emit('frame', {
    reqId: req.reqId,
    frameId: getFrameId(),
    isClient: fromClient,
    length: chunk.length,
    ignore: ignore,
    bin: chunk
  });
}

function handleConnSend(ctx, reqTrans, sendStatus) {
  var req = ctx.req;
  var res = ctx.res;
  var hasEvent = ctx.hasEvent;
  var writer = res.pipeWriter || res;
  var url = ctx.url;
  ctx.sendToServer = function(data) {
    data = data.data;
    writer.write(data);
    emitDataToProxy(req, data, true);
  };
  sendStatus.emitData = function() {
    if (sendStatus.chunk) {
      emitDataToProxy(req, sendStatus.chunk, true, sendStatus.ignore);
      sendStatus.chunk = null;
    }
  };
  reqTrans._transform = function(chunk, _, cb) {
    hasEvent && proxy.emit('tunnelRequest', url);
    if (sendStatus.pause) {
      sendStatus.chunk = chunk;
      sendStatus.callback = cb;
      return;
    }
    emitDataToProxy(req, chunk, true, sendStatus.ignore);
    if (sendStatus.ignore) {
      chunk = null;
    }
    cb(null, chunk);
  };
}

function handleConnReceive(ctx, resTrans, receiveStatus) {
  var req = ctx.req;
  var hasEvent = ctx.hasEvent;
  var url = ctx.url;
  var writer = req.pipeWriter || req;
  ctx.sendToClient = function(data) {
    data = data.data;
    writer.write(data);
    emitDataToProxy(req, data);
  };
  receiveStatus.emitData = function() {
    if (receiveStatus.chunk) {
      emitDataToProxy(req, receiveStatus.chunk, undefined, receiveStatus.ignore);
      receiveStatus.chunk = null;
    }
  };
          
  resTrans._transform = function(chunk, _, cb) {
    hasEvent && proxy.emit('tunnelRequest', url);
    if (receiveStatus.pause) {
      receiveStatus.chunk = chunk;
      receiveStatus.callback = cb;
      return;
    }
    emitDataToProxy(req, chunk, undefined, receiveStatus.ignore);
    if (receiveStatus.ignore) {
      chunk = null;
    }
    cb(null, chunk);
  };
}

function clearupStatus(conns, reqId, sendStatus, receiveStatus) {
  delete conns[reqId];
  sendStatus.callback = null;
  receiveStatus.callback = null;
  sendStatus.addToReceiver = null;
  receiveStatus.addToReceiver = null;
  clearInterval(sendStatus.timer);
  clearInterval(receiveStatus.timer);
}

function getBinary(data, len) {
  return len > MAX_PAYLOAD ? data.slice(0, MAX_PAYLOAD) : data;
}

function drainData(status, socket, receiver, toServer) {
  if (!status.sender) {
    try {
      status.sender = wsParser.getSender(socket.pipeWriter || socket, toServer);
    } catch (e) {}
  }
  status.sender && status.data.forEach(function(item) {
    status.sender.send(item.data, item);
    receiver.onData(item.data, item);
  });
  status.data = [];
  receiver.ping();
}

function handleFrame(receiver, socket, status, chunk, cb, toServer) {
  if (!receiver.existsCacheData) {
    status.ignoring = status.ignore;
    if (status.pause) {
      status.callback = cb;
      status.chunk = chunk;
      status.addToReceiver = function() {
        receiver.add(chunk);
      };
      drainData(status, socket, receiver, toServer);
      return;
    }
    if (status.ignore) {
      drainData(status, socket, receiver, toServer);
    }
  }
  var toRead = receiver.add(chunk);
  if (status.ignoring) {
    if (!status.ignore && toRead >= 0) {
      status.ignoring = false;
      socket.write(chunk.slice(toRead));
    }
    chunk = null;
  } else if (toRead >= 0 && (status.pause || 
    status.ignore || status.data.length)) {
    if (toRead) {
      var readAll = toRead === chunk.length;
      status.chunk = readAll ? null : chunk.slice(toRead);
      socket.write(readAll ? chunk : chunk.slice(0, toRead));
    } else {
      status.chunk = chunk;
    }
    if (status.pause) {
      status.callback = cb;
      drainData(status, socket, receiver, toServer);
      return;
    }
    if (status.ignore) {
      status.ignoring = true;
      chunk = null;
      drainData(status, socket, receiver, toServer);
    }
  }
  if (chunk && status.timer) {
    clearInterval(status.timer);
    status.timer = null;
  }
  cb(null, chunk);
}

function clearTimer(status) {
  if (!status.ignore && !status.pause) {
    clearInterval(status.timer);
    status.timer = null;
  }
}

function handleWsSend(ctx, reqTrans, sendStatus, isTunnel) {
  var req = ctx.req;
  var res = ctx.res;
  var url = ctx.url;
  var eventName = isTunnel ? 'tunnelRequest' : 'wsRequest';
  var hideWs = isHide(req);
  var reqReceiver;
  if (!hideWs) {
    try {
      reqReceiver = wsParser.getReceiver(res);
    } catch (e) {
      hideWs = true;
    }
  }
  if (hideWs) {
    reqTrans._transform = function(chunk, _, cb) {
      proxy.emit(eventName, url);
      cb(null, chunk);
    };
    return;
  }
  var hasEvent = ctx.hasEvent;
  var reqId = req.reqId;
  handleSocketEnd(req, res, reqReceiver.cleanup.bind(reqReceiver));
  ctx.sendToServer = function(data) {
    if (sendStatus.data.length > MAX_COMPOSE_FRAME_COUNT) {
      return false;
    }
    sendStatus.data.push(data);
    if (sendStatus.ignoring || sendStatus.callback || !reqReceiver.existsCacheData) {
      drainData(sendStatus, res, reqReceiver, true);
    }
  };
  reqReceiver.ping = function() {
    if (isTunnel || sendStatus.timer || req.disable.pong) {
      return;
    }
    res.write(PONG);
    sendStatus.timer = setInterval(function() {
      res.write(PONG);
      clearTimer(sendStatus);
    }, INTERVAL);
  };
  reqReceiver.onclose = function(code) {
    ctx.req._errorCode = code;
  };
  reqReceiver.onData = function(data, opts) {
    if (req._hasClosed) {
      return;
    }
    var opcode = opts.opcode;
    if (!opcode) {
      opcode = opts.binary ? 2 : 1;
    }
    proxy.emit('frame', {
      reqId: reqId,
      frameId: getFrameId(),
      isClient: true,
      mask: isTunnel ? undefined : opts.mask,
      ignore: opts.data ? undefined : sendStatus.ignoring,
      bin: getBinary(data, opts.length),
      compressed: isTunnel ? undefined : opts.compressed,
      length: opts.length,
      opcode: isTunnel ? undefined : opcode
    });
  };
  reqReceiver.onerror = function(err) {
    if (req._hasClosed) {
      return;
    }
    req._hasClosed = true;
    proxy.emit('frame', {
      reqId: reqId,
      frameId: getFrameId(),
      isClient: true,
      err: err.message,
      bin: ''
    });
  };
  reqTrans._transform = function(chunk, _, cb) {
    hasEvent && proxy.emit(eventName, url);
    handleFrame(reqReceiver, res, sendStatus, chunk, cb, true);
  };
}

function handleWsReceive(ctx, resTrans, receiveStatus, isTunnel) {
  var req = ctx.req;
  var res = ctx.res;
  var url = ctx.url;
  var eventName = isTunnel ? 'tunnelRequest' : 'wsRequest';
  var hideWs = isHide(req);
  var resReceiver;
  if (!hideWs) {
    try {
      resReceiver = wsParser.getReceiver(res, true);
    } catch (e) {
      hideWs = true;
    }
  }
  if (hideWs) {
    resTrans._transform = function(chunk, _, cb) {
      proxy.emit(eventName, url);
      cb(null, chunk);
    };
    return;
  }
  var hasEvent = ctx.hasEvent;
  var reqId = req.reqId;
  handleSocketEnd(req, res, resReceiver.cleanup.bind(resReceiver));
  ctx.sendToClient = function (data) {
    if (receiveStatus.data.length > MAX_COMPOSE_FRAME_COUNT) {
      return false;
    }
    receiveStatus.data.push(data);
    if (receiveStatus.ignoring || receiveStatus.callback || !resReceiver.existsCacheData) {
      drainData(receiveStatus, req, resReceiver);
    }
  };
  resReceiver.ping = function() {
    if (isTunnel || receiveStatus.timer || req.disable.ping) {
      return;
    }
    req.write(PING);
    receiveStatus.timer = setInterval(function() {
      req.write(PING);
      clearTimer(receiveStatus);
    }, INTERVAL);
  };
  resReceiver.onclose = function(code) {
    ctx.res._errorCode = code;
  };
  resReceiver.onData = function(data, opts) {
    if (req._hasClosed) {
      return;
    }
    var opcode = opts.opcode;
    if (!opcode) {
      opcode = opts.binary ? 2 : 1;
    }
    proxy.emit('frame', {
      reqId: reqId,
      frameId: getFrameId(),
      bin: getBinary(data, opts.length),
      mask: isTunnel ? undefined : opts.mask,
      ignore: opts.data ? undefined : receiveStatus.ignoring,
      compressed: isTunnel ? undefined : opts.compressed,
      length: opts.length,
      opcode: isTunnel ? undefined : opcode
    });
  };
  resReceiver.onerror = function(err) {
    if (req._hasClosed) {
      return;
    }
    req._hasClosed = true;
    proxy.emit('frame', {
      reqId: reqId,
      frameId: getFrameId(),
      err: err.message,
      bin: ''
    });
  };
  resTrans._transform = function(chunk, _, cb) {
    hasEvent && proxy.emit(eventName, url);
    handleFrame(resReceiver, req, receiveStatus, chunk, cb);
  };
}

function getContext(req, res, hasEvent, sendStatus, receiveStatus) {
  var reqId = req.reqId;
  var ctx = conns[reqId] = {
    customParser: req.customParser,
    req: req,
    res: res,
    hasEvent: hasEvent,
    url: req.fullUrl,
    charset: util.getCharset(res.headers['content-type']) || '',
    clearup: function() {
      clearupStatus(conns, reqId, sendStatus, receiveStatus);
    },
    setSendStatus: function(status) {
      setConnStatus(ctx, status, sendStatus, 'sendStatus');
    },
    setReceiveStatus: function(status) {
      setConnStatus(ctx, status, receiveStatus);
    }
  };
  initStatus(ctx, req.enable);
  return ctx;
}

exports.handleUpgrade = function(req, res) {
  var hasEvent = util.listenerCount(proxy, 'wsRequest');
  var hide = isHide(req);
  if (req.isPluginReq || (!hasEvent && hide)) {
    handleClose(req, res, true);
    return req.pipe(res).pipe(req);
  }
  var url = req.fullUrl;
  var customParser = req.customParser;
  var reqId = req.reqId;
  var sendStatus = { data: [] };
  var receiveStatus = { data: [] };
  var emitError = function(err) {
    req.emit('error', err);
  };
  var reqTrans = new PassThrough();
  var resTrans = new PassThrough();
 
  reqTrans.on('error', emitError);
  resTrans.on('error', emitError);
  res.headers = res.headers || {};
  req.wsExts = res.headers['sec-websocket-extensions'] || '';
  handleClose(req, res);
  pluginMgr.getWsPipe(req, res, function(reqRead, reqWrite, resRead, resWrite) {
    if (req._hasClosed) {
      return;
    }
    var hasPipeStream = reqRead || reqWrite || resRead || resWrite;
    if (hasPipeStream) {
      req.customParser = customParser = false;
    }
    var ctx = conns[reqId] = getContext(req, res, hasEvent, sendStatus, receiveStatus);
    if (customParser) {
      removePending(reqId);
      var handleInspect = function(chunk, _, cb) {
        hasEvent && proxy.emit('wsRequest', url);
        cb(null, chunk);
      };
      reqTrans._transform = handleInspect;
      resTrans._transform = handleInspect;
    } else {
      if (reqWrite) {
        reqWrite.headers = res.headers;
        res.pipeWriter = reqWrite;
      }
      if (resWrite) {
        resWrite.headers = res.headers;
        req.pipeWriter = resWrite;
      }
      handleWsSend(ctx, reqTrans, sendStatus);
      handleWsReceive(ctx, resTrans, receiveStatus);
    }
    pipeStream(req, reqRead).pipe(reqTrans).pipe(pipeStream(reqWrite, res, true));
    pipeStream(res, resRead).pipe(resTrans).pipe(pipeStream(resWrite, req, true));
  });
};

exports.handleConnect = function(req, res) {
  var hasEvent = util.listenerCount(proxy, 'tunnelRequest');
  var isConn = req.inspectFrames;
  if (req.isPluginReq || (!isConn && !hasEvent)) {
    handleClose(req, res, true);
    return req.pipe(res).pipe(req);
  }
  var url = req.fullUrl;
  var reqId = req.reqId;
  var customParser = req.customParser;
  var sendStatus = { data: [] };
  var receiveStatus = { data: [] };
  var ctx = '';
  var reqTrans = new PassThrough();
  var resTrans = new PassThrough();
  var emitError = function(err) {
    req.emit('error', err);
  };
  reqTrans.on('error', emitError);
  resTrans.on('error', emitError);
  res.headers = res.headers || req.headers;
  handleClose(req, res, !isConn);
  pluginMgr.getTunnelPipe(req, res, function(reqRead, reqWrite, resRead, resWrite) {
    if (req._hasClosed) {
      return;
    }
    var hide = isHide(req);
    if (isConn && !hide) {
      ctx = getContext(req, res, hasEvent, sendStatus, receiveStatus);
    }
        
    if (customParser || hide || !isConn) {
      removePending(reqId);
      var handleInspect = function(chunk, _, cb) {
        hasEvent && proxy.emit('tunnelRequest', url);
        cb(null, chunk);
      };
      reqTrans._transform = handleInspect;
      resTrans._transform = handleInspect;
    } else {
      req.wsExts = res.wsExts = '';
      if (reqWrite) {
        reqWrite.wsExts = '';
        res.pipeWriter = reqWrite;
      }
      if (resWrite) {
        resWrite.wsExts = '';
        req.pipeWriter = resWrite;
      }
      if (reqRead && reqWrite) {
        handleWsSend(ctx, reqTrans, sendStatus, true);
      } else {
        handleConnSend(ctx, reqTrans, sendStatus);
      }
      if (resRead && resWrite) {
        handleWsReceive(ctx, resTrans, receiveStatus, true);
      } else {
        handleConnReceive(ctx, resTrans, receiveStatus);
      }
    }
            
    pipeStream(req, reqRead).pipe(reqTrans).pipe(pipeStream(reqWrite, res, true));
    pipeStream(res, resRead).pipe(resTrans).pipe(pipeStream(resWrite, req, true));
  });
};

exports.abort = function(reqId) {
  var ctx = conns[reqId];
  if (!ctx) {
    return;
  }
  delete conns[reqId];
  ctx.req.destroy();
  ctx.res.destroy();
};

exports.getStatus = function(reqId) {
  var ctx = reqId && conns[reqId];
  if (!ctx) {
    return;
  }
  return {
    sendStatus: ctx.sendStatus,
    receiveStatus: ctx.receiveStatus
  };
};

exports.setPending = function(req) {
  var reqId = req.customParser && req.reqId;
  if (reqId && pendingReqList.indexOf(reqId) === -1) {
    pendingReqList.push(reqId);
    if (pendingReqList.length > 500) {
      pendingReqList = pendingReqList.slice(-300);
    } 
  }
};

exports.exists = function(reqId) {
  return reqId && conns[reqId];
};

exports.getData = function(reqId) {
  var ctx = reqId && conns[reqId];
  if (ctx) {
    var result = {
      sendStatus: ctx.sendStatus,
      receiveStatus: ctx.receiveStatus,
      toServer: ctx.toServerData,
      toClient: ctx.toClientData
    };
    delete ctx.toServerData;
    delete ctx.toClientData;
    return result;
  }
  return pendingReqList.indexOf(reqId) === -1 ? undefined : 1;
};

exports.changeStatus = function(data) {
  var ctx = conns[data.reqId];
  if (!ctx) {
    return;
  }
  if (data.sendStatus >= 0) {
    ctx.setSendStatus(data.sendStatus);
  } else {
    ctx.setReceiveStatus(data.receiveStatus);
  }
  return true;
};

function getBuffer(data, charset) {
  if (data.base64) {
    try {
      return Buffer.from(data.base64, 'base64');
    } catch(e) {}
  } else if (data.text) {
    return util.toBuffer(data.text, charset);
  }
}

exports.sendData = function(data) {
  var ctx = conns[data.reqId];
  if (!ctx) {
    return;
  }
  var buf = getBuffer(data, ctx.charset);
  if (!buf) {
    return;
  }
  var isServer = data.target === 'server';
  var binary = data.type === 'bin';
  data = { binary: binary };
  if (ctx.customParser) {
    var name = isServer ? 'toServerData' : 'toClientData';
    var dataList = ctx[name] = ctx[name] || [];
    if (dataList.length > MAX_COMPOSE_FRAME_COUNT) {
      return false;
    }
    data.base64 = buf.toString('base64');
    dataList.push(data);
    return;
  }
  data.binary = binary;
  data.length = buf.length;
  data.data = buf;
  if (isServer) {
    data.mask = true;
    return ctx.sendToServer(data);
  } else {
    return ctx.sendToClient(data);
  }
};