cmd/dht-ping/main.go

   1 package main
   2
   3 import (
   4         "flag"
   5         "fmt"
   6         "log"
   7         "net"
   8         "os"
   9         "time"
  10
  11         "github.com/anacrolix/torrent/dht"
  12 )
  13
  14 type pingResponse struct {
  15         addr string
  16         krpc dht.Msg
  17         rtt  time.Duration
  18 }
  19
  20 func main() {
  21         log.SetFlags(log.LstdFlags | log.Lshortfile)
  22         timeout := flag.Duration("timeout", -1, "maximum timeout")
  23         flag.Parse()
  24         pingStrAddrs := flag.Args()
  25         if len(pingStrAddrs) == 0 {
  26                 os.Stderr.WriteString("u must specify addrs of nodes to ping e.g. router.bittorrent.com:6881\n")
  27                 os.Exit(2)
  28         }
  29         s, err := dht.NewServer(nil)
  30         if err != nil {
  31                 log.Fatal(err)
  32         }
  33         log.Printf("dht server on %s", s.Addr())
  34         pingResponses := make(chan pingResponse)
  35         timeoutChan := make(chan struct{})
  36         go func() {
  37                 for i, netloc := range pingStrAddrs {
  38                         if i != 0 {
  39                                 time.Sleep(1 * time.Millisecond)
  40                         }
  41                         addr, err := net.ResolveUDPAddr("udp4", netloc)
  42                         if err != nil {
  43                                 log.Fatal(err)
  44                         }
  45                         t, err := s.Ping(addr)
  46                         if err != nil {
  47                                 log.Fatal(err)
  48                         }
  49                         start := time.Now()
  50                         t.SetResponseHandler(func(addr string) func(dht.Msg) {
  51                                 return func(resp dht.Msg) {
  52                                         pingResponses <- pingResponse{
  53                                                 addr: addr,
  54                                                 krpc: resp,
  55                                                 rtt:  time.Now().Sub(start),
  56                                         }
  57                                 }
  58                         }(netloc))
  59                 }
  60                 if *timeout >= 0 {
  61                         time.Sleep(*timeout)
  62                         close(timeoutChan)
  63                 }
  64         }()
  65         responses := 0
  66 pingResponses:
  67         for _ = range pingStrAddrs {
  68                 select {
  69                 case resp := <-pingResponses:
  70                         if resp.krpc == nil {
  71                                 break
  72                         }
  73                         responses++
  74                         fmt.Printf("%-65s %s\n", fmt.Sprintf("%x (%s):", resp.krpc["r"].(map[string]interface{})["id"].(string), resp.addr), resp.rtt)
  75                 case <-timeoutChan:
  76                         break pingResponses
  77                 }
  78         }
  79         // timeouts := len(pingStrAddrs) - responses
  80         fmt.Printf("%d/%d responses (%f%%)\n", responses, len(pingStrAddrs), 100*float64(responses)/float64(len(pingStrAddrs)))
  81 }