pinecone ping demo

ping.go

package main

import (
	"bytes"
	"context"
	"crypto/ed25519"
	"encoding/hex"
	"encoding/json"
	"flag"
	"fmt"
	"log"
	"net"
	"os"
	"os/signal"
	"strings"
	"syscall"
	"time"

	"github.com/matrix-org/pinecone/connections"
	"github.com/matrix-org/pinecone/multicast"
	"github.com/matrix-org/pinecone/router"
	"github.com/matrix-org/pinecone/types"
)

type Payload struct {
	Seq   int
	Begin time.Time
}

func main() {
	sigs := make(chan os.Signal, 1)
	signal.Notify(sigs, syscall.SIGINT, syscall.SIGTERM)

	listentcp := flag.String("listen", ":0", "address to listen for TCP connections")
	connect := flag.String("connect", "", "peers to connect to")
	count := flag.Int("count", 4, "ping count")
	pubKeyStr := flag.String("pubkey", "", "peer public key")
	flag.Parse()

	_, sk, err := ed25519.GenerateKey(nil)
	if err != nil {
		panic(err)
	}

	logger := log.New(os.Stdout, "", 0)
	pineconeRouter := router.NewRouter(logger, sk)
	pineconeRouter.EnableHopLimiting()
	pineconeRouter.EnableWakeupBroadcasts()
	// pineconeEvents := make(chan events.Event)
	// pineconeRouter.Subscribe(pineconeEvents)

	pineconeMulticast := multicast.NewMulticast(logger, pineconeRouter)
	pineconeMulticast.Start()

	connMgr := connections.NewConnectionManager(pineconeRouter, nil)

	listener := net.ListenConfig{}
	if listentcp != nil && *listentcp != "" {
		go func() {
			listener, err := listener.Listen(context.Background(), "tcp", *listentcp)
			if err != nil {
				panic(err)
			}
			fmt.Printf("listening on: %s", listener.Addr())

			for {
				conn, err := listener.Accept()
				if err != nil {
					panic(err)
				}

				if _, err := pineconeRouter.Connect(
					conn,
					router.ConnectionURI(conn.RemoteAddr().String()),
					router.ConnectionPeerType(router.PeerTypeRemote),
				); err != nil {
					fmt.Println("Inbound TCP connection", conn.RemoteAddr(), "error:", err)
					conn.Close()
				} else {
					fmt.Println("Inbound TCP connection", conn.RemoteAddr(), "is connected")
				}
			}
		}()
	}

	if connect == nil || *connect == "" {
		go func() { // echo msg
			for {
				var msg [types.MaxPayloadSize]byte
				_, remoteAddr, err := pineconeRouter.ReadFrom(msg[:])
				if err != nil {
					fmt.Printf("got msg failed: %s\n", err)
					continue
				}
				content := bytes.TrimRight(msg[:], "\x00")
				fmt.Printf("got msg <%s> from <%s>\n", content, remoteAddr)
				if _, err := pineconeRouter.WriteTo(content, remoteAddr); err != nil {
					panic(err)
				}
			}
		}()
	}

	if connect != nil && *connect != "" {
		for _, item := range strings.Split(*connect, ",") {
			connMgr.AddPeer(strings.TrimSpace(item))
		}
	}

	if pubKeyStr != nil && len(*pubKeyStr) > 0 {
		for {
			if len(pineconeRouter.Peers()) < 2 {
				fmt.Printf("peers: %+v\n", pineconeRouter.Peers())
				time.Sleep(100 * time.Millisecond)
				continue
			}
			break
		}

		pkBytes, err := hex.DecodeString(*pubKeyStr)
		if err != nil {
			log.Fatalf("decode hex public key failed: %s\n", err)
		}
		pk := types.PublicKey{}
		copy(pk[:], pkBytes)

		for i := 0; i < *count; i++ {
			fmt.Printf("peers: %+v\n", pineconeRouter.Peers())
			payload := Payload{
				Seq:   i,
				Begin: time.Now(),
			}
			msg, err := json.Marshal(payload)
			if err != nil {
				panic(err)
			}

			fmt.Printf("send %+v to %s\n", payload, pk)
			if !writeToWithTimeout(pineconeRouter, msg, pk) {
				continue
			}

			// read msg
			var readMsg [types.MaxPayloadSize]byte
			_, addr, err := pineconeRouter.ReadFrom(readMsg[:])
			if err != nil {
				fmt.Printf("read msg failed: %s\n", err)
				continue
			}
			data := bytes.TrimRight(msg[:], "\x00")
			var content Payload
			if err := json.Unmarshal(data, &content); err != nil {
				panic(err)
			}

			fmt.Printf("got msg from %s, seq=%d time=%s\n", addr, content.Seq, time.Now().Sub(content.Begin))
		}
	}

	<-sigs
}

func writeToWithTimeout(_router *router.Router, msg []byte, pk net.Addr) (sent bool) {
	ctx, cancel := context.WithTimeout(context.Background(), time.Second*10)
	defer cancel()

	go func() {
		if _, err := _router.WriteTo(
			msg, pk,
		); err != nil {
			fmt.Printf("write to %s failed: %s\n", pk, err)
			return
		}
		sent = true
	}()

	<-ctx.Done()

	return sent
}

test command:

pubC:

go run ping.go -listen ":2222"

privB:

go run ping.go -listen :3333 -connect <pubC-public-ip>:2222 

privA:

go run ping.go -count 10 -listen :3333 -connect <pubC-public-ip>:2222  -pubkey <privB-router-identity>