feat: add COUNT message validation, response helpers, repository counting, and handler routing

Author: a-khushal

src/@types/repositories.ts

@@ -33,6 +33,7 @@ export interface IEventRepository {
   upsert(event: Event): Promise<number>
   upsertMany(events: Event[]): Promise<number>
   findByFilters(filters: SubscriptionFilter[]): IQueryResult<DBEvent[]>
+  countByFilters(filters: SubscriptionFilter[]): Promise<number>
   deleteByPubkeyAndIds(pubkey: Pubkey, ids: EventId[]): Promise<number>
   deleteByPubkeyExceptKinds(pubkey: Pubkey, excludedKinds: number[]): Promise<number>
   hasActiveRequestToVanish(pubkey: Pubkey): Promise<boolean>

src/factories/message-handler-factory.ts

@@ -2,6 +2,7 @@ import { ICacheAdapter, IWebSocketAdapter } from '../@types/adapters'
 import { IEventRepository, INip05VerificationRepository, IUserRepository } from '../@types/repositories'
 import { IncomingMessage, MessageType } from '../@types/messages'
 import { createSettings } from './settings-factory'
+import { CountMessageHandler } from '../handlers/count-message-handler'
 import { EventMessageHandler } from '../handlers/event-message-handler'
 import { eventStrategyFactory } from './event-strategy-factory'
 import { getCacheClient } from '../cache/client'
@@ -42,6 +43,8 @@ export const messageHandlerFactory =
         return new SubscribeMessageHandler(adapter, eventRepository, createSettings)
       case MessageType.CLOSE:
         return new UnsubscribeMessageHandler(adapter)
+      case MessageType.COUNT:
+        return new CountMessageHandler(adapter, eventRepository, createSettings)
       default:
         throw new Error(`Unknown message type: ${String(message[0]).substring(0, 64)}`)
     }

src/handlers/count-message-handler.ts

@@ -0,0 +1,61 @@
+import { equals, uniqWith } from 'ramda'
+
+import { IWebSocketAdapter } from '../@types/adapters'
+import { IMessageHandler } from '../@types/message-handlers'
+import { CountMessage } from '../@types/messages'
+import { IEventRepository } from '../@types/repositories'
+import { Settings } from '../@types/settings'
+import { SubscriptionFilter } from '../@types/subscription'
+import { WebSocketAdapterEvent } from '../constants/adapter'
+import { createLogger } from '../factories/logger-factory'
+import { createClosedMessage, createCountResultMessage } from '../utils/messages'
+
+const debug = createLogger('count-message-handler')
+
+export class CountMessageHandler implements IMessageHandler {
+  public constructor(
+    private readonly webSocket: IWebSocketAdapter,
+    private readonly eventRepository: IEventRepository,
+    private readonly settings: () => Settings,
+  ) {}
+
+  public async handleMessage(message: CountMessage): Promise<void> {
+    const queryId = message[1]
+    // Some clients send the same filter more than once.
+    // We remove duplicates so we do less DB work.
+    const filters = uniqWith(equals, message.slice(2)) as SubscriptionFilter[]
+
+    const reason = this.canCount(queryId, filters)
+    if (reason) {
+      debug('count request %s with %o rejected: %s', queryId, filters, reason)
+      // NIP-45 says we should close rejected COUNT requests with a reason.
+      this.webSocket.emit(WebSocketAdapterEvent.Message, createClosedMessage(queryId, reason))
+      return
+    }
+
+    try {
+      const count = await this.eventRepository.countByFilters(filters)
+      this.webSocket.emit(WebSocketAdapterEvent.Message, createCountResultMessage(queryId, { count }))
+    } catch (error) {
+      debug('count request %s failed: %o', queryId, error)
+      // Keep this message generic so internal errors are not leaked to clients.
+      this.webSocket.emit(WebSocketAdapterEvent.Message, createClosedMessage(queryId, 'error: unable to count events'))
+    }
+  }
+
+  private canCount(queryId: string, filters: SubscriptionFilter[]): string | undefined {
+    const subscriptionLimits = this.settings().limits?.client?.subscription
+    const maxFilters = subscriptionLimits?.maxFilters ?? 0
+
+    if (maxFilters > 0 && filters.length > maxFilters) {
+      return `Too many filters: Number of filters per count query must be less then or equal to ${maxFilters}`
+    }
+
+    if (
+      typeof subscriptionLimits?.maxSubscriptionIdLength === 'number' &&
+      queryId.length > subscriptionLimits.maxSubscriptionIdLength
+    ) {
+      return `Query ID too long: Query ID must be less or equal to ${subscriptionLimits.maxSubscriptionIdLength}`
+    }
+  }
+}

src/repositories/event-repository.ts

@@ -168,6 +168,109 @@ export class EventRepository implements IEventRepository {
     return query
   }
 
