Updated functionality for updating Yggdrasil config in iOS demo

2 files changed, 69 insertions, 18 deletions

diff --git a/cmd/dendrite-demo-yggdrasil/main.go b/cmd/dendrite-demo-yggdrasil/main.go
index d3e9e549..9c8c5055 100644
--- a/cmd/dendrite-demo-yggdrasil/main.go
+++ b/cmd/dendrite-demo-yggdrasil/main.go
@@ -54,7 +54,7 @@ func main() {
 	flag.Parse()
 	internal.SetupPprof()
 
-	ygg, err := yggconn.Setup(*instanceName, *instancePeer, ".")
+	ygg, err := yggconn.Setup(*instanceName, *instancePeer, ".", true)
 	if err != nil {
 		panic(err)
 	}
diff --git a/cmd/dendrite-demo-yggdrasil/yggconn/node.go b/cmd/dendrite-demo-yggdrasil/yggconn/node.go
index da184fd5..4b08f1f3 100644
--- a/cmd/dendrite-demo-yggdrasil/yggconn/node.go
+++ b/cmd/dendrite-demo-yggdrasil/yggconn/node.go
@@ -80,7 +80,7 @@ func (n *Node) DialerContext(ctx context.Context, network, address string) (net.
 }
 
 // nolint:gocyclo
-func Setup(instanceName, instancePeer, storageDirectory string) (*Node, error) {
+func Setup(instanceName, instancePeer, storageDirectory string, enableMulticast bool) (*Node, error) {
 	n := &Node{
 		core:      &yggdrasil.Core{},
 		config:    yggdrasilconfig.GenerateConfig(),
@@ -100,22 +100,30 @@ func Setup(instanceName, instancePeer, storageDirectory string) (*Node, error) {
 		if err := json.Unmarshal([]byte(yggconf), &n.config); err != nil {
 			panic(err)
 		}
+	}
+
+	if instancePeer != "" {
+		n.config.Peers = []string{instancePeer}
 	} else {
-		n.config.AdminListen = "none" // fmt.Sprintf("unix://%s/%s-yggdrasil.sock", storageDirectory, instanceName)
+		n.config.Peers = []string{}
+	}
+	n.config.AdminListen = "none"
+	if enableMulticast {
 		n.config.MulticastInterfaces = []string{".*"}
-		n.config.EncryptionPrivateKey = hex.EncodeToString(n.EncryptionPrivateKey())
-		n.config.EncryptionPublicKey = hex.EncodeToString(n.EncryptionPublicKey())
+	} else {
+		n.config.MulticastInterfaces = []string{}
+	}
+	n.config.EncryptionPrivateKey = hex.EncodeToString(n.EncryptionPrivateKey())
+	n.config.EncryptionPublicKey = hex.EncodeToString(n.EncryptionPublicKey())
 
-		j, err := json.MarshalIndent(n.config, "", "  ")
-		if err != nil {
-			panic(err)
-		}
-		if e := ioutil.WriteFile(yggfile, j, 0600); e != nil {
-			n.log.Printf("Couldn't write private key to file '%s': %s\n", yggfile, e)
-		}
+	j, err := json.MarshalIndent(n.config, "", "  ")
+	if err != nil {
+		panic(err)
+	}
+	if e := ioutil.WriteFile(yggfile, j, 0600); e != nil {
+		n.log.Printf("Couldn't write private key to file '%s': %s\n", yggfile, e)
 	}
 
-	var err error
 	n.log.EnableLevel("error")
 	n.log.EnableLevel("warn")
 	n.log.EnableLevel("info")
@@ -123,11 +131,6 @@ func Setup(instanceName, instancePeer, storageDirectory string) (*Node, error) {
 	if err != nil {
 		panic(err)
 	}
-	if instancePeer != "" {
-		if err = n.core.AddPeer(instancePeer, ""); err != nil {
-			panic(err)
-		}
-	}
 	if err = n.multicast.Init(n.core, n.state, n.log, nil); err != nil {
 		panic(err)
 	}
@@ -210,3 +213,51 @@ func (n *Node) KnownNodes() []gomatrixserverlib.ServerName {
 	}
 	return nodes
 }
+
+func (n *Node) SetMulticastEnabled(enabled bool) {
+	if enabled {
+		n.config.MulticastInterfaces = []string{".*"}
+	} else {
+		n.config.MulticastInterfaces = []string{}
+	}
+	n.multicast.UpdateConfig(n.config)
+	if !enabled {
+		n.DisconnectMulticastPeers()
+	}
+}
+
+func (n *Node) DisconnectMulticastPeers() {
+	for _, sp := range n.core.GetSwitchPeers() {
+		if !strings.HasPrefix(sp.Endpoint, "fe80") {
+			continue
+		}
+		if err := n.core.DisconnectPeer(sp.Port); err != nil {
+			n.log.Printf("Failed to disconnect port %d: %s", sp.Port, err)
+		}
+	}
+}
+
+func (n *Node) DisconnectNonMulticastPeers() {
+	for _, sp := range n.core.GetSwitchPeers() {
+		if strings.HasPrefix(sp.Endpoint, "fe80") {
+			continue
+		}
+		if err := n.core.DisconnectPeer(sp.Port); err != nil {
+			n.log.Printf("Failed to disconnect port %d: %s", sp.Port, err)
+		}
+	}
+}
+
+func (n *Node) SetStaticPeer(uri string) error {
+	n.config.Peers = []string{}
+	n.core.UpdateConfig(n.config)
+	n.DisconnectNonMulticastPeers()
+	if uri != "" {
+		n.log.Infoln("Adding static peer", uri)
+		if err := n.core.AddPeer(uri, ""); err != nil {
+			n.log.Infoln("Adding static peer failed:", err)
+			return err
+		}
+	}
+	return nil
+}