mirror of
				https://source.quilibrium.com/quilibrium/ceremonyclient.git
				synced 2025-11-04 00:27:41 +00:00 
			
		
		
		
	
		
			
				
	
	
		
			115 lines
		
	
	
		
			2.3 KiB
		
	
	
	
		
			Go
		
	
	
	
	
	
			
		
		
	
	
			115 lines
		
	
	
		
			2.3 KiB
		
	
	
	
		
			Go
		
	
	
	
	
	
package mocks
 | 
						|
 | 
						|
import (
 | 
						|
	"context"
 | 
						|
	"sync"
 | 
						|
	"time"
 | 
						|
 | 
						|
	"github.com/libp2p/go-libp2p/core/discovery"
 | 
						|
	"github.com/libp2p/go-libp2p/core/host"
 | 
						|
	"github.com/libp2p/go-libp2p/core/peer"
 | 
						|
)
 | 
						|
 | 
						|
type clock interface {
 | 
						|
	Now() time.Time
 | 
						|
}
 | 
						|
 | 
						|
type MockDiscoveryServer struct {
 | 
						|
	mx    sync.Mutex
 | 
						|
	db    map[string]map[peer.ID]*discoveryRegistration
 | 
						|
	clock clock
 | 
						|
}
 | 
						|
 | 
						|
type discoveryRegistration struct {
 | 
						|
	info       peer.AddrInfo
 | 
						|
	expiration time.Time
 | 
						|
}
 | 
						|
 | 
						|
func NewDiscoveryServer(clock clock) *MockDiscoveryServer {
 | 
						|
	return &MockDiscoveryServer{
 | 
						|
		db:    make(map[string]map[peer.ID]*discoveryRegistration),
 | 
						|
		clock: clock,
 | 
						|
	}
 | 
						|
}
 | 
						|
 | 
						|
func (s *MockDiscoveryServer) Advertise(ns string, info peer.AddrInfo, ttl time.Duration) (time.Duration, error) {
 | 
						|
	s.mx.Lock()
 | 
						|
	defer s.mx.Unlock()
 | 
						|
 | 
						|
	peers, ok := s.db[ns]
 | 
						|
	if !ok {
 | 
						|
		peers = make(map[peer.ID]*discoveryRegistration)
 | 
						|
		s.db[ns] = peers
 | 
						|
	}
 | 
						|
	peers[info.ID] = &discoveryRegistration{info, s.clock.Now().Add(ttl)}
 | 
						|
	return ttl, nil
 | 
						|
}
 | 
						|
 | 
						|
func (s *MockDiscoveryServer) FindPeers(ns string, limit int) (<-chan peer.AddrInfo, error) {
 | 
						|
	s.mx.Lock()
 | 
						|
	defer s.mx.Unlock()
 | 
						|
 | 
						|
	peers, ok := s.db[ns]
 | 
						|
	if !ok || len(peers) == 0 {
 | 
						|
		emptyCh := make(chan peer.AddrInfo)
 | 
						|
		close(emptyCh)
 | 
						|
		return emptyCh, nil
 | 
						|
	}
 | 
						|
 | 
						|
	count := len(peers)
 | 
						|
	if limit != 0 && count > limit {
 | 
						|
		count = limit
 | 
						|
	}
 | 
						|
 | 
						|
	iterTime := s.clock.Now()
 | 
						|
	ch := make(chan peer.AddrInfo, count)
 | 
						|
	numSent := 0
 | 
						|
	for p, reg := range peers {
 | 
						|
		if numSent == count {
 | 
						|
			break
 | 
						|
		}
 | 
						|
		if iterTime.After(reg.expiration) {
 | 
						|
			delete(peers, p)
 | 
						|
			continue
 | 
						|
		}
 | 
						|
 | 
						|
		numSent++
 | 
						|
		ch <- reg.info
 | 
						|
	}
 | 
						|
	close(ch)
 | 
						|
 | 
						|
	return ch, nil
 | 
						|
}
 | 
						|
 | 
						|
type MockDiscoveryClient struct {
 | 
						|
	host   host.Host
 | 
						|
	server *MockDiscoveryServer
 | 
						|
}
 | 
						|
 | 
						|
func NewDiscoveryClient(h host.Host, server *MockDiscoveryServer) *MockDiscoveryClient {
 | 
						|
	return &MockDiscoveryClient{
 | 
						|
		host:   h,
 | 
						|
		server: server,
 | 
						|
	}
 | 
						|
}
 | 
						|
 | 
						|
func (d *MockDiscoveryClient) Advertise(ctx context.Context, ns string, opts ...discovery.Option) (time.Duration, error) {
 | 
						|
	var options discovery.Options
 | 
						|
	err := options.Apply(opts...)
 | 
						|
	if err != nil {
 | 
						|
		return 0, err
 | 
						|
	}
 | 
						|
 | 
						|
	return d.server.Advertise(ns, *host.InfoFromHost(d.host), options.Ttl)
 | 
						|
}
 | 
						|
 | 
						|
func (d *MockDiscoveryClient) FindPeers(ctx context.Context, ns string, opts ...discovery.Option) (<-chan peer.AddrInfo, error) {
 | 
						|
	var options discovery.Options
 | 
						|
	err := options.Apply(opts...)
 | 
						|
	if err != nil {
 | 
						|
		return nil, err
 | 
						|
	}
 | 
						|
 | 
						|
	return d.server.FindPeers(ns, options.Limit)
 | 
						|
}
 |