dht/announce.go | 2 +- dht/dht.go | 21 +++++++++++++++++---- dht/dht_test.go | 2 +- diff --git a/dht/announce.go b/dht/announce.go index 84b3c7bc958b052327efb622b4254771e63303ac..a7c55868d618a4e881c69350ebd728f260a3be50 100644 --- a/dht/announce.go +++ b/dht/announce.go @@ -178,7 +178,7 @@ } nodeInfo := NodeInfo{ Addr: t.remoteAddr, } - copy(nodeInfo.ID[:], m.ID()) + copy(nodeInfo.ID[:], m.SenderID()) select { case me.values <- PeersValues{ Peers: vs, diff --git a/dht/dht.go b/dht/dht.go index ed45395ca0d57648f7753e3fa999d11a8c4641eb..bca1c7783a41dfc7400690def0e307fd39868908 100644 --- a/dht/dht.go +++ b/dht/dht.go @@ -289,11 +289,24 @@ t, _ = tif.(string) return } -func (m Msg) ID() string { +func (m Msg) Args() map[string]interface{} { + defer func() { + recover() + }() + return m["a"].(map[string]interface{}) +} + +func (m Msg) SenderID() string { defer func() { recover() }() - return m[m["y"].(string)].(map[string]interface{})["id"].(string) + switch m["y"].(string) { + case "q": + return m.Args()["id"].(string) + case "r": + return m["r"].(map[string]interface{})["id"].(string) + } + return "" } // Suggested nodes in a response. @@ -647,7 +660,7 @@ if t == nil { //log.Printf("unexpected message: %#v", d) return } - node := s.getNode(addr, d.ID()) + node := s.getNode(addr, d.SenderID()) node.lastGotResponse = time.Now() // TODO: Update node ID as this is an authoritative packet. go t.handleResponse(d) @@ -1090,7 +1103,7 @@ t, err = s.query(addr, "get_peers", map[string]interface{}{"info_hash": infoHash}, func(m Msg) { s.liftNodes(m) at, ok := m.AnnounceToken() if ok { - s.getNode(addr, m.ID()).announceToken = at + s.getNode(addr, m.SenderID()).announceToken = at } }) return diff --git a/dht/dht_test.go b/dht/dht_test.go index 6636ad6f252a3a702744c564936f1764b505040a..f84ccb68f8a9788518fe07acc6b26ffe9ca3ac8c 100644 --- a/dht/dht_test.go +++ b/dht/dht_test.go @@ -145,7 +145,7 @@ } defer tn.Close() ok := make(chan bool) tn.SetResponseHandler(func(msg Msg) { - ok <- msg.ID() == srv0.ID() + ok <- msg.SenderID() == srv0.ID() }) if !<-ok { t.FailNow()