feat: implement caching for user admission check

Fixes #447

Author: Mohit-Davar

src/@types/adapters.ts

@@ -22,6 +22,7 @@ export interface ICacheAdapter {
   getKey(key: string): Promise<string>
   hasKey(key: string): Promise<boolean>
   setKey(key: string, value: string): Promise<boolean>
+  setKeyWithExpiry(key: string, value: string, ttl: number): Promise<boolean>
   addToSortedSet(key: string, set: Record<string, string> | Record<string, string>[]): Promise<number>
   removeRangeByScoreFromSortedSet(key: string, min: number, max: number): Promise<number>
   getRangeFromSortedSet(key: string, start: number, stop: number): Promise<string[]>

src/adapters/redis-adapter.ts

@@ -60,10 +60,16 @@ export class RedisAdapter implements ICacheAdapter {
 
   public async setKey(key: string, value: string): Promise<boolean> {
     await this.connection
-    debug('get %s key', key)
+    debug('set %s key', key)
     return 'OK' === await this.client.set(key, value)
   }
 
+  public async setKeyWithExpiry(key: string, value: string, ttl: number): Promise<boolean> {
+    await this.connection
+    debug('set %s key with expiry %d', key, ttl)
+    return 'OK' === await this.client.set(key, value, { EX: ttl })
+  }
+
   public async removeRangeByScoreFromSortedSet(key: string, min: number, max: number): Promise<number> {
     await this.connection
     debug('remove %d..%d range from sorted set %s', min, max, key)

src/factories/message-handler-factory.ts

@@ -1,13 +1,23 @@
+import { ICacheAdapter, IWebSocketAdapter } from '../@types/adapters'
 import { IEventRepository, IUserRepository } from '../@types/repositories'
 import { IncomingMessage, MessageType } from '../@types/messages'
 import { createSettings } from './settings-factory'
 import { EventMessageHandler } from '../handlers/event-message-handler'
 import { eventStrategyFactory } from './event-strategy-factory'
-import { IWebSocketAdapter } from '../@types/adapters'
+import { getCacheClient } from '../cache/client'
+import { RedisAdapter } from '../adapters/redis-adapter'
 import { slidingWindowRateLimiterFactory } from './rate-limiter-factory'
 import { SubscribeMessageHandler } from '../handlers/subscribe-message-handler'
 import { UnsubscribeMessageHandler } from '../handlers/unsubscribe-message-handler'
 
+let cacheAdapter: ICacheAdapter | undefined = undefined
+const getCache = (): ICacheAdapter => {
+  if (!cacheAdapter) {
+    cacheAdapter = new RedisAdapter(getCacheClient())
+  }
+  return cacheAdapter
+}
+
 export const messageHandlerFactory = (
   eventRepository: IEventRepository,
   userRepository: IUserRepository,
@@ -22,6 +32,7 @@ export const messageHandlerFactory = (
           userRepository,
           createSettings,
           slidingWindowRateLimiterFactory,
+          getCache(),
         )
       }
     case MessageType.REQ:

src/handlers/event-message-handler.ts

@@ -18,13 +18,20 @@ import { IEventStrategy, IMessageHandler } from '../@types/message-handlers'
 import { createCommandResult } from '../utils/messages'
 import { createLogger } from '../factories/logger-factory'
 import { Factory } from '../@types/base'
+import { ICacheAdapter } from '../@types/adapters'
 import { IncomingEventMessage } from '../@types/messages'
 import { IRateLimiter } from '../@types/utils'
 import { IWebSocketAdapter } from '../@types/adapters'
 import { WebSocketAdapterEvent } from '../constants/adapter'
 
 const debug = createLogger('event-message-handler')
 
+export enum CacheAdmissionState {
+  ADMITTED = 'admitted',
+  NOT_ADMITTED = 'not-admitted',
+  INSUFFICIENT_BALANCE = 'insufficient-balance',
+}
+
 export class EventMessageHandler implements IMessageHandler {
   public constructor(
     protected readonly webSocket: IWebSocketAdapter,
@@ -33,6 +40,7 @@ export class EventMessageHandler implements IMessageHandler {
     protected readonly userRepository: IUserRepository,
     private readonly settings: () => Settings,
     private readonly slidingWindowRateLimiter: Factory<IRateLimiter>,
+    private readonly cache: ICacheAdapter,
   ) {}
 
   public async handleMessage(message: IncomingEventMessage): Promise<void> {
@@ -313,17 +321,45 @@ export class EventMessageHandler implements IMessageHandler {
       return
     }
 
-    // const hasKey = await this.cache.hasKey(`${event.pubkey}:is-admitted`)
-    // TODO: use cache
+    const cacheKey = `${event.pubkey}:is-admitted`
+
+    try {
+      const cachedAdmission = await this.cache.getKey(cacheKey)
+      switch (cachedAdmission) {
+        case CacheAdmissionState.ADMITTED:
+          debug('cache hit for %s admission: admitted', event.pubkey)
+          return
+        case CacheAdmissionState.NOT_ADMITTED:
+          debug('cache hit for %s admission: blocked', event.pubkey)
+          return 'blocked: pubkey not admitted'
+        case CacheAdmissionState.INSUFFICIENT_BALANCE:
+          debug('cache hit for %s admission: blocked (insufficient balance)', event.pubkey)
+          return 'blocked: insufficient balance'
+        default:
+          break
+      }
+    } catch (error) {
+      debug('cache error for %s: %o', event.pubkey, error)
+    }
+
     const user = await this.userRepository.findByPubkey(event.pubkey)
     if (!user || !user.isAdmitted) {
+      this.cacheSet(cacheKey, CacheAdmissionState.NOT_ADMITTED, 60)
       return 'blocked: pubkey not admitted'
     }
 
     const minBalance = currentSettings.limits?.event?.pubkey?.minBalance ?? 0n
     if (minBalance > 0n && user.balance < minBalance) {
+      this.cacheSet(cacheKey, CacheAdmissionState.INSUFFICIENT_BALANCE, 60)
       return 'blocked: insufficient balance'
     }
+
+    this.cacheSet(cacheKey, CacheAdmissionState.ADMITTED, 300)
+  }
+
+  private cacheSet(key: string, value: string, ttl: number): void {
+    this.cache.setKeyWithExpiry(key, value, ttl)
+      .catch((error) => debug('unable to cache %s: %o', key, error))
   }
 
   protected addExpirationMetadata(event: Event): Event | ExpiringEvent {

test/unit/handlers/event-message-handler.spec.ts

@@ -8,11 +8,11 @@ import sinonChai from 'sinon-chai'
 chai.use(sinonChai)
 chai.use(chaiAsPromised)
 
+import { CacheAdmissionState, EventMessageHandler } from '../../../src/handlers/event-message-handler'
 import { EventLimits, Settings } from '../../../src/@types/settings'
 import { IncomingEventMessage, MessageType } from '../../../src/@types/messages'
 import { Event } from '../../../src/@types/event'
 import { EventKinds } from '../../../src/constants/base'
-import { EventMessageHandler } from '../../../src/handlers/event-message-handler'
 import { IUserRepository } from '../../../src/@types/repositories'
 import { IWebSocketAdapter } from '../../../src/@types/adapters'
 import { WebSocketAdapterEvent } from '../../../src/constants/adapter'
@@ -88,7 +88,8 @@ describe('EventMessageHandler', () => {
         () => ({
           info: { relay_url: 'relay_url' },
         }) as any,
-        () => ({ hit: async () => false })
+        () => ({ hit: async () => false }),
+        { hasKey: async () => false, setKey: async () => true } as any,
       )
     })
 
@@ -262,7 +263,8 @@ describe('EventMessageHandler', () => {
         { hasActiveRequestToVanish: async () => false } as any,
         userRepository,
         () => settings,
-        () => ({ hit: async () => false })
+        () => ({ hit: async () => false }),
+        { hasKey: async () => false, setKey: async () => true } as any,
       )
     })
 
@@ -738,7 +740,8 @@ describe('EventMessageHandler', () => {
         { hasActiveRequestToVanish: async () => false } as any,
         userRepository,
         () => settings,
-        () => ({ hit: rateLimiterHitStub })
+        () => ({ hit: rateLimiterHitStub }),
+        { hasKey: async () => false, setKey: async () => true, setKeyWithExpiry: async () => true } as any,
       )
     })
 
@@ -953,6 +956,7 @@ describe('EventMessageHandler', () => {
     let webSocket: IWebSocketAdapter
     let getRelayPublicKeyStub: SinonStub
     let userRepositoryFindByPubkeyStub: SinonStub
+    let cacheGetKeyStub: SinonStub
 
     beforeEach(() => {
       settings = {
@@ -994,6 +998,7 @@ describe('EventMessageHandler', () => {
       getRelayPublicKeyStub = sandbox.stub(EventMessageHandler.prototype, 'getRelayPublicKey' as any)
       getClientAddressStub = sandbox.stub()
       userRepositoryFindByPubkeyStub = sandbox.stub()
+      cacheGetKeyStub = sandbox.stub().resolves(null)
       webSocket = {
         getClientAddress: getClientAddressStub,
       } as any
@@ -1006,11 +1011,18 @@ describe('EventMessageHandler', () => {
         { hasActiveRequestToVanish: async () => false } as any,
         userRepository,
         () => settings,
-        () => ({ hit: async () => false })
+        () => ({ hit: async () => false }),
+        {
+          hasKey: async () => false,
+          getKey: cacheGetKeyStub,
+          setKey: async () => true,
+          setKeyWithExpiry: async () => true,
+          setKeyExpiry: async () => undefined,
+        } as any,
       )
     })
 
-    it ('fulfills with undefined if payments are disabled', async () => {
+    it('fulfills with undefined if payments are disabled', async () => {
       settings.payments.enabled = false
 
       return expect((handler as any).isUserAdmitted(event)).to.eventually.be.undefined
@@ -1089,12 +1101,6 @@ describe('EventMessageHandler', () => {
       return expect((handler as any).isUserAdmitted(event)).to.eventually.equal('blocked: pubkey not admitted')
     })
 
-    it('fulfills with reason if user is not admitted', async () => {
-      userRepositoryFindByPubkeyStub.resolves({ isAdmitted: false })
-
-      return expect((handler as any).isUserAdmitted(event)).to.eventually.equal('blocked: pubkey not admitted')
-    })
-
     it('fulfills with reason if user does not meet minimum balance', async () => {
       settings.limits.event.pubkey.minBalance = 1000n
       userRepositoryFindByPubkeyStub.resolves({ isAdmitted: true, balance: 999n })
@@ -1108,5 +1114,23 @@ describe('EventMessageHandler', () => {
 
       return expect((handler as any).isUserAdmitted(event)).to.eventually.be.undefined
     })
+
+    it('fulfills with undefined if user is admitted in cache', async () => {
+      cacheGetKeyStub.resolves(CacheAdmissionState.ADMITTED)
+
+      return expect((handler as any).isUserAdmitted(event)).to.eventually.be.undefined
+    })
+
+    it('fulfills with reason if user is blocked in cache', async () => {
+      cacheGetKeyStub.resolves(CacheAdmissionState.NOT_ADMITTED)
+
+      return expect((handler as any).isUserAdmitted(event)).to.eventually.equal('blocked: pubkey not admitted')
+    })
+
+    it('fulfills with reason if user has insufficient balance in cache', async () => {
+      cacheGetKeyStub.resolves(CacheAdmissionState.INSUFFICIENT_BALANCE)
+
+      return expect((handler as any).isUserAdmitted(event)).to.eventually.equal('blocked: insufficient balance')
+    })
   })
 })