edge-utils
Version:
Platform-agnostic utilities for edge computing and serverless environments
3 lines (2 loc) • 4.57 kB
JavaScript
const t=require("crypto");class e{constructor(t={}){this.refillRate=t.refillRate||10,this.capacity=t.capacity||100,this.tokens=this.capacity,this.lastRefill=Date.now(),this.storage=t.storage||new Map}_refill(){const t=Date.now(),e=(t-this.lastRefill)/1e3,s=Math.floor(e*this.refillRate);s>0&&(this.tokens=Math.min(this.capacity,this.tokens+s),this.lastRefill=t)}async check(t,e=1){const s=await this._getState(t);return this._refillState(s),s.tokens>=e&&(s.tokens-=e,await this._setState(t,s),!0)}async remaining(t){const e=await this._getState(t);return this._refillState(e),Math.max(0,e.tokens)}async resetTime(t){const e=await this._getState(t),s=(this.capacity-e.tokens)/this.refillRate;return Math.ceil(e.lastRefill/1e3+s)}async _getState(t){try{const e=await(this.storage.get?.(`ratelimit:${t}`))||this.storage.get(`ratelimit:${t}`);if(e)return JSON.parse(e)}catch(t){}return{tokens:this.capacity,lastRefill:Date.now()}}async _setState(t,e){try{const s=JSON.stringify(e);await(this.storage.put?.(`ratelimit:${t}`,s))||this.storage.set(`ratelimit:${t}`,s)}catch(t){}}_refillState(t){const e=Date.now(),s=(e-t.lastRefill)/1e3,i=Math.floor(s*this.refillRate);i>0&&(t.tokens=Math.min(this.capacity,t.tokens+i),t.lastRefill=e)}}class s{constructor(t={}){this.windowSize=t.windowSize||60,this.maxRequests=t.maxRequests||100,this.storage=t.storage||new Map}async check(t){const e=Date.now(),s=e-1e3*this.windowSize;try{await this._cleanup(t,s);return!(await this._count(t,s)>=this.maxRequests)&&(await this._add(t,e),!0)}catch(t){return!0}}async remaining(t){const e=Date.now()-1e3*this.windowSize;try{const s=await this._count(t,e);return Math.max(0,this.maxRequests-s)}catch(t){return this.maxRequests}}async resetTime(t){try{const e=await this._getTimestamps(t);if(0===e.length)return Math.ceil(Date.now()/1e3);const s=Math.min(...e);return Math.ceil((s+1e3*this.windowSize)/1e3)}catch(t){return Math.ceil(Date.now()/1e3+this.windowSize)}}async _cleanup(t,e){const s=(await this._getTimestamps(t)).filter(t=>t>=e);await this._setTimestamps(t,s)}async _count(t,e){return(await this._getTimestamps(t)).filter(t=>t>=e).length}async _add(t,e){const s=await this._getTimestamps(t);s.push(e),await this._setTimestamps(t,s)}async _getTimestamps(t){try{const e=await(this.storage.get?.(`ratelimit:${t}`))||this.storage.get(`ratelimit:${t}`);return e?JSON.parse(e):[]}catch(t){return[]}}async _setTimestamps(t,e){try{const s=JSON.stringify(e);await(this.storage.put?.(`ratelimit:${t}`,s))||this.storage.set(`ratelimit:${t}`,s)}catch(t){}}}module.exports={TokenBucketLimiter:e,SlidingWindowLimiter:s,RateLimitManager:class{constructor(t={}){this.strategies=new Map,this.storage=t.storage,this.defaultStrategy=t.defaultStrategy||"token-bucket",this.exemptions=new Set(t.exemptions||[]),this.whitelist=new Set(t.whitelist||[])}addStrategy(t,i){const{type:a,...r}=i;switch(r.storage=this.storage,a){case"token-bucket":this.strategies.set(t,new e(r));break;case"sliding-window":this.strategies.set(t,new s(r));break;default:throw new Error(`Unknown rate limit type: ${a}`)}}async check(t,e={}){const s="ip"===e.by?this._getClientIP(t):"api-key"===e.by?t.headers?.["x-api-key"]||"anonymous":"user"===e.by?t.user?.id||"anonymous":"global";if(this.exemptions.has(s)||this.whitelist.has(s))return{allowed:!0,remaining:-1,resetTime:-1,headers:{}};const i=this._generateKey(t,e),a=e.strategy||this.defaultStrategy,r=this.strategies.get(a);if(!r)throw new Error(`Rate limit strategy not found: ${a}`);const n=await r.check(i),h=await r.remaining(i),o=await r.resetTime(i);return{allowed:n,remaining:h,resetTime:o,headers:{"X-RateLimit-Limit":e.limit||r.capacity||r.maxRequests,"X-RateLimit-Remaining":h,"X-RateLimit-Reset":o}}}_generateKey(e,s){const i=[];return"ip"===s.by?i.push(this._getClientIP(e)):"api-key"===s.by?i.push(e.headers?.["x-api-key"]||"anonymous"):"user"===s.by?i.push(e.user?.id||"anonymous"):"composite"===s.by?(i.push(this._getClientIP(e)),i.push(e.method),i.push(e.url?.pathname||"/")):i.push("global"),t.createHash("sha256").update(i.join(":")).digest("hex")}_getClientIP(t){return t.headers?.["cf-connecting-ip"]||t.headers?.["x-forwarded-for"]?.split(",")[0]||t.headers?.["x-real-ip"]||"127.0.0.1"}middleware(t={}){return async(e,s)=>{const i=await this.check(e,t);if(!i.allowed){const t=Math.max(1,i.resetTime-Math.floor(Date.now()/1e3));return new Response("Rate limit exceeded",{status:429,headers:{...i.headers,"Retry-After":t.toString(),"Content-Type":"text/plain"}})}return s.rateLimitHeaders=i.headers,null}}}};
//# sourceMappingURL=rate-limiting.esm.js.map