« index

Coverage for /Users/yunong/workspace/node-restify/lib/plugins/throttle.js : 96%

261 lines | 252 run | 9 missing | 1 partial | 25 blocks | 18 blocks run | 7 blocks missing

1

2

3

4

5

6

7

8

9

10

11

12

13

14

15

16

17

18

19

20

21

22

23

24

25

26

27

28

29

30

31

32

33

34

35

36

37

38

39

40

41

42

43

44

45

46

47

48

49

50

51

52

53

54

55

56

57

58

59

60

61

62

63

64

65

66

67

68

69

70

71

72

73

74

75

76

77

78

79

80

81

82

83

84

85

86

87

88

89

90

91

92

93

94

95

96

97

98

99

100

101

102

103

104

105

106

107

108

109

110

111

112

113

114

115

116

117

118

119

120

121

122

123

124

125

126

127

128

129

130

131

132

133

134

135

136

137

138

139

140

141

142

143

144

145

146

147

148

149

150

151

152

153

154

155

156

157

158

159

160

161

162

163

164

165

166

167

168

169

170

171

172

173

174

175

176

177

178

179

180

181

182

183

184

185

186

187

188

189

190

191

192

193

194

195

196

197

198

199

200

201

202

203

204

205

206

207

208

209

210

211

212

213

214

215

216

217

218

219

220

221

222

223

224

225

226

227

228

229

230

231

232

233

234

235

236

237

238

239

240

241

242

243

244

245

246

247

248

249

250

251

252

253

254

255

256

257

258

259

260

