devp2p -> DPT: improved peer related typings

Author: holgerd77

packages/devp2p/src/dpt/dpt.ts

@@ -10,6 +10,13 @@ import { Server as DPTServer } from './server'
 
 const debug = createDebugLogger('devp2p:dpt')
 
+export interface PeerInfo {
+  id?: Uint8Array | Buffer
+  address?: string
+  udpPort?: number | null
+  tcpPort?: number | null
+}
+
 export class DPT extends EventEmitter {
   privateKey: Buffer
   banlist: BanList
@@ -28,8 +35,8 @@ export class DPT extends EventEmitter {
     this.banlist = new BanList()
 
     this._kbucket = new KBucket(this._id)
-    this._kbucket.on('added', (peer) => this.emit('peer:added', peer))
-    this._kbucket.on('removed', (peer) => this.emit('peer:removed', peer))
+    this._kbucket.on('added', (peer: PeerInfo) => this.emit('peer:added', peer))
+    this._kbucket.on('removed', (peer: PeerInfo) => this.emit('peer:removed', peer))
     this._kbucket.on('ping', this._onKBucketPing)
 
     this._server = new DPTServer(this, this.privateKey, {
@@ -55,7 +62,7 @@ export class DPT extends EventEmitter {
     this._server.destroy(...args)
   }
 
-  _onKBucketPing(oldPeers: any[], newPeer: any): void {
+  _onKBucketPing(oldPeers: PeerInfo[], newPeer: PeerInfo): void {
     if (this.banlist.has(newPeer)) return
 
     let count = 0
@@ -81,15 +88,15 @@ export class DPT extends EventEmitter {
     for (const peer of peers) this.addPeer(peer).catch(() => {})
   }
 
-  async bootstrap(peer: any): Promise<void> {
+  async bootstrap(peer: PeerInfo): Promise<void> {
     debug(`bootstrap with peer ${peer.address}:${peer.udpPort}`)
 
     peer = await this.addPeer(peer)
     if (!this._id) return
     this._server.findneighbours(peer, this._id)
   }
 
-  async addPeer(obj: any): Promise<any> {
+  async addPeer(obj: PeerInfo): Promise<any> {
     if (this.banlist.has(obj)) throw new Error('Peer is banned')
     debug(`attempt adding peer ${obj.address}:${obj.udpPort}`)
 
@@ -109,23 +116,23 @@ export class DPT extends EventEmitter {
     }
   }
 
-  getPeer(obj: any): any {
+  getPeer(obj: string | Buffer | PeerInfo) {
     return this._kbucket.get(obj)
   }
 
-  getPeers(): any[] {
+  getPeers() {
     return this._kbucket.getAll()
   }
 
-  getClosestPeers(id: string): any {
+  getClosestPeers(id: string) {
     return this._kbucket.closest(id)
   }
 
   removePeer(obj: any) {
     this._kbucket.remove(obj)
   }
 
-  banPeer(obj: any, maxAge?: number) {
+  banPeer(obj: string | Buffer | PeerInfo, maxAge?: number) {
     this.banlist.add(obj, maxAge)
     this._kbucket.remove(obj)
   }

packages/devp2p/src/dpt/kbucket.ts

@@ -1,56 +1,58 @@
 import { EventEmitter } from 'events'
 import _KBucket = require('k-bucket')
+import { PeerInfo } from './dpt'
 
 const KBUCKET_SIZE = 16
 const KBUCKET_CONCURRENCY = 3
 
-export interface KObj {
-  id?: string
-  port?: string
-  address?: string
+export interface CustomContact extends PeerInfo {
+  id: Uint8Array | Buffer
+  vectorClock: number
 }
 
 export class KBucket extends EventEmitter {
-  _peers: Map<string, any> = new Map()
+  _peers: Map<string, PeerInfo> = new Map()
   _kbucket: _KBucket
   constructor(id: string | Buffer) {
     super()
 
-    this._kbucket = new _KBucket({
+    this._kbucket = new _KBucket<CustomContact>({
       localNodeId: typeof id === 'string' ? Buffer.from(id) : id,
       numberOfNodesPerKBucket: KBUCKET_SIZE,
       numberOfNodesToPing: KBUCKET_CONCURRENCY,
     })
 
-    this._kbucket.on('added', (peer: any) => {
+    this._kbucket.on('added', (peer: PeerInfo) => {
       KBucket.getKeys(peer).forEach((key) => this._peers.set(key, peer))
       this.emit('added', peer)
     })
 
-    this._kbucket.on('removed', (peer: any) => {
+    this._kbucket.on('removed', (peer: PeerInfo) => {
       KBucket.getKeys(peer).forEach((key) => this._peers.delete(key))
       this.emit('removed', peer)
     })
 
-    this._kbucket.on('ping', (...args: any[]) => this.emit('ping', ...args))
+    this._kbucket.on('ping', (oldPeers: PeerInfo[], newPeer: PeerInfo) => {
+      this.emit('ping', { oldPeers, newPeer })
+    })
   }
 
-  static getKeys(obj: Buffer | string | KObj): string[] {
+  static getKeys(obj: Buffer | string | PeerInfo): string[] {
     if (Buffer.isBuffer(obj)) return [obj.toString('hex')]
     if (typeof obj === 'string') return [obj]
 
     const keys = []
     if (Buffer.isBuffer(obj.id)) keys.push(obj.id.toString('hex'))
-    if (obj.address && obj.port) keys.push(`${obj.address}:${obj.port}`)
+    //if (obj.address && obj.port) keys.push(`${obj.address}:${obj.port}`)
     return keys
   }
 
-  add(peer: any) {
+  add(peer: PeerInfo) {
     const isExists = KBucket.getKeys(peer).some((key) => this._peers.has(key))
-    if (!isExists) this._kbucket.add(peer)
+    if (!isExists) this._kbucket.add(peer as CustomContact)
   }
 
-  get(obj: Buffer | string | KObj) {
+  get(obj: Buffer | string | PeerInfo) {
     for (const key of KBucket.getKeys(obj)) {
       const peer = this._peers.get(key)
       if (peer !== undefined) return peer
@@ -59,16 +61,16 @@ export class KBucket extends EventEmitter {
     return null
   }
 
-  getAll(): Array<any> {
+  getAll(): Array<PeerInfo> {
     return this._kbucket.toArray()
   }
 
-  closest(id: string): any {
+  closest(id: string): PeerInfo[] {
     return this._kbucket.closest(Buffer.from(id), KBUCKET_SIZE)
   }
 
-  remove(obj: Buffer | string | KObj) {
+  remove(obj: Buffer | string | PeerInfo) {
     const peer = this.get(obj)
-    if (peer !== null) this._kbucket.remove(peer.id)
+    if (peer !== null) this._kbucket.remove((peer as CustomContact).id)
   }
 }

packages/devp2p/src/dpt/message.ts

@@ -3,20 +3,14 @@ import ip from 'ip'
 import * as rlp from 'rlp'
 import secp256k1 from 'secp256k1'
 import { keccak256, int2buffer, buffer2int, assertEq, unstrictDecode } from '../util'
+import { PeerInfo } from './dpt'
 
 const debug = createDebugLogger('devp2p:dpt:server')
 
 function getTimestamp() {
   return (Date.now() / 1000) | 0
 }
 
-export interface PeerInfo {
-  id?: Buffer
-  address?: string
-  udpPort?: number | null
-  tcpPort?: number | null
-}
-
 const timestamp = {
   encode: function (value = getTimestamp() + 60) {
     const buffer = Buffer.allocUnsafe(4)
@@ -133,7 +127,7 @@ type OutNeighborMsg = { [0]: Buffer[][]; [1]: Buffer }
 const neighbours = {
   encode: function (obj: InNeighborMsg): OutNeighborMsg {
     return [
-      obj.peers.map((peer: PeerInfo) => endpoint.encode(peer).concat(peer.id!)),
+      obj.peers.map((peer: PeerInfo) => endpoint.encode(peer).concat(peer.id! as Buffer)),
       timestamp.encode(obj.timestamp),
     ]
   },

packages/devp2p/src/dpt/server.ts

@@ -5,9 +5,8 @@ import { debug as createDebugLogger } from 'debug'
 import LRUCache = require('lru-cache')
 import { encode, decode } from './message'
 import { keccak256, pk2id, createDeferred, formatLogId } from '../util'
-import { DPT } from './dpt'
+import { DPT, PeerInfo } from './dpt'
 import { Socket as DgramSocket, RemoteInfo } from 'dgram'
-import { PeerInfo } from './message'
 
 const debug = createDebugLogger('devp2p:dpt:server')
 const verbose = createDebugLogger('verbose').enabled