+  public async countByFilters(filters: SubscriptionFilter[]): Promise<number> {
+    debug('counting events for %o', filters)
+
+    if (!Array.isArray(filters) || !filters.length) {
+      throw new Error('Filters cannot be empty')
+    }
+
+    const now = Math.floor(Date.now() / 1000)
+
+    const queries = filters.map((currentFilter) => {
+      const builder = this.readReplicaDbClient<DBEvent>('events').select('events.event_id')
+
+      forEachObjIndexed((tableFields: string[], filterName: string | number) => {
+        builder.andWhere((bd) => {
+          cond([
+            [isEmpty, () => void bd.whereRaw('1 = 0')],
+            [
+              complement(isNil),
+              pipe(
+                groupByLengthSpec,
+                evolve({
+                  exact: (pubkeys: string[]) =>
+                    tableFields.forEach((tableField) => bd.orWhereIn(tableField, pubkeys.map(toBuffer))),
+                  even: forEach((prefix: string) =>
+                    tableFields.forEach((tableField) =>
+                      bd.orWhereRaw(`substring("${tableField}" from 1 for ?) = ?`, [prefix.length >> 1, toBuffer(prefix)]),
+                    ),
+                  ),
+                  odd: forEach((prefix: string) =>
+                    tableFields.forEach((tableField) =>
+                      bd.orWhereRaw(`substring("${tableField}" from 1 for ?) BETWEEN ? AND ?`, [
+                        (prefix.length >> 1) + 1,
+                        `\\x${prefix}0`,
+                        `\\x${prefix}f`,
+                      ]),
+                    ),
+                  ),
+                } as any),
+              ),
+            ],
+          ])(currentFilter[filterName] as string[])
+        })
+      })({ authors: ['event_pubkey'], ids: ['event_id'] })
+
+      if (Array.isArray(currentFilter.kinds)) {
+        builder.whereIn('event_kind', currentFilter.kinds)
+      }
+
+      if (typeof currentFilter.since === 'number') {
+        builder.where('event_created_at', '>=', currentFilter.since)
+      }
+
+      if (typeof currentFilter.until === 'number') {
+        builder.where('event_created_at', '<=', currentFilter.until)
+      }
+
+      if (typeof currentFilter.limit === 'number') {
+        builder.limit(currentFilter.limit).orderBy('event_created_at', 'DESC').orderBy('event_id', 'asc')
+      }
+
+      const andWhereRaw = invoker(1, 'andWhereRaw')
+      const orWhereRaw = invoker(2, 'orWhereRaw')
+
+      let isTagQuery = false
+      pipe(
+        toPairs,
+        filter(pipe(nth(0) as () => string, isGenericTagQuery)) as any,
+        forEach(([filterName, criteria]: [string, string[]]) => {
+          isTagQuery = true
+          builder.andWhere((bd) => {
+            ifElse(
+              isEmpty,
+              () => andWhereRaw('1 = 0', bd),
+              forEach(
+                (criterion: string) =>
+                  void orWhereRaw('event_tags.tag_name = ? AND event_tags.tag_value = ?', [filterName[1], criterion], bd),
+              ),
+            )(criteria)
+          })
+        }),
+      )(currentFilter as any)
+
+      if (isTagQuery) {
+        builder.leftJoin('event_tags', 'events.event_id', 'event_tags.event_id').select('events.event_id')
+      }
+
+      builder.whereNull('events.deleted_at').andWhere((bd) => {
+        bd.whereNull('events.expires_at').orWhere('events.expires_at', '>', now)
+      })
+
+      return builder
+    })
+
+    const [query, ...subqueries] = queries
+    if (subqueries.length) {
+      query.union(subqueries, true)
+    }
+
+    const result = await this.readReplicaDbClient.from(query.as('matching_events')).countDistinct({ count: 'event_id' }).first()
+
+    return Number(result?.count ?? 0)
+  }
+
   public async create(event: Event): Promise<number> {
     return this.insert(event).then(prop('rowCount') as () => number, () => 0)
   }

src/schemas/message-schema.ts

@@ -30,6 +30,29 @@ export const reqMessageSchema = z
     }
   })
 
+export const countMessageSchema = z
+  .tuple([z.literal(MessageType.COUNT), z.string().max(256).min(1)])
+  .rest(filterSchema)
+  .superRefine((val, ctx) => {
+    if (val.length < 3) {
+      ctx.addIssue({
+        code: z.ZodIssueCode.too_small,
+        minimum: 3,
+        type: 'array',
+        inclusive: true,
+        message: 'COUNT message must contain at least one filter',
+      })
+    } else if (val.length > 12) {
+      ctx.addIssue({
+        code: z.ZodIssueCode.too_big,
+        maximum: 12,
+        type: 'array',
+        inclusive: true,
+        message: 'COUNT message must contain at most 12 elements',
+      })
+    }
+  })
+
 export const closeMessageSchema = z.tuple([z.literal(MessageType.CLOSE), subscriptionSchema])
 
-export const messageSchema = z.union([eventMessageSchema, reqMessageSchema, closeMessageSchema])
+export const messageSchema = z.union([eventMessageSchema, reqMessageSchema, closeMessageSchema, countMessageSchema])

src/utils/messages.ts

@@ -1,4 +1,7 @@
 import {
+  ClosedMessage,
+  CountResultMessage,
+  CountResultPayload,
   EndOfStoredEventsNotice,
   IncomingEventMessage,
   IncomingRelayedEventMessage,
@@ -29,6 +32,15 @@ export const createCommandResult = (eventId: EventId, successful: boolean, messa
   return [MessageType.OK, eventId, successful, message]
 }
 
+// NIP-45
+export const createCountResultMessage = (queryId: SubscriptionId, payload: CountResultPayload): CountResultMessage => {
+  return [MessageType.COUNT, queryId, payload]
+}
+
+export const createClosedMessage = (queryId: SubscriptionId, reason: string): ClosedMessage => {
+  return [MessageType.CLOSED, queryId, reason]
+}
+
 export const createSubscriptionMessage = (
   subscriptionId: SubscriptionId,
   filters: SubscriptionFilter[],