dht: save tokens from get_peers for announce

TODO: implement announce_peer
This commit is contained in:
Hojun-Cho 2026-01-22 10:42:55 +09:00
parent c70d24be5c
commit 60ec9c19ea
2 changed files with 28 additions and 16 deletions

View File

@ -63,7 +63,8 @@ func (dp *dialPool) stop() {
func (dp *dialPool) connectLoop(ctx context.Context) { func (dp *dialPool) connectLoop(ctx context.Context) {
dhtCtx, cancel := context.WithTimeout(ctx, dhtTimeout) dhtCtx, cancel := context.WithTimeout(ctx, dhtTimeout)
addrs, err := dp.dht.GetPeers(dhtCtx, dp.infoHash) // TODO: handle tokens
addrs, _, err := dp.dht.GetPeers(dhtCtx, dp.infoHash)
cancel() cancel()
if err != nil { if err != nil {

View File

@ -160,17 +160,23 @@ func sortByDist(nodes []node, target [20]byte) {
}) })
} }
func (d *DHT) GetPeers(ctx context.Context, h [20]byte) ([]*net.TCPAddr, error) { func (d *DHT) GetPeers(ctx context.Context, h [20]byte) ([]*net.TCPAddr, map[string]string, error) {
queried := make(map[string]bool) type reply struct {
candidates := make([]node, len(d.nodes)) from *net.UDPAddr
copy(candidates, d.nodes) data *resp
sortByDist(candidates, h) }
results := make(chan *resp)
inflight := 0 inflight := 0
queryCount := 0 queryCount := 0
queried := make(map[string]bool)
candidates := make([]node, len(d.nodes))
tokens := make(map[string]string)
replies := make(chan reply)
var peers []*net.TCPAddr var peers []*net.TCPAddr
copy(candidates, d.nodes)
sortByDist(candidates, h)
for queryCount < maxQueries { for queryCount < maxQueries {
if ctx.Err() != nil { if ctx.Err() != nil {
break break
@ -186,23 +192,28 @@ func (d *DHT) GetPeers(ctx context.Context, h [20]byte) ([]*net.TCPAddr, error)
queryCount++ queryCount++
go func(addr *net.UDPAddr) { go func(addr *net.UDPAddr) {
r, _ := d.query(ctx, addr, getPeers, args{InfoHash: h}) r, _ := d.query(ctx, addr, getPeers, args{InfoHash: h})
results <- r replies <- reply{addr, r}
}(n.Addr) }(n.Addr)
} }
if inflight == 0 { if inflight == 0 {
break break
} }
r := <-results r := <-replies
inflight-- inflight--
if r == nil { if r.data == nil || r.data.Token == "" {
continue continue
} }
for _, p := range r.Peers { tokens[r.from.String()] = r.data.Token
for _, p := range r.data.Peers {
if a := decodePeer(p); a != nil { if a := decodePeer(p); a != nil {
if !slices.ContainsFunc(peers, func(p *net.TCPAddr) bool {
return p.String() == a.String()
}) {
peers = append(peers, a) peers = append(peers, a)
} }
} }
for _, c := range r.Nodes { }
for _, c := range r.data.Nodes {
if !queried[c.Addr.String()] { if !queried[c.Addr.String()] {
candidates = append(candidates, c) candidates = append(candidates, c)
} }
@ -210,12 +221,12 @@ func (d *DHT) GetPeers(ctx context.Context, h [20]byte) ([]*net.TCPAddr, error)
sortByDist(candidates, h) sortByDist(candidates, h)
} }
for range inflight { for range inflight {
<-results <-replies
} }
if len(peers) > 0 { if len(peers) > 0 {
return peers, nil return peers, tokens, nil
} }
return nil, fmt.Errorf("no peers found") return nil, nil, fmt.Errorf("no peers found")
} }
func nextCandidate(candidates []node, queried map[string]bool) (node, int) { func nextCandidate(candidates []node, queried map[string]bool) (node, int) {