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