main.go 1.71 KB
Newer Older
1 2 3
package main

import (
Matt Joiner's avatar
Matt Joiner committed
4
	"context"
5
	"log"
6
	"net"
Matt Joiner's avatar
Matt Joiner committed
7
	"net/http"
8 9 10
	"os"
	"os/signal"

Matt Joiner's avatar
Matt Joiner committed
11
	_ "github.com/anacrolix/envpprof"
Matt Joiner's avatar
Matt Joiner committed
12
	"github.com/anacrolix/tagflag"
Matt Joiner's avatar
Matt Joiner committed
13 14

	"github.com/anacrolix/dht/v2"
15 16 17
)

var (
Matt Joiner's avatar
Matt Joiner committed
18
	flags = struct {
Matt Joiner's avatar
Matt Joiner committed
19 20 21
		TableFile   string `help:"name of file for storing node info"`
		Addr        string `help:"local UDP address"`
		NoBootstrap bool
Matt Joiner's avatar
Matt Joiner committed
22 23 24
	}{
		Addr: ":0",
	}
25 26 27
	s *dht.Server
)

28 29
func loadTable() (err error) {
	added, err := s.AddNodesFromFile(flags.TableFile)
30
	log.Printf("loaded %d nodes from table file", added)
31
	return
32 33 34
}

func saveTable() error {
35
	return dht.WriteNodesToFile(s.Nodes(), flags.TableFile)
36 37 38
}

func main() {
Matt Joiner's avatar
Matt Joiner committed
39 40
	log.SetFlags(log.LstdFlags | log.Lshortfile)
	tagflag.Parse(&flags)
41 42 43 44 45
	conn, err := net.ListenPacket("udp", flags.Addr)
	if err != nil {
		log.Fatal(err)
	}
	defer conn.Close()
Matt Joiner's avatar
Matt Joiner committed
46
	s, err = dht.NewServer(&dht.ServerConfig{
47 48
		Conn:          conn,
		StartingNodes: dht.GlobalBootstrapAddrs,
Matt Joiner's avatar
Matt Joiner committed
49 50 51 52
	})
	if err != nil {
		log.Fatal(err)
	}
Matt Joiner's avatar
Matt Joiner committed
53 54 55
	http.HandleFunc("/debug/dht", func(w http.ResponseWriter, r *http.Request) {
		s.WriteStatus(w)
	})
56 57 58 59 60
	if flags.TableFile != "" {
		err = loadTable()
		if err != nil {
			log.Fatalf("error loading table: %s", err)
		}
Matt Joiner's avatar
Matt Joiner committed
61 62
	}
	log.Printf("dht server on %s, ID is %x", s.Addr(), s.ID())
63

Matt Joiner's avatar
Matt Joiner committed
64
	ctx, cancel := context.WithCancel(context.Background())
65
	go func() {
Matt Joiner's avatar
Matt Joiner committed
66
		ch := make(chan os.Signal, 1)
Matt Joiner's avatar
Matt Joiner committed
67
		signal.Notify(ch)
68
		<-ch
Matt Joiner's avatar
Matt Joiner committed
69
		cancel()
70
	}()
Matt Joiner's avatar
Matt Joiner committed
71 72 73 74 75 76 77 78 79
	if !flags.NoBootstrap {
		go func() {
			if tried, err := s.Bootstrap(); err != nil {
				log.Printf("error bootstrapping: %s", err)
			} else {
				log.Printf("finished bootstrapping: crawled %#v addrs", tried)
			}
		}()
	}
Matt Joiner's avatar
Matt Joiner committed
80 81
	<-ctx.Done()
	s.Close()
82 83 84 85 86 87

	if flags.TableFile != "" {
		if err := saveTable(); err != nil {
			log.Printf("error saving node table: %s", err)
		}
	}
88
}