Implement in other custom transports (#713)

Author: ScottyPoi

packages/portalnetwork/src/client/client.ts

@@ -127,7 +127,7 @@ export class PortalNetwork extends EventEmitter<PortalNetworkEvents> {
     switch (opts.transport) {
       case TransportLayer.WEB: {
         opts.proxyAddress = opts.proxyAddress ?? 'ws://127.0.0.1:5050'
-        config.transport = new WebSocketTransportService(ma, config.enr.nodeId, opts.proxyAddress)
+        config.transport = new WebSocketTransportService(ma, config.enr.nodeId, opts.proxyAddress, new RateLimiter())
         break
       }
       case TransportLayer.MOBILE:
@@ -478,19 +478,19 @@ export class PortalNetwork extends EventEmitter<PortalNetworkEvents> {
 
   public addToBlackList = (ma: Multiaddr) => {
     (<RateLimiter>(
-      (<UDPTransportService>this.discv5.sessionService.transport)['rateLimiter']
+      (<any>this.discv5.sessionService.transport)['rateLimiter']
     )).addToBlackList(ma.nodeAddress().address)
   }
 
   public isBlackListed = (ma: Multiaddr) => {
     return (<RateLimiter>(
-      (<UDPTransportService>this.discv5.sessionService.transport)['rateLimiter']
+      (<any>this.discv5.sessionService.transport)['rateLimiter']
     )).isBlackListed(ma.nodeAddress().address)
   }
 
   public removeFromBlackList = (ma: Multiaddr) => {
     (<RateLimiter>(
-      (<UDPTransportService>this.discv5.sessionService.transport)['rateLimiter']
+      (<any>this.discv5.sessionService.transport)['rateLimiter']
     )).removeFromBlackList(ma.nodeAddress().address)
   }
 }

packages/portalnetwork/src/transports/capacitorUdp.ts

@@ -4,6 +4,7 @@ import { decodePacket, encodePacket } from '@chainsafe/discv5/packet'
 import { UDP } from '@frontall/capacitor-udp'
 import { multiaddr as ma } from '@multiformats/multiaddr'
 
+import type { IRateLimiter } from './rateLimiter.js'
 import type { SocketAddress } from '@chainsafe/discv5'
 import type { IPacket } from '@chainsafe/discv5/packet'
 import type {
@@ -33,11 +34,13 @@ export class CapacitorUDPTransportService
     ip4: true,
     ip6: false,
   }
-  public constructor(multiaddr: Multiaddr, srcId: string) {
+  private rateLimiter?: IRateLimiter
+  public constructor(multiaddr: Multiaddr, srcId: string, rateLimiter?: IRateLimiter) {
     //eslint-disable-next-line constructor-super
     super()
     this.bindAddrs = [multiaddr]
     this.srcId = srcId
+    this.rateLimiter = rateLimiter
   }
 
   public async start(): Promise<void> {
@@ -74,6 +77,9 @@ export class CapacitorUDPTransportService
   }
 
   public handleIncoming = (data: Uint8Array, rinfo: IRemoteInfo): void => {
+    if (this.rateLimiter && !this.rateLimiter.allowEncodedPacket(rinfo.address)) {
+      return
+    }
     const multiaddr = ma(
       `/${rinfo.family === 'IPv4' ? 'ip4' : 'ip6'}/${rinfo.address}/udp/${rinfo.port}`,
     )

packages/portalnetwork/src/transports/websockets.ts

@@ -18,7 +18,7 @@ import type { ENR } from '@chainsafe/enr'
 import type { Multiaddr } from '@multiformats/multiaddr'
 import type { Debugger } from 'debug'
 import type StrictEventEmitter from 'strict-event-emitter-types/types/src'
-
+import type { IRateLimiter } from './rateLimiter.js'
 const log = debug('discv5:transport')
 
 interface WebSocketTransportEvents extends ITransportEvents {
@@ -40,12 +40,13 @@ export class WebSocketTransportService
   private socket: WebSocketAsPromised
   private srcId: string
   private log: Debugger
+  private rateLimiter?: IRateLimiter
   ipMode: IPMode = {
     ip4: true,
     ip6: false,
   }
   bindAddrs: Multiaddr[] = []
-  public constructor(multiaddr: Multiaddr, srcId: string, proxyAddress: string) {
+  public constructor(multiaddr: Multiaddr, srcId: string, proxyAddress: string, rateLimiter?: IRateLimiter) {
     //eslint-disable-next-line constructor-super
     super()
     this.log = debug('Portal').extend('WebSocketTransportService')
@@ -58,6 +59,7 @@ export class WebSocketTransportService
       createWebSocket: (url) => new WebSocket(url),
       extractMessageData: (event) => event,
     })
+    this.rateLimiter = rateLimiter
   }
 
   public async start(): Promise<void> {
@@ -108,6 +110,9 @@ export class WebSocketTransportService
     const rinfo = JSON.parse(
       new TextDecoder().decode(data.slice(2, rinfoLength + 2)),
     ) as IRemoteInfo
+    if (this.rateLimiter && !this.rateLimiter.allowEncodedPacket(rinfo.address)) {
+      return;
+    }
     const multiaddr = ma(
       `/${rinfo.family === 'IPv4' ? 'ip4' : 'ip6'}/${rinfo.address}/udp/${rinfo.port}`,
     )