261

  // Copyright 2012 Mark Cavage <mcavage@gmail.com> All rights reserved.
  
  var sprintf = require('util').format;
  
  var assert = require('assert-plus');
  var LRU = require('lru-cache');
  
  var errors = require('../errors');
  
  
  ///--- Globals
  
  var TooManyRequestsError = errors.TooManyRequestsError;
  
  var MESSAGE = 'You have exceeded your request rate of %s r/s.';
  
  
  ///--- Helpers
  
  function xor() {
      var x = false;
      for (var i = 0; i < arguments.length; i++) {
          if (arguments[i] && !x)
              x = true;
          else if (arguments[i] && x)
              return (false);
      }
      return (x);
  }
  
  
  ///--- Internal Class (TokenBucket)
  
  /**
   * An implementation of the Token Bucket algorithm.
   *
   * Basically, in network throttling, there are two "mainstream"
   * algorithms for throttling requests, Token Bucket and Leaky Bucket.
   * For restify, I went with Token Bucket.  For a good description of the
   * algorithm, see: http://en.wikipedia.org/wiki/Token_bucket
   *
   * In the options object, you pass in the total tokens and the fill rate.
   * Practically speaking, this means "allow `fill rate` requests/second,
   * with bursts up to `total tokens`".  Note that the bucket is initialized
   * to full.
   *
   * Also, in googling, I came across a concise python implementation, so this
   * is just a port of that. Thanks http://code.activestate.com/recipes/511490 !
   *
   * @param {Object} options contains the parameters:
   *                   - {Number} capacity the maximum burst.
   *                   - {Number} fillRate the rate to refill tokens.
   */
  function TokenBucket(options) {
      assert.object(options, 'options');
      assert.number(options.capacity, 'options.capacity');
      assert.number(options.fillRate, 'options.fillRate');
  
      this.tokens = this.capacity = options.capacity;
      this.fillRate = options.fillRate;
      this.time = Date.now();
  }
  
  
  /**
   * Consume N tokens from the bucket.
   *
   * If there is not capacity, the tokens are not pulled from the bucket.
   *
   * @param {Number} tokens the number of tokens to pull out.
   * @return {Boolean} true if capacity, false otherwise.
   */
  TokenBucket.prototype.consume = function consume(tokens) {
      if (tokens <= this._fill()) {
          this.tokens -= tokens;
          return (true);
      }
  
      return (false);
  };
  
  
  /**
   * Fills the bucket with more tokens.
   *
   * Rather than do some whacky setTimeout() deal, we just approximate refilling
   * the bucket by tracking elapsed time from the last time we touched the bucket.
   *
   * Simply, we set the bucket size to min(totalTokens,
   *                                       current + (fillRate * elapsed time)).
   *
   * @return {Number} the current number of tokens in the bucket.
   */
  TokenBucket.prototype._fill = function _fill() {
      var now = Date.now();
      if (now < this.time) // reset account for clock drift (like DST)
          this.time = now - 1000;
  
      if (this.tokens < this.capacity) {
          var delta = this.fillRate * ((now - this.time) / 1000);
          this.tokens = Math.min(this.capacity, this.tokens + delta);
      }
      this.time = now;
  
      return (this.tokens);
  };
  
  
  ///--- Internal Class (TokenTable)
  // Just a wrapper over LRU that supports put/get to store token -> bucket
  // mappings
  
  function TokenTable(options) {
      assert.object(options, 'options');
  
      this.table = new LRU(options.size || 10000);
  }
  
  
  TokenTable.prototype.put = function put(key, value) {
      this.table.set(key, value);
  };
  
  
  TokenTable.prototype.get = function get(key) {
      return (this.table.get(key));
  };
  
  
  ///--- Exported API
  
  /**
   * Creates an API rate limiter that can be plugged into the standard
   * restify request handling pipeline.
   *
   * This throttle gives you three options on which to throttle:
   * username, IP address and 'X-Forwarded-For'. IP/XFF is a /32 match,
   * so keep that in mind if using it.  Username takes the user specified
   * on req.username (which gets automagically set for supported Authorization
   * types; otherwise set it yourself with a filter that runs before this).
   *
   * In both cases, you can set a `burst` and a `rate` (in requests/seconds),
   * as an integer/float.  Those really translate to the `TokenBucket`
   * algorithm, so read up on that (or see the comments above...).
   *
   * In either case, the top level options burst/rate set a blanket throttling
   * rate, and then you can pass in an `overrides` object with rates for
   * specific users/IPs.  You should use overrides sparingly, as we make a new
   * TokenBucket to track each.
   *
   * On the `options` object ip and username are treated as an XOR.
   *
   * An example options object with overrides:
   *
   *  {
   *    burst: 10,  // Max 10 concurrent requests (if tokens)
   *    rate: 0.5,  // Steady state: 1 request / 2 seconds
   *    ip: true,   // throttle per IP
   *    overrides: {
   *      '192.168.1.1': {
   *        burst: 0,
   *        rate: 0    // unlimited
   *    }
   *  }
   *
   *
   * @param {Object} options required options with:
   *                   - {Number} burst (required).
   *                   - {Number} rate (required).
   *                   - {Boolean} ip (optional).
   *                   - {Boolean} username (optional).
   *                   - {Boolean} xff (optional).
   *                   - {Object} overrides (optional).
   *                   - {Object} tokensTable: a storage engine this plugin will
   *                              use to store throttling keys -> bucket mappings.
   *                              If you don't specify this, the default is to
   *                              use an in-memory O(1) LRU, with 10k distinct
   *                              keys.  Any implementation just needs to support
   *                              put/get.
   *                   - {Number} maxKeys: If using the default implementation,
   *                              you can specify how large you want the table to
   *                              be.  Default is 10000.
   * @return {Function} of type f(req, res, next) to be plugged into a route.
   * @throws {TypeError} on bad input.
   */
  function throttle(options) {
      assert.object(options, 'options');
      assert.number(options.burst, 'options.burst');
      assert.number(options.rate, 'options.rate');
      if (!xor(options.ip, options.xff, options.username))
          throw new Error('(ip ^ username ^ xff)');
  
      var table = options.tokensTable ||
          new TokenTable({size: options.maxKeys});
  
      function rateLimit(req, res, next) {
          var attr;
          var burst = options.burst;
          var rate = options.rate;
  
          if (options.ip) {
              attr = req.connection.remoteAddress;
          } else if (options.xff) {
              attr = req.headers['x-forwarded-for'];
          } else if (options.username) {
              attr = req.username;
          } else {
              req.log.warn({config: options},
                  'Invalid throttle configuration');
              return (next());
          }
  
          // Before bothering with overrides, see if this request
          // even matches
          if (!attr)
              return (next());
  
          // Check the overrides
          if (options.overrides &&
              options.overrides[attr] &&
              options.overrides[attr].burst !== undefined &&
              options.overrides[attr].rate !== undefined) {
  
              burst = options.overrides[attr].burst;
              rate = options.overrides[attr].rate;
          }
  
          if (!rate || !burst)
              return (next());
  
          var bucket = table.get(attr);
          if (!bucket) {
              bucket = new TokenBucket({
                  capacity: burst,
                  fillRate: rate
              });
              table.put(attr, bucket);
          }
  
          req.log.trace('Throttle(%s): num_tokens= %d',
              attr, bucket.tokens);
  
          if (!bucket.consume(1)) {
              req.log.info({
                  address: req.connection.remoteAddress || '?',
                  method: req.method,
                  url: req.url,
                  user: req.username || '?'
              }, 'Throttling');
  
              var msg = sprintf(MESSAGE, rate);
              return (next(new TooManyRequestsError(msg)));
          }
  
          return (next());
      }
  
      return (rateLimit);
  }
  
  module.exports = throttle;
« index | cover.io