| 1 | package harness |
| 2 | |
| 3 | import ( |
| 4 | "context" |
| 5 | "encoding/hex" |
| 6 | "sync" |
| 7 | |
| 8 | "github.com/libp2p/go-libp2p" |
| 9 | dht "github.com/libp2p/go-libp2p-kad-dht" |
| 10 | "github.com/libp2p/go-libp2p-kad-dht/records" |
| 11 | "github.com/libp2p/go-libp2p/core/host" |
| 12 | "github.com/libp2p/go-libp2p/core/peer" |
| 13 | ) |
| 14 | |
| 15 | // stubPeerPool manages ephemeral in-process libp2p/DHT peers for |
| 16 | // TEST_DHT_STUB mode. |
| 17 | // |
| 18 | // All peers share a single in-memory ProviderStore. This store is |
| 19 | // NOT shared with the kubo daemons; it lives in the test process. |
| 20 | // When a kubo daemon sends ADD_PROVIDER to any ephemeral peer, the |
| 21 | // record is stored in this shared store. When another kubo daemon |
| 22 | // queries GET_PROVIDERS from any peer, it finds the record because |
| 23 | // all peers see the same store. The kubo daemons communicate with |
| 24 | // the ephemeral peers via real DHT protocol messages over loopback |
| 25 | // TCP. |
| 26 | type stubPeerPool struct { |
| 27 | hosts []host.Host |
| 28 | dhts []*dht.IpfsDHT |
| 29 | store *sharedMemStore |
| 30 | cancel context.CancelFunc |
| 31 | } |
| 32 | |
| 33 | // stubDHTPeerCount is the number of ephemeral DHT peers to create. |
| 34 | // Matches amino.DefaultBucketSize (K=20 in Kademlia), ensuring |
| 35 | // GetClosestPeers always finds enough peers for provide replication. |
| 36 | const stubDHTPeerCount = 20 |
| 37 | |
| 38 | // newStubPeerPool creates count ephemeral DHT peers on loopback and |
| 39 | // mesh-connects them. |
| 40 | func newStubPeerPool(count int) (*stubPeerPool, error) { |
| 41 | ctx, cancel := context.WithCancel(context.Background()) |
| 42 | |
| 43 | store := &sharedMemStore{data: make(map[string][]peer.AddrInfo)} |
| 44 | |
| 45 | hosts := make([]host.Host, 0, count) |
| 46 | dhts := make([]*dht.IpfsDHT, 0, count) |
| 47 | |
| 48 | cleanup := func() { |
| 49 | for _, d := range dhts { |
| 50 | d.Close() |
| 51 | } |
| 52 | for _, h := range hosts { |
| 53 | h.Close() |
| 54 | } |
| 55 | cancel() |
| 56 | } |
| 57 | |
| 58 | for range count { |
| 59 | h, err := libp2p.New(libp2p.ListenAddrStrings("/ip4/127.0.0.1/tcp/0")) |
| 60 | if err != nil { |
| 61 | cleanup() |
| 62 | return nil, err |
| 63 | } |
| 64 | d, err := dht.New(ctx, h, |
| 65 | dht.Mode(dht.ModeServer), |
| 66 | dht.ProviderStore(store), |
| 67 | dht.AddressFilter(nil), |
| 68 | dht.DisableAutoRefresh(), |
| 69 | dht.BootstrapPeers(), |
| 70 | ) |
| 71 | if err != nil { |
| 72 | h.Close() |
| 73 | cleanup() |
| 74 | return nil, err |
| 75 | } |
| 76 | hosts = append(hosts, h) |
| 77 | dhts = append(dhts, d) |
| 78 | } |
| 79 | |
| 80 | // Full-mesh connect so routing tables are populated. |
| 81 | for i, h := range hosts { |
| 82 | for j, other := range hosts { |
| 83 | if i == j { |
| 84 | continue |
| 85 | } |
| 86 | ai := peer.AddrInfo{ID: other.ID(), Addrs: other.Addrs()} |
| 87 | if err := h.Connect(ctx, ai); err != nil { |
| 88 | cleanup() |
| 89 | return nil, err |
| 90 | } |
| 91 | } |
| 92 | } |
| 93 | |
| 94 | return &stubPeerPool{ |
| 95 | hosts: hosts, |
| 96 | dhts: dhts, |
| 97 | store: store, |
| 98 | cancel: cancel, |
| 99 | }, nil |
| 100 | } |
| 101 | |
| 102 | func (p *stubPeerPool) Close() { |
| 103 | if p == nil { |
| 104 | return |
| 105 | } |
| 106 | for _, d := range p.dhts { |
| 107 | d.Close() |
| 108 | } |
| 109 | for _, h := range p.hosts { |
| 110 | h.Close() |
| 111 | } |
| 112 | p.cancel() |
| 113 | } |
| 114 | |
| 115 | // sharedMemStore implements records.ProviderStore with a shared |
| 116 | // in-memory map. All ephemeral peers reference the same instance |
| 117 | // so any peer can answer provider queries for any CID. |
| 118 | type sharedMemStore struct { |
| 119 | mu sync.RWMutex |
| 120 | data map[string][]peer.AddrInfo |
| 121 | } |
| 122 | |
| 123 | var _ records.ProviderStore = (*sharedMemStore)(nil) |
| 124 | |
| 125 | func (s *sharedMemStore) AddProvider(_ context.Context, key []byte, prov peer.AddrInfo) error { |
| 126 | h := hex.EncodeToString(key) |
| 127 | s.mu.Lock() |
| 128 | defer s.mu.Unlock() |
| 129 | for _, existing := range s.data[h] { |
| 130 | if existing.ID == prov.ID { |
| 131 | return nil |
| 132 | } |
| 133 | } |
| 134 | s.data[h] = append(s.data[h], prov) |
| 135 | return nil |
| 136 | } |
| 137 | |
| 138 | func (s *sharedMemStore) GetProviders(_ context.Context, key []byte) ([]peer.AddrInfo, error) { |
| 139 | h := hex.EncodeToString(key) |
| 140 | s.mu.RLock() |
| 141 | defer s.mu.RUnlock() |
| 142 | return s.data[h], nil |
| 143 | } |
| 144 | |
| 145 | func (s *sharedMemStore) Close() error { return nil } |