import { SchemaDirectiveVisitor } from '@graphql-tools/utils' ; import { defaultFieldResolver , GraphQLField } from 'graphql' ; import Redis from 'ioredis' ; import { CACHE_PREFIX } from 'common/enums' ; import { ActionLimitExceededError } from 'common/errors' ; import { CacheService } from 'connectors' ; const checkOperationLimit = async ( { user , operation , limit , period , } ) => { const cacheService = new CacheService ( CACHE_PREFIX . OPERATION_LOG ) ; const cacheKey = cacheService . genKey ( { id : user , field : operation , } ) ; const redisClient = cacheService . redis . client as Redis . Redis ; const operationLog = await redisClient . lrange ( cacheKey , 0 , - 1 ) ; const current = Math . floor ( Date . now ( ) / 1000 ) ; if ( ! operationLog ) { redisClient . lpush ( cacheKey , current ) . then ( ( ) => { redisClient . expire ( cacheKey , period ) ; } ) ; return true ; } const valid = operationLog . map ( ( timestamp ) => parseInt ( timestamp , 10 ) >= current - period ) ; const times = valid . reduce < any > ( ( a , b ) => ( a ? 1 : 0 ) + ( b ? 1 : 0 ) , 0 ) ; if ( times >= limit ) { return false ; } redisClient . lpush ( cacheKey , current ) ; redisClient . ltrim ( cacheKey , 0 , times ) ; redisClient . expire ( cacheKey , period ) ; return true ; } ; export class RateLimitDirective extends SchemaDirectiveVisitor { public visitFieldDefinition ( field ) { const { resolve = defaultFieldResolver , name } = field ; const { limit , period } = this . args ; field . resolve = async function ( ... args ) { const [ , , { viewer } ] = args ; const pass = await checkOperationLimit ( { user : viewer . id || viewer . ip , operation : name , limit , period , } ) ; if ( ! pass ) { throw new ActionLimitExceededError ( `rate exceeded for operation ${ name }` ) ; } return resolve . apply ( this , args ) ; } ; } }