1#!/usr/bin/env node
  2
  3"use strict";
  4
  5var http = require('http');
  6var httpProxy = require('./network_simulation/lib/proxy');
  7var checkServer = require('./network_simulation/lib/check_server');
  8var nopt = require('nopt');
  9var url = require('url');
 10var pending = 0; // gotta love global vars
 11
 12var RequestData = function (request, response, buffer) {
 13  this.request = request;
 14  this.response = response;
 15  this.buffer = buffer;
 16};
 17
 18function ipFromRequest(reqData) {
 19  return reqData.request.headers['x-forwarded-for'];
 20}
 21
 22function rejectRequest(reqData) {
 23  reqData.response.writeHead(400);
 24  reqData.response.end();
 25}
 26
 27var Queue = function (proxies, parameters, users, last_sent) {
 28  this.proxies = proxies;
 29  this.parameters = parameters;
 30  this.users = users;
 31  this.last_sent = last_sent;
 32};
 33
 34Queue.prototype.takeRequest = function (reqData) {
 35  // Reject traffic as necessary:
 36  var req_ip = ipFromRequest(reqData);
 37  var now = new Date().getTime();
 38  var bursty_threshold = 120;
 39  var idle_threshold = 4;
 40
 41  // maximize utility always
 42  if (pending < idle_threshold) {
 43    // proxy it through
 44    var rand = Math.floor(Math.random()*this.proxies.length);
 45    this.proxies[rand].proxyRequest(reqData.request, reqData.response, reqData.buffer);
 46    pending += 1;
 47
 48  } else {
 49    if (req_ip in this.users) {
 50     
 51      // check for blacklist
 52      if (this.users[req_ip]["blacklist"]) {
 53        rejectRequest(reqData);
 54        return;
 55      }
 56
 57      // check for bad behavior 
 58      else if (now - this.users[req_ip]["last"] < bursty_threshold) {
 59        //console.log("BLACKLISTED! %s, last:%d(ms), total:%d(requests)", req_ip.substring(0,8),  now - this.users[req_ip]["last"], this.users[req_ip]["total"]+1);
 60        this.users[req_ip]["blacklist"] = true;
 61        this.users[req_ip]["total"] = 0;
 62        rejectRequest(reqData);
 63        return;
 64      }
 65
 66    } else {
 67      // new user to track
 68      this.users[req_ip] = new Object();
 69      this.users[req_ip]["total"] =  0;
 70      this.users[req_ip]["blacklist"] = false;
 71    } 
 72
 73    // update user stats
 74    this.users[req_ip]["total"] = this.users[req_ip]["total"] + 1;
 75    this.users[req_ip]["last"] = new Date().getTime();
 76    
 77    // proxy it through
 78    var rand = Math.floor(Math.random()*this.proxies.length);
 79    this.proxies[rand].proxyRequest(reqData.request, reqData.response, reqData.buffer);
 80    pending += 1;
 81  }
 82};
 83
 84Queue.prototype.requestFinished = function () {
 85  // remove one from the pending counter
 86  if (pending > 0){
 87    pending -= 1;
 88  }
 89  return;
 90};
 91
 92function checkBackends(targets, path, response) {
 93  var toCheck = targets.map(function (target) {
 94    var output = {};
 95    output['host'] = target['host'];
 96    output['port'] = target['port'];
 97    output['path'] = path;
 98    return output;
 99  });
100  var success = function () {
101    response.writeHead(200, {"Content-Type": "application/json"});
102    response.end()
103  };
104  var error = function () {
105    response.writeHead(500, {"Content-Type": "application/json"});
106    response.end()
107  };
108  checkServer.checkServers(toCheck, success, error);
109}
110
111function main() {
112  var opts = {
113    "out-ports": String,
114    "in-port": String,
115  };
116  var parsed = nopt(opts),
117      inPort = parsed['in-port'] || '3000',
118      outPorts = parsed['out-ports'] ? parsed['out-ports'].split(",") : ['3001'],
119      targets = [],
120      target,
121      proxies = [],
122      proxy,
123      i;
124
125  for (i = 0; i < outPorts.length; i++) {
126    target = {'host': 'localhost', 'port': outPorts[i]};
127    targets.push(target);
128    proxy = new httpProxy.HttpProxy({'target': target});
129    proxy.identifier = i;
130    proxies.push(proxy);
131  }
132 
133  var users = new Object();
134  var queue = new Queue(proxies, {}, users, new Date().getTime());
135  for (i = 0; i < proxies.length; i++) {
136    proxy = proxies[i];
137    proxy.on("end", queue.requestFinished);
138  }
139
140  var server = http.createServer(function (req, res) {
141    if (req.method === "HEAD") {
142      // HEAD requests are used to monitor the status of the simulation
143      // proxies[0].proxyRequest(reqData.request, reqData.response, reqData.buffer);
144      checkBackends(targets, url.parse(req.url)['pathname'], res);
145    } else {
146      var buffer = httpProxy.buffer(req);
147      var reqData = new RequestData(req, res, buffer);
148      queue.takeRequest(reqData);
149    }
150  });
151
152  server.on('close', function () {
153    for (i = 0; i < proxies.length; i++) {
154      proxies[i].close();
155    }
156  });
157  console.log("The shield is up and listening.");
158  server.listen(inPort);
159}
160
161main();