core-chain/cmd/p2p/main.go

package main

import (
    "bytes"
    "encoding/base64"
    "encoding/hex"
    "fmt"
    "github.com/ethereum/go-ethereum/crypto"
    "github.com/ethereum/go-ethereum/p2p/discover"
    "github.com/ethereum/go-ethereum/p2p/enode"
    "github.com/ethereum/go-ethereum/p2p/enr"
    "github.com/ethereum/go-ethereum/params"
    "github.com/ethereum/go-ethereum/rlp"
    "gopkg.in/urfave/cli.v1"
    "net"
    "os"
    "path/filepath"
    "sort"
    "strings"
)

var app = &cli.App{
	Name:   filepath.Base(os.Args[0]),
	Usage:  "p2p tool",
	Writer: os.Stdout,
}

var (
    bootnodesFlag = cli.StringFlag{
        Name:  "bootnodes",
        Usage: "Comma separated nodes used for bootstrapping",
    }
    nodekeyFlag = cli.StringFlag{
        Name:  "nodekey",
        Usage: "Hex-encoded node key",
    }
    nodedbFlag = cli.StringFlag{
        Name:  "nodedb",
        Usage: "Nodes database location",
    }
    listenAddrFlag = cli.StringFlag{
        Name:  "addr",
        Usage: "Listening address",
    }
)

func main() {
	app.CommandNotFound = func(ctx *cli.Context, cmd string) {
		fmt.Fprintf(os.Stderr, "No such command: %s\n", cmd)
		os.Exit(1)
	}

	app.Commands = []cli.Command{
		onePeerCommand,
	}

	exit(app.Run(os.Args))
}

func exit(err interface{}) {
	if err == nil {
		os.Exit(0)
	}
	fmt.Fprintln(os.Stderr, err)
	os.Exit(1)
}

// commandHasFlag returns true if the current command supports the given flag.
func commandHasFlag(ctx *cli.Context, flag cli.Flag) bool {
    flags := ctx.FlagNames()
    sort.Strings(flags)
    i := sort.SearchStrings(flags, flag.GetName())
    return i != len(flags) && flags[i] == flag.GetName()
}

// startV4 starts an ephemeral discovery V4 node.
func startV4(ctx *cli.Context) *discover.UDPv4 {
    ln, config := makeDiscoveryConfig(ctx)
    socket := listen(ln, ctx.String(listenAddrFlag.Name))
    disc, err := discover.ListenV4(socket, ln, config)
    if err != nil {
        exit(err)
    }
    return disc
}

func makeDiscoveryConfig(ctx *cli.Context) (*enode.LocalNode, discover.Config) {
    var cfg discover.Config

    if ctx.IsSet(nodekeyFlag.Name) {
        key, err := crypto.HexToECDSA(ctx.String(nodekeyFlag.Name))
        if err != nil {
            exit(fmt.Errorf("-%s: %v", nodekeyFlag.Name, err))
        }
        cfg.PrivateKey = key
    } else {
        cfg.PrivateKey, _ = crypto.GenerateKey()
    }

    if commandHasFlag(ctx, bootnodesFlag) {
        bn, err := parseBootnodes(ctx)
        if err != nil {
            exit(err)
        }
        cfg.Bootnodes = bn
    }

    dbpath := ctx.String(nodedbFlag.Name)
    db, err := enode.OpenDB(dbpath)
    if err != nil {
        exit(err)
    }
    ln := enode.NewLocalNode(db, cfg.PrivateKey)
    return ln, cfg
}

func listen(ln *enode.LocalNode, addr string) *net.UDPConn {
    if addr == "" {
        addr = "0.0.0.0:0"
    }
    socket, err := net.ListenPacket("udp4", addr)
    if err != nil {
        exit(err)
    }
    usocket := socket.(*net.UDPConn)
    uaddr := socket.LocalAddr().(*net.UDPAddr)
    if uaddr.IP.IsUnspecified() {
        ln.SetFallbackIP(net.IP{127, 0, 0, 1})
    } else {
        ln.SetFallbackIP(uaddr.IP)
    }
    ln.SetFallbackUDP(uaddr.Port)
    return usocket
}

func parseBootnodes(ctx *cli.Context) ([]*enode.Node, error) {
    s := params.MainnetBootnodes
    if ctx.IsSet(bootnodesFlag.Name) {
        input := ctx.String(bootnodesFlag.Name)
        if input == "" {
            return nil, nil
        }
        s = strings.Split(input, ",")
    }
    nodes := make([]*enode.Node, len(s))
    var err error
    for i, record := range s {
        nodes[i], err = parseNode(record)
        if err != nil {
            return nil, fmt.Errorf("invalid bootstrap node: %v", err)
        }
    }
    return nodes, nil
}

// parseNode parses a node record and verifies its signature.
func parseNode(source string) (*enode.Node, error) {
    if strings.HasPrefix(source, "enode://") {
        return enode.ParseV4(source)
    }
    r, err := parseRecord(source)
    if err != nil {
        return nil, err
    }
    return enode.New(enode.ValidSchemes, r)
}

// parseRecord parses a node record from hex, base64, or raw binary input.
func parseRecord(source string) (*enr.Record, error) {
    bin := []byte(source)
    if d, ok := decodeRecordHex(bytes.TrimSpace(bin)); ok {
        bin = d
    } else if d, ok := decodeRecordBase64(bytes.TrimSpace(bin)); ok {
        bin = d
    }
    var r enr.Record
    err := rlp.DecodeBytes(bin, &r)
    return &r, err
}

func decodeRecordHex(b []byte) ([]byte, bool) {
    if bytes.HasPrefix(b, []byte("0x")) {
        b = b[2:]
    }
    dec := make([]byte, hex.DecodedLen(len(b)))
    _, err := hex.Decode(dec, b)
    return dec, err == nil
}

func decodeRecordBase64(b []byte) ([]byte, bool) {
    if bytes.HasPrefix(b, []byte("enr:")) {
        b = b[4:]
    }
    dec := make([]byte, base64.RawURLEncoding.DecodedLen(len(b)))
    n, err := base64.RawURLEncoding.Decode(dec, b)
    return dec[:n], err == nil
}

// getNodeArg handles the common case of a single node descriptor argument.
func getNodeArg(ctx *cli.Context) *enode.Node {
    if ctx.NArg() < 1 {
        exit("missing node as command-line argument")
    }
    n, err := parseNode(ctx.Args()[0])
    if err != nil {
        exit(err)
    }
    return n
}