Fix token for response

1 files changed, 60 insertions, 50 deletions

diff --git a/dht/messages.go b/dht/messages.go
index be7d6b6..023e27d 100644
--- a/dht/messages.go
+++ b/dht/messages.go
@@ -3,50 +3,58 @@ package dht
 import (
 	"fmt"
 	"net"
-	"strings"
+	//"strings"
 )
 
-func (n *Node) onPingQuery(rn remoteNode, msg map[string]interface{}) {
-	t := msg["t"].(string)
-	//n.log.Debug("ping", "source", rn)
+func (n *Node) onPingQuery(rn remoteNode, msg map[string]interface{}) error {
+	t, err := getStringKey(msg, "t")
+	if err != nil {
+		return err
+	}
 	n.queueMsg(rn, makeResponse(t, map[string]interface{}{
 		"id": string(n.id),
 	}))
+	return nil
 }
 
-func (n *Node) onGetPeersQuery(rn remoteNode, msg map[string]interface{}) {
-	a := msg["a"].(map[string]interface{})
-	if err := checkKey(a, "info_hash", "string"); err != nil {
-		//n.queueMsg(addr, makeError(t, protocolError, err.Error()))
-		return
+func (n *Node) onGetPeersQuery(rn remoteNode, msg map[string]interface{}) error {
+	a, err := getMapKey(msg, "a")
+	if err != nil {
+		return err
 	}
 
 	// This is the ih of the torrent
-	th, err := InfohashFromString(a["info_hash"].(string))
+	torrent, err := getStringKey(a, "info_hash")
 	if err != nil {
-		n.log.Warn("invalid torrent", "infohash", a["info_hash"])
+		return err
+	}
+	th, err := InfohashFromString(torrent)
+	if err != nil {
+		return err
 	}
 	n.log.Debug("get_peers query", "source", rn, "torrent", th)
 
-	token := []byte(*th)[:2]
-
-	id := generateNeighbour(n.id, *th)
-	nodes := n.rTable.get(8)
-	compactNS := []string{}
-	for _, rn := range nodes {
-		ns := encodeCompactNodeAddr(rn.address.String())
-		if ns == "" {
-			n.log.Warn("failed to compact node", "address", rn.address.String())
-			continue
+	token := torrent[:2]
+	neighbour := generateNeighbour(n.id, *th)
+	/*
+		nodes := n.rTable.get(8)
+		compactNS := []string{}
+		for _, rn := range nodes {
+			ns := encodeCompactNodeAddr(rn.address.String())
+			if ns == "" {
+				n.log.Warn("failed to compact node", "address", rn.address.String())
+				continue
+			}
+			compactNS = append(compactNS, ns)
 		}
-		compactNS = append(compactNS, ns)
-	}
+	*/
 
 	t := msg["t"].(string)
 	n.queueMsg(rn, makeResponse(t, map[string]interface{}{
-		"id":    string(id),
+		"id":    string(neighbour),
 		"token": token,
-		"nodes": strings.Join(compactNS, ""),
+		"nodes": "",
+		//"nodes": strings.Join(compactNS, ""),
 	}))
 
 	//nodes := n.rTable.get(50)
@@ -60,46 +68,47 @@ func (n *Node) onGetPeersQuery(rn remoteNode, msg map[string]interface{}) {
 			n.queueMsg(*o, q)
 		}
 	*/
+	return nil
 }
 
-func (n *Node) onAnnouncePeerQuery(rn remoteNode, msg map[string]interface{}) {
-	a := msg["a"].(map[string]interface{})
-	err := checkKeys(a, [][]string{
+func (n *Node) onAnnouncePeerQuery(rn remoteNode, msg map[string]interface{}) error {
+	a, err := getMapKey(msg, "a")
+	if err != nil {
+		return err
+	}
+	err = checkKeys(a, [][]string{
 		{"info_hash", "string"},
 		{"port", "int"},
 		{"token", "string"},
 	})
-	if err != nil {
-		//n.queueMsg(addr, makeError(t, protocolError, err.Error()))
-		return
-	}
 
 	n.log.Debug("announce_peer", "source", rn)
 
-	// TODO
-	if impliedPort, ok := a["implied_port"]; ok && impliedPort.(int) != 0 {
-		// Use the port from the network
-	} else {
-		// Use the port in the message
-		host, _, err := net.SplitHostPort(rn.address.String())
-		if err != nil {
-			n.log.Warn("failed to split host/port", "error", err)
-			return
-		}
-		newPort := a["port"]
-		if newPort == 0 {
-			n.log.Warn("sent port 0", "source", rn)
-			return
+	if impliedPort, err := getIntKey(a, "implied_port"); err == nil {
+		if impliedPort != 0 {
+			// Use the port in the message
+			host, _, err := net.SplitHostPort(rn.address.String())
+			if err != nil {
+				return err
+			}
+			newPort := a["port"]
+			if newPort == 0 {
+				return fmt.Errorf("ignoring port 0")
+			}
+			addr, err := net.ResolveUDPAddr(n.family, fmt.Sprintf("%s:%d", host, newPort))
+			rn = remoteNode{address: addr, id: rn.id}
 		}
-		addr, err := net.ResolveUDPAddr(n.family, fmt.Sprintf("%s:%d", host, newPort))
-		rn = remoteNode{address: addr, id: rn.id}
 	}
 
 	// TODO do we reply?
 
-	ih, err := InfohashFromString(a["info_hash"].(string))
+	ihStr, err := getStringKey(a, "info_hash")
+	if err != nil {
+		return err
+	}
+	ih, err := InfohashFromString(ihStr)
 	if err != nil {
-		n.log.Warn("invalid torrent", "infohash", a["info_hash"])
+		n.log.Warn("invalid torrent", "infohash", ihStr)
 	}
 
 	p := Peer{Node: rn, Infohash: *ih}
@@ -107,6 +116,7 @@ func (n *Node) onAnnouncePeerQuery(rn remoteNode, msg map[string]interface{}) {
 	if n.OnAnnouncePeer != nil {
 		go n.OnAnnouncePeer(p)
 	}
+	return nil
 }
 
 func (n *Node) onFindNodeResponse(rn remoteNode, msg map[string]interface{}) {