Update P2P demos

2 files changed, 65 insertions, 42 deletions

diff --git a/cmd/dendrite-demo-yggdrasil/main.go b/cmd/dendrite-demo-yggdrasil/main.go
index d2fddf8b..7445b6bf 100644
--- a/cmd/dendrite-demo-yggdrasil/main.go
+++ b/cmd/dendrite-demo-yggdrasil/main.go
@@ -16,7 +16,9 @@ package main
 
 import (
 	"context"
+	"crypto/ed25519"
 	"crypto/tls"
+	"encoding/hex"
 	"flag"
 	"fmt"
 	"net"
@@ -42,6 +44,7 @@ import (
 	"github.com/matrix-org/dendrite/setup/base"
 	"github.com/matrix-org/dendrite/setup/config"
 	"github.com/matrix-org/dendrite/setup/mscs"
+	"github.com/matrix-org/dendrite/test"
 	"github.com/matrix-org/dendrite/userapi"
 	"github.com/sirupsen/logrus"
 
@@ -49,19 +52,18 @@ import (
 )
 
 var (
-	instanceName = flag.String("name", "dendrite-p2p-ygg", "the name of this P2P demo instance")
-	instancePort = flag.Int("port", 8008, "the port that the client API will listen on")
-	instancePeer = flag.String("peer", "", "an internet Yggdrasil peer to connect to")
+	instanceName   = flag.String("name", "dendrite-p2p-ygg", "the name of this P2P demo instance")
+	instancePort   = flag.Int("port", 8008, "the port that the client API will listen on")
+	instancePeer   = flag.String("peer", "", "the static Yggdrasil peers to connect to, comma separated-list")
+	instanceListen = flag.String("listen", "tcp://:0", "the port Yggdrasil peers can connect to")
 )
 
 func main() {
 	flag.Parse()
 	internal.SetupPprof()
 
-	ygg, err := yggconn.Setup(*instanceName, ".", *instancePeer)
-	if err != nil {
-		panic(err)
-	}
+	var pk ed25519.PublicKey
+	var sk ed25519.PrivateKey
 
 	// iterate through the cli args and check if the config flag was set
 	configFlagSet := false
@@ -74,6 +76,34 @@ func main() {
 
 	cfg := &config.Dendrite{}
 
+	keyfile := *instanceName + ".pem"
+	if _, err := os.Stat(keyfile); os.IsNotExist(err) {
+		oldkeyfile := *instanceName + ".key"
+		if _, err = os.Stat(oldkeyfile); os.IsNotExist(err) {
+			if err = test.NewMatrixKey(keyfile); err != nil {
+				panic("failed to generate a new PEM key: " + err.Error())
+			}
+			if _, sk, err = config.LoadMatrixKey(keyfile, os.ReadFile); err != nil {
+				panic("failed to load PEM key: " + err.Error())
+			}
+		} else {
+			if sk, err = os.ReadFile(oldkeyfile); err != nil {
+				panic("failed to read the old private key: " + err.Error())
+			}
+			if len(sk) != ed25519.PrivateKeySize {
+				panic("the private key is not long enough")
+			}
+			if err := test.SaveMatrixKey(keyfile, sk); err != nil {
+				panic("failed to convert the private key to PEM format: " + err.Error())
+			}
+		}
+	} else {
+		var err error
+		if _, sk, err = config.LoadMatrixKey(keyfile, os.ReadFile); err != nil {
+			panic("failed to load PEM key: " + err.Error())
+		}
+	}
+
 	// use custom config if config flag is set
 	if configFlagSet {
 		cfg = setup.ParseFlags(true)
@@ -82,6 +112,7 @@ func main() {
 			Generate:   true,
 			Monolithic: true,
 		})
+		cfg.Global.PrivateKey = sk
 		cfg.Global.JetStream.StoragePath = config.Path(fmt.Sprintf("%s/", *instanceName))
 		cfg.UserAPI.AccountDatabase.ConnectionString = config.DataSource(fmt.Sprintf("file:%s-account.db", *instanceName))
 		cfg.MediaAPI.Database.ConnectionString = config.DataSource(fmt.Sprintf("file:%s-mediaapi.db", *instanceName))
@@ -93,19 +124,23 @@ func main() {
 		cfg.MSCs.Database.ConnectionString = config.DataSource(fmt.Sprintf("file:%s-mscs.db", *instanceName))
 		cfg.ClientAPI.RegistrationDisabled = false
 		cfg.ClientAPI.OpenRegistrationWithoutVerificationEnabled = true
-		if err = cfg.Derive(); err != nil {
+		if err := cfg.Derive(); err != nil {
 			panic(err)
 		}
 	}
 
-	// always override ServerName, PrivateKey and KeyID
-	cfg.Global.ServerName = gomatrixserverlib.ServerName(ygg.DerivedServerName())
-	cfg.Global.PrivateKey = ygg.PrivateKey()
-	cfg.Global.KeyID = signing.KeyID
+	pk = sk.Public().(ed25519.PublicKey)
+	cfg.Global.ServerName = gomatrixserverlib.ServerName(hex.EncodeToString(pk))
+	cfg.Global.KeyID = gomatrixserverlib.KeyID(signing.KeyID)
 
 	base := base.NewBaseDendrite(cfg, "Monolith")
 	defer base.Close() // nolint: errcheck
 
+	ygg, err := yggconn.Setup(sk, *instanceName, ".", *instancePeer, *instanceListen)
+	if err != nil {
+		panic(err)
+	}
+
 	federation := ygg.CreateFederationClient(base)
 
 	serverKeyAPI := &signing.YggdrasilKeys{}
diff --git a/cmd/dendrite-demo-yggdrasil/yggconn/node.go b/cmd/dendrite-demo-yggdrasil/yggconn/node.go
index ff3c73ec..83b4cdf9 100644
--- a/cmd/dendrite-demo-yggdrasil/yggconn/node.go
+++ b/cmd/dendrite-demo-yggdrasil/yggconn/node.go
@@ -18,15 +18,13 @@ import (
 	"context"
 	"crypto/ed25519"
 	"encoding/hex"
-	"encoding/json"
 	"fmt"
-	"log"
 	"net"
-	"os"
 	"strings"
 
 	"github.com/matrix-org/gomatrixserverlib"
 	"github.com/neilalexander/utp"
+	"github.com/sirupsen/logrus"
 
 	ironwoodtypes "github.com/Arceliar/ironwood/types"
 	yggdrasilconfig "github.com/yggdrasil-network/yggdrasil-go/src/config"
@@ -57,48 +55,38 @@ func (n *Node) DialerContext(ctx context.Context, _, address string) (net.Conn,
 	return n.utpSocket.DialAddrContext(ctx, pk)
 }
 
-func Setup(instanceName, storageDirectory, peerURI string) (*Node, error) {
+func Setup(sk ed25519.PrivateKey, instanceName, storageDirectory, peerURI, listenURI string) (*Node, error) {
 	n := &Node{
 		core:      &yggdrasilcore.Core{},
 		config:    yggdrasildefaults.GenerateConfig(),
 		multicast: &yggdrasilmulticast.Multicast{},
-		log:       gologme.New(os.Stdout, "YGG ", log.Flags()),
+		log:       gologme.New(logrus.StandardLogger().Writer(), "", 0),
 		incoming:  make(chan net.Conn),
 	}
 
-	yggfile := fmt.Sprintf("%s/%s-yggdrasil.conf", storageDirectory, instanceName)
-	if _, err := os.Stat(yggfile); !os.IsNotExist(err) {
-		yggconf, e := os.ReadFile(yggfile)
-		if e != nil {
-			panic(err)
-		}
-		if err := json.Unmarshal([]byte(yggconf), &n.config); err != nil {
-			panic(err)
-		}
+	options := []yggdrasilcore.SetupOption{
+		yggdrasilcore.AdminListenAddress("none"),
+	}
+	if listenURI != "" {
+		options = append(options, yggdrasilcore.ListenAddress(listenURI))
 	}
-
-	n.config.Peers = []string{}
 	if peerURI != "" {
-		n.config.Peers = append(n.config.Peers, peerURI)
+		for _, uri := range strings.Split(peerURI, ",") {
+			options = append(options, yggdrasilcore.Peer{
+				URI: uri,
+			})
+		}
 	}
-	n.config.AdminListen = "none"
 
-	j, err := json.MarshalIndent(n.config, "", "  ")
-	if err != nil {
+	var err error
+	if n.core, err = yggdrasilcore.New(sk, options...); err != nil {
 		panic(err)
 	}
-	if e := os.WriteFile(yggfile, j, 0600); e != nil {
-		n.log.Printf("Couldn't write private key to file '%s': %s\n", yggfile, e)
-	}
-
 	n.log.EnableLevel("error")
 	n.log.EnableLevel("warn")
 	n.log.EnableLevel("info")
-	if err = n.core.Start(n.config, n.log); err != nil {
-		panic(err)
-	}
-	n.utpSocket, err = utp.NewSocketFromPacketConnNoClose(n.core)
-	if err != nil {
+	n.core.SetLogger(n.log)
+	if n.utpSocket, err = utp.NewSocketFromPacketConnNoClose(n.core); err != nil {
 		panic(err)
 	}
 	if err = n.multicast.Init(n.core, n.config, n.log, nil); err != nil {
@@ -108,7 +96,7 @@ func Setup(instanceName, storageDirectory, peerURI string) (*Node, error) {
 		panic(err)
 	}
 
-	n.log.Println("Public key:", n.core.PublicKey())
+	n.log.Printf("Public key: %x", n.core.PublicKey())
 	go n.listenFromYgg()
 
 	return n, nil