Wrap server UDP sends to ensure the node last send is updated

diff --git a/dht/dht.go b/dht/dht.go
index c251cbcb0bc7eaa7ad3e15aa2f0a9afbbc128ddd..e7e4fd941feca7a31ba32de4aced006943a48a71 100644 (file)
--- a/dht/dht.go
+++ b/dht/dht.go
@@ -147,7 +147,7 @@ func (s *Server) reply(addr *net.UDPAddr, t string) {
        if err != nil {
                panic(err)
        }
-       _, err = s.Socket.WriteTo(b, addr)
+       err = s.writeToNode(b, addr)
        if err != nil {
                panic(err)
        }
@@ -170,6 +170,19 @@ func (s *Server) getNode(addr *net.UDPAddr) (n *Node) {
        return
 }
 
+func (s *Server) writeToNode(b []byte, node *net.UDPAddr) (err error) {
+       n, err := s.Socket.WriteTo(b, node)
+       if err != nil {
+               return
+       }
+       if n != len(b) {
+               err = io.ErrShortWrite
+               return
+       }
+       s.sentToNode(node)
+       return
+}
+
 func (s *Server) sentToNode(addr *net.UDPAddr) {
        n := s.getNode(addr)
        n.lastSentTo = time.Now()
@@ -237,17 +250,10 @@ func (s *Server) query(node *net.UDPAddr, q string, a map[string]string) (t *tra
        }
        t.response = t.Response
        s.addTransaction(t)
-       n, err := s.Socket.WriteTo(b, node)
+       err = s.writeToNode(b, node)
        if err != nil {
                s.removeTransaction(t)
-               return
-       }
-       if n != len(b) {
-               err = io.ErrShortWrite
-               s.removeTransaction(t)
-               return
        }
-       s.sentToNode(node)
        return
 }