Update P2P demos

1 files changed, 47 insertions, 12 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{}