cmd/dht-server/main.go

   1 package main
   2
   3 import (
   4         "flag"
   5         "fmt"
   6         "io"
   7         "log"
   8         "os"
   9         "os/signal"
  10
  11         "bitbucket.org/anacrolix/go.torrent/dht"
  12 )
  13
  14 type pingResponse struct {
  15         addr string
  16         krpc dht.Msg
  17 }
  18
  19 var (
  20         tableFileName = flag.String("tableFile", "", "name of file for storing node info")
  21         serveAddr     = flag.String("serveAddr", ":0", "local UDP address")
  22
  23         s *dht.Server
  24 )
  25
  26 func loadTable() error {
  27         if *tableFileName == "" {
  28                 return nil
  29         }
  30         f, err := os.Open(*tableFileName)
  31         if os.IsNotExist(err) {
  32                 return nil
  33         }
  34         if err != nil {
  35                 return fmt.Errorf("error opening table file: %s", err)
  36         }
  37         defer f.Close()
  38         added := 0
  39         for {
  40                 b := make([]byte, dht.CompactNodeInfoLen)
  41                 _, err := io.ReadFull(f, b)
  42                 if err == io.EOF {
  43                         break
  44                 }
  45                 if err != nil {
  46                         return fmt.Errorf("error reading table file: %s", err)
  47                 }
  48                 var ni dht.NodeInfo
  49                 err = ni.UnmarshalCompact(b)
  50                 if err != nil {
  51                         return fmt.Errorf("error unmarshaling compact node info: %s", err)
  52                 }
  53                 s.AddNode(ni)
  54                 added++
  55         }
  56         log.Printf("loaded %d nodes from table file", added)
  57         return nil
  58 }
  59
  60 func init() {
  61         log.SetFlags(log.LstdFlags | log.Lshortfile)
  62         flag.Parse()
  63         var err error
  64         s, err = dht.NewServer(&dht.ServerConfig{
  65                 Addr: *serveAddr,
  66         })
  67         if err != nil {
  68                 log.Fatal(err)
  69         }
  70         err = loadTable()
  71         if err != nil {
  72                 log.Fatalf("error loading table: %s", err)
  73         }
  74         log.Printf("dht server on %s, ID is %q", s.LocalAddr(), s.IDString())
  75         setupSignals()
  76 }
  77
  78 func saveTable() error {
  79         goodNodes := s.Nodes()
  80         if *tableFileName == "" {
  81                 if len(goodNodes) != 0 {
  82                         log.Printf("discarding %d good nodes!", len(goodNodes))
  83                 }
  84                 return nil
  85         }
  86         f, err := os.OpenFile(*tableFileName, os.O_CREATE|os.O_TRUNC|os.O_WRONLY, 0666)
  87         if err != nil {
  88                 return fmt.Errorf("error opening table file: %s", err)
  89         }
  90         defer f.Close()
  91         for _, nodeInfo := range goodNodes {
  92                 var b [dht.CompactNodeInfoLen]byte
  93                 err := nodeInfo.PutCompact(b[:])
  94                 if err != nil {
  95                         return fmt.Errorf("error compacting node info: %s", err)
  96                 }
  97                 _, err = f.Write(b[:])
  98                 if err != nil {
  99                         return fmt.Errorf("error writing compact node info: %s", err)
 100                 }
 101         }
 102         log.Printf("saved %d nodes to table file", len(goodNodes))
 103         return nil
 104 }
 105
 106 func setupSignals() {
 107         ch := make(chan os.Signal)
 108         signal.Notify(ch)
 109         go func() {
 110                 <-ch
 111                 s.Close()
 112         }()
 113 }
 114
 115 func main() {
 116         select {}
 117         // if err := saveTable(); err != nil {
 118         //      log.Printf("error saving node table: %s", err)
 119         // }
 120 }