From 041014e5682c32e4499e946944b10f7780fcc8b3 Mon Sep 17 00:00:00 2001
From: dozyio <github@dozy.io>
Date: Fri, 7 Mar 2025 01:42:25 +0000
Subject: [PATCH] feat: allow custom registrar

---
 packages/libp2p/src/index.ts                  |  6 ++
 packages/libp2p/src/libp2p.ts                 |  5 ++
 .../test/registrar/custom-registrar.spec.ts   | 77 +++++++++++++++++++
 3 files changed, 88 insertions(+)
 create mode 100644 packages/libp2p/test/registrar/custom-registrar.spec.ts

diff --git a/packages/libp2p/src/index.ts b/packages/libp2p/src/index.ts
index fde1a69ed1..80a0c0122d 100644
--- a/packages/libp2p/src/index.ts
+++ b/packages/libp2p/src/index.ts
@@ -24,6 +24,7 @@ import type { ConnectionManagerInit } from './connection-manager/index.js'
 import type { ConnectionMonitorInit } from './connection-monitor.js'
 import type { TransportManagerInit } from './transport-manager.js'
 import type { Libp2p, ServiceMap, ComponentLogger, NodeInfo, ConnectionProtector, ConnectionEncrypter, ConnectionGater, ContentRouting, Metrics, PeerDiscovery, PeerRouting, StreamMuxerFactory, Transport, PrivateKey } from '@libp2p/interface'
+import type { Registrar } from '@libp2p/interface-internal'
 import type { PersistentPeerStoreInit } from '@libp2p/peer-store'
 import type { DNS } from '@multiformats/dns'
 import type { Datastore } from 'interface-datastore'
@@ -132,6 +133,11 @@ export interface Libp2pInit<T extends ServiceMap = ServiceMap> {
    */
   connectionProtector?(components: Components): ConnectionProtector
 
+  /**
+   * A registrar implementation that can be used to modify protocol handling
+   */
+  registrar?(components: Components): Registrar
+
   /**
    * Arbitrary libp2p modules
    */
diff --git a/packages/libp2p/src/libp2p.ts b/packages/libp2p/src/libp2p.ts
index 4d6b62ba82..0abc98c6ae 100644
--- a/packages/libp2p/src/libp2p.ts
+++ b/packages/libp2p/src/libp2p.ts
@@ -135,6 +135,11 @@ export class Libp2p<T extends ServiceMap = ServiceMap> extends TypedEventEmitter
     // Create the Registrar
     this.configureComponent('registrar', new Registrar(this.components))
 
+    // Setup a custom registrar that can override the default registrar
+    if (init.registrar != null) {
+      this.configureComponent('registrar', init.registrar(this.components))
+    }
+
     // Addresses {listen, announce, noAnnounce}
     this.configureComponent('addressManager', new AddressManager(this.components, init.addresses))
 
diff --git a/packages/libp2p/test/registrar/custom-registrar.spec.ts b/packages/libp2p/test/registrar/custom-registrar.spec.ts
new file mode 100644
index 0000000000..485af8bcd1
--- /dev/null
+++ b/packages/libp2p/test/registrar/custom-registrar.spec.ts
@@ -0,0 +1,77 @@
+/* eslint-env mocha */
+
+import { expect } from 'aegir/chai'
+import sinon from 'sinon'
+import { createLibp2p } from '../../src/index.js'
+import type { Libp2p } from '@libp2p/interface'
+import type { Registrar } from '@libp2p/interface-internal'
+
+describe('custom registrar functionality', () => {
+  let libp2p: Libp2p
+  let customRegistrar: {
+    handle: sinon.SinonStub<any, Promise<void>>
+    unhandle: sinon.SinonStub<any, Promise<void>>
+    getProtocols: sinon.SinonStub<any, string[]>
+    register: sinon.SinonStub<any, Promise<string>>
+    unregister: sinon.SinonStub<any, void>
+    getHandler: sinon.SinonStub<any, { handler(): void, options: Record<string, unknown> }>
+    getTopologies: sinon.SinonStub<any, any[]>
+  }
+
+  beforeEach(async (): Promise<void> => {
+    customRegistrar = {
+      handle: sinon.stub().resolves(),
+      unhandle: sinon.stub().resolves(),
+      getProtocols: sinon.stub().returns(['/custom/1.0.0']),
+      register: sinon.stub().resolves('custom-topology-id'),
+      unregister: sinon.stub(),
+      getHandler: sinon.stub().returns({ handler: (): void => {}, options: {} }),
+      getTopologies: sinon.stub().returns([])
+    }
+
+    libp2p = await createLibp2p({
+      registrar: (_components) => customRegistrar as unknown as Registrar
+    })
+  })
+
+  afterEach(async (): Promise<void> => {
+    await libp2p?.stop()
+    sinon.restore()
+  })
+
+  it('should reflect custom getProtocols implementation', (): void => {
+    const protocols: string[] = libp2p.getProtocols()
+    expect(protocols).to.deep.equal(['/custom/1.0.0'])
+    expect(customRegistrar.getProtocols.callCount).to.equal(1)
+  })
+
+  it('should call custom registrar handle method when registering a protocol handler', async (): Promise<void> => {
+    const testHandler = (): void => {}
+    await libp2p.handle('/custom/1.0.0', testHandler)
+    expect(customRegistrar.handle.callCount).to.equal(1)
+    expect(customRegistrar.handle.firstCall.args).to.deep.equal(['/custom/1.0.0', testHandler, undefined])
+  })
+
+  it('should call custom registrar unhandle method when unregistering a protocol handler', async (): Promise<void> => {
+    await libp2p.unhandle('/custom/1.0.0')
+    expect(customRegistrar.unhandle.callCount).to.equal(1)
+    expect(customRegistrar.unhandle.firstCall.args).to.deep.equal(['/custom/1.0.0'])
+  })
+
+  it('should use custom registrar register method for topologies', async (): Promise<void> => {
+    const topology = {
+      onConnect: (): void => {},
+      onDisconnect: (): void => {}
+    }
+    const topologyId: string = await libp2p.register('/custom/1.0.0', topology)
+    expect(topologyId).to.equal('custom-topology-id')
+    expect(customRegistrar.register.callCount).to.equal(1)
+    expect(customRegistrar.register.firstCall.args).to.deep.equal(['/custom/1.0.0', topology])
+  })
+
+  it('should call custom registrar unregister method for topologies', (): void => {
+    libp2p.unregister('custom-topology-id')
+    expect(customRegistrar.unregister.callCount).to.equal(1)
+    expect(customRegistrar.unregister.firstCall.args).to.deep.equal(['custom-topology-id'])
+  })
+})