dht: Improve on on Msg methods

diff --git a/dht/announce.go b/dht/announce.go
index 84b3c7bc958b052327efb622b4254771e63303ac..a7c55868d618a4e881c69350ebd728f260a3be50 100644 (file)
--- a/dht/announce.go
+++ b/dht/announce.go
@@ -178,7 +178,7 @@ func (me *Announce) getPeers(addr dHTAddr) error {
                        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 (file)
--- a/dht/dht.go
+++ b/dht/dht.go
@@ -289,11 +289,24 @@ func (m Msg) T() (t string) {
        return
 }
 
-func (m Msg) ID() string {
+func (m Msg) Args() map[string]interface{} {
        defer func() {
                recover()
        }()
-       return m[m["y"].(string)].(map[string]interface{})["id"].(string)
+       return m["a"].(map[string]interface{})
+}
+
+func (m Msg) SenderID() string {
+       defer func() {
+               recover()
+       }()
+       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 @@ func (s *Server) processPacket(b []byte, addr dHTAddr) {
                //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 @@ func (s *Server) getPeers(addr dHTAddr, infoHash string) (t *Transaction, err er
                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 (file)
--- a/dht/dht_test.go
+++ b/dht/dht_test.go
@@ -145,7 +145,7 @@ func TestPing(t *testing.T) {
        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()