Skip to content

Commit

Permalink
feat: support pubkey/ip whitelisting for event rate limits
Browse files Browse the repository at this point in the history
  • Loading branch information
cameri committed Dec 27, 2022
1 parent 3620026 commit 845dedc
Show file tree
Hide file tree
Showing 2 changed files with 104 additions and 9 deletions.
17 changes: 15 additions & 2 deletions src/handlers/event-message-handler.ts
Original file line number Diff line number Diff line change
Expand Up @@ -126,9 +126,22 @@ export class EventMessageHandler implements IMessageHandler {
}

protected async isRateLimited(event: Event): Promise<boolean> {
const rateLimits = this.settings().limits?.event?.rateLimits
const { whitelists, rateLimits } = this.settings().limits?.event ?? {}
if (!rateLimits || !rateLimits.length) {
return
return false
}

if (
Array.isArray(whitelists?.pubkeys)
&& whitelists.pubkeys.includes(event.pubkey)
) {
return false
}

if (Array.isArray(whitelists?.ipAddresses)
&& whitelists.ipAddresses.includes(this.webSocket.getClientAddress())
) {
return false
}

const rateLimiter = this.slidingWindowRateLimiter()
Expand Down
96 changes: 89 additions & 7 deletions test/unit/handlers/event-message-handler.spec.ts
Original file line number Diff line number Diff line change
Expand Up @@ -10,12 +10,13 @@ import { EventLimits, ISettings } from '../../../src/@types/settings'
import { IncomingEventMessage, MessageType } from '../../../src/@types/messages'
import { Event } from '../../../src/@types/event'
import { EventMessageHandler } from '../../../src/handlers/event-message-handler'
import { IWebSocketAdapter } from '../../../src/@types/adapters'
import { WebSocketAdapterEvent } from '../../../src/constants/adapter'

const { expect } = chai

describe('EventMessageHandler', () => {
let webSocket: EventEmitter
let webSocket: IWebSocketAdapter
let handler: EventMessageHandler
let event: Event
let message: IncomingEventMessage
Expand Down Expand Up @@ -59,7 +60,7 @@ describe('EventMessageHandler', () => {
execute: strategyExecuteStub,
})
onMessageSpy = sandbox.fake.returns(undefined)
webSocket = new EventEmitter()
webSocket = new EventEmitter() as any
webSocket.on(WebSocketAdapterEvent.Message, onMessageSpy)
message = [MessageType.EVENT, event]
isRateLimitedStub = sandbox.stub(EventMessageHandler.prototype, 'isRateLimited' as any)
Expand Down Expand Up @@ -548,6 +549,8 @@ describe('EventMessageHandler', () => {
let eventLimits: EventLimits
let settings: ISettings
let rateLimiterHitStub: SinonStub
let getClientAddressStub: Sinon.SinonStub
let webSocket: IWebSocketAdapter

beforeEach(() => {
eventLimits = {
Expand All @@ -559,22 +562,101 @@ describe('EventMessageHandler', () => {
},
} as any
rateLimiterHitStub = sandbox.stub()
getClientAddressStub = sandbox.stub()
webSocket = {
getClientAddress: getClientAddressStub,
} as any
handler = new EventMessageHandler(
{} as any,
webSocket,
() => null,
() => settings,
() => ({ hit: rateLimiterHitStub })
)
})

it('returns undefined if rate limits setting is not set', async () => {
it('fulfills with false if limits setting is not set', async () => {
settings.limits = undefined
return expect((handler as any).isRateLimited(event)).to.eventually.be.false
})


it('fulfills with false if event limits setting is not set', async () => {
settings.limits.event = undefined
return expect((handler as any).isRateLimited(event)).to.eventually.be.false
})

it('fulfills with false if rate limits setting is not set', async () => {
eventLimits.rateLimits = undefined
return expect((handler as any).isRateLimited(event)).to.eventually.be.undefined
return expect((handler as any).isRateLimited(event)).to.eventually.be.false
})

it('returns undefined if rate limits setting is empty', async () => {
it('fulfills with false if rate limits setting is empty', async () => {
eventLimits.rateLimits = []
return expect((handler as any).isRateLimited(event)).to.eventually.be.undefined
return expect((handler as any).isRateLimited(event)).to.eventually.be.false
})

it('skips rate limiter if IP is whitelisted', async () => {
eventLimits.rateLimits = [
{
period: 60000,
rate: 1,
},
]
eventLimits.whitelists = {}
eventLimits.whitelists.ipAddresses = ['2604:a880:cad:d0::e7e:7001']
getClientAddressStub.returns('2604:a880:cad:d0::e7e:7001')

const actualResult = await (handler as any).isRateLimited(event)

expect(actualResult).to.be.false
expect(rateLimiterHitStub).not.to.have.been.called
})

it('calls rate limiter if IP is not whitelisted', async () => {
eventLimits.rateLimits = [
{
period: 60000,
rate: 1,
},
]
eventLimits.whitelists = {}
eventLimits.whitelists.ipAddresses = ['::1']
getClientAddressStub.returns('2604:a880:cad:d0::e7e:7001')

await (handler as any).isRateLimited(event)

expect(rateLimiterHitStub).to.have.been.called
})

it('skips rate limiter if pubkey is whitelisted', async () => {
eventLimits.rateLimits = [
{
period: 60000,
rate: 1,
},
]
eventLimits.whitelists = {}
eventLimits.whitelists.pubkeys = [event.pubkey]

const actualResult = await (handler as any).isRateLimited(event)

expect(actualResult).to.be.false
expect(rateLimiterHitStub).not.to.have.been.called
})

it('calls rate limiter if pubkey is not whitelisted', async () => {
eventLimits.rateLimits = [
{
period: 60000,
rate: 1,
},
]
eventLimits.whitelists = {}
eventLimits.whitelists.pubkeys = ['other']

await (handler as any).isRateLimited(event)

expect(rateLimiterHitStub).to.have.been.called
})

it('calls hit with given rate limit settings', async () => {
Expand Down

0 comments on commit 845dedc

Please sign in to comment.