reflector.go/cmd/dht.go

90 lines
2.3 KiB
Go
Raw Normal View History

package cmd
import (
2018-07-13 19:31:54 +02:00
"log"
"math/big"
"net"
2018-07-13 19:31:54 +02:00
"net/http"
"os"
"os/signal"
"strconv"
"syscall"
"time"
2018-07-13 19:31:54 +02:00
"github.com/lbryio/reflector.go/dht"
2018-06-14 17:48:02 +02:00
"github.com/lbryio/reflector.go/dht/bits"
2018-07-13 19:31:54 +02:00
"github.com/spf13/cobra"
)
type NodeRPC string
type PingArgs struct {
2018-07-13 19:31:54 +02:00
nodeID string
address string
2018-07-13 19:31:54 +02:00
port int
}
type PingResult string
func (n *NodeRPC) Ping(r *http.Request, args *PingArgs, result *PingResult) error {
*result = PingResult("pong")
return nil
}
var dhtPort int
var rpcPort int
func init() {
var cmd = &cobra.Command{
2018-07-10 23:30:47 +02:00
Use: "dht [bootstrap|connect]",
Short: "Run dht node",
ValidArgs: []string{"start", "bootstrap"},
Args: argFuncs(cobra.ExactArgs(1), cobra.OnlyValidArgs),
Run: dhtCmd,
}
cmd.PersistentFlags().StringP("nodeID", "n", "", "nodeID in hex")
cmd.PersistentFlags().IntVar(&dhtPort, "port", 4567, "Port to start DHT on")
cmd.PersistentFlags().IntVar(&rpcPort, "rpc_port", 1234, "Port to listen for rpc commands on")
rootCmd.AddCommand(cmd)
}
func dhtCmd(cmd *cobra.Command, args []string) {
if args[0] == "bootstrap" {
node := dht.NewBootstrapNode(bits.Rand(), 1*time.Millisecond, 1*time.Minute)
listener, err := net.ListenPacket(dht.Network, "127.0.0.1:"+strconv.Itoa(dhtPort))
checkErr(err)
conn := listener.(*net.UDPConn)
err = node.Connect(conn)
checkErr(err)
interruptChan := make(chan os.Signal, 1)
signal.Notify(interruptChan, os.Interrupt, syscall.SIGTERM)
<-interruptChan
log.Printf("shutting down bootstrap node")
node.Shutdown()
} else {
nodeIDStr := cmd.Flag("nodeID").Value.String()
nodeID := bits.Bitmap{}
if nodeIDStr == "" {
nodeID = bits.Rand()
} else {
nodeID = bits.FromHexP(nodeIDStr)
}
log.Println(nodeID.String())
node := dht.NewBootstrapNode(nodeID, 1*time.Millisecond, 1*time.Minute)
listener, err := net.ListenPacket(dht.Network, "127.0.0.1:"+strconv.Itoa(dhtPort))
checkErr(err)
conn := listener.(*net.UDPConn)
err = node.Connect(conn)
checkErr(err)
log.Println("started node")
_, _, err = dht.FindContacts(&node.Node, nodeID.Sub(bits.FromBigP(big.NewInt(1))), false, nil)
rpcServer := dht.RunRPCServer("127.0.0.1:"+strconv.Itoa(rpcPort), "/", node)
interruptChan := make(chan os.Signal, 1)
signal.Notify(interruptChan, os.Interrupt, syscall.SIGTERM, syscall.SIGINT)
<-interruptChan
rpcServer.Wg.Done()
node.Shutdown()
2018-04-25 00:12:17 +02:00
}
}