refactor(*): rewritten to use koa-compose and koa-ip-filter

index.js

@@ -7,6 +7,69 @@
 
 'use strict'
 
-module.exports = function () {
+const convert = require('koa-convert')
+const compose = require('koa-compose')
+const filter = require('koa-ip-filter')
+const extend = require('extend-shallow')
 
+module.exports = function koaBetterRatelimit (opts) {
+  opts.filter = opts.filter || '**'
+
+  const ipFilter = convert(filter(opts))
+  let store = opts.store && typeof opts.store === 'object'
+    ? opts.store
+    : {}
+
+  return compose([ipFilter, function (ctx, done) {
+    const headers = extend({
+      remaining: 'X-RateLimit-Remaining',
+      reset: 'X-RateLimit-Reset',
+      limit: 'X-RateLimit-Limit'
+    }, opts.headers)
+
+    opts = extend({
+      duration: 1000 * 60 * 60 * 24,
+      limited: 'Too Many Requests',
+      max: 500
+    }, opts)
+
+    const id = ctx.identifier
+    const data = store[id]
+    const now = Date.now()
+    const reset = now + opts.duration
+    const needReset = data && data.remaining === false && data.reset < now
+
+    if (!store.hasOwnProperty(id) || needReset) {
+      store[id] = {
+        id: id,
+        reset: reset,
+        remaining: opts.max
+      }
+    }
+
+    const entry = store[id]
+    entry.remaining = entry.remaining > 0
+        ? entry.remaining - 1
+        : false
+
+    ctx.set(headers.remaining, entry.remaining || 0)
+    ctx.set(headers.reset, entry.reset)
+    ctx.set(headers.limit, opts.max)
+
+    if (typeof entry.remaining === 'number') {
+      return done()
+    }
+
+    const after = entry.reset - (now / 1000) | 0
+    ctx.set('Retry-After', after)
+
+    ctx.status = 429
+    ctx.body = typeof opts.limited === 'function'
+        ? opts.limited(ctx, entry)
+        : opts.limited
+
+    if (opts.throw) {
+      ctx.throw(ctx.status, ctx.body)
+    }
+  }])
 }

package.json

@@ -1,7 +1,7 @@
 {
   "name": "koa-better-ratelimit",
   "version": "2.1.2",
-  "description": "Request limiter by IP or custom ID for your APIs built on [koa][] - works both for v1 and v2.",
+  "description": "Smart and easy request rate limiter for your APIs built on [koa][], using [koa-ip-filter][]. Support custom stores, custom ID, custom error messages and custom headers.",
   "repository": "tunnckoCore/koa-better-ratelimit",
   "author": "Charlike Mike Reagent <@tunnckoCore> (http://www.tunnckocore.tk)",
   "precommit.silent": true,
@@ -20,13 +20,17 @@
     "precommit": "git add --all",
     "commit": "git-cz"
   },
-  "dependencies": {},
+  "dependencies": {
+    "extend-shallow": "^2.0.1",
+    "koa-compose": "^3.1.0",
+    "koa-convert": "^1.2.0",
+    "koa-ip-filter": "^3.0.0"
+  },
   "devDependencies": {
     "commitizen": "^2.8.6",
     "coveralls": "^2.11.12",
     "cz-conventional-changelog": "^1.2.0",
     "koa": "^1.2.4",
-    "koa-better-body": "^3.0.2",
     "koa2": "^2.0.0-alpha.7",
     "mukla": "^0.4.1",
     "nyc": "^8.1.0",

test.js

@@ -8,3 +8,31 @@
 /* jshint asi:true */
 
 'use strict'
+
+const ratelimit = require('./index')
+const Koa = require('koa2')
+const app = new Koa()
+
+app
+  .use(function * (next) {
+    try {
+      yield * next
+    } catch (e) {
+      console.log('err', e)
+    }
+  })
+  .use(ratelimit({
+    // it can be double star - globstar
+    // or *.*.*.*
+    id: function () {
+      return '0.0.0.0'
+    },
+    duration: 1000 * 20,
+    // filter: '*.*.*.*',
+    max: 2,
+    throw: true
+  }))
+
+// .listen(2222, function () {
+//   console.log('Foo bar baz!')
+// })