From 68719cabd86e9d89d19bdfa471878571404a255a Mon Sep 17 00:00:00 2001
From: Raphael Reyna <raphaelreyna@protonmail.com>
Date: Sun, 19 May 2024 12:55:16 -0700
Subject: [PATCH] Fix discovery server config bugs (#40)

* fix discovery server bugs

* fix bug where oneshot couldnt connect to discovery server if using ipv6 based redirect
---
 v2/integration_testing/testing.go             |  8 +++-
 v2/pkg/commands/config/get/cobra.go           | 39 ++++++++++++++++++-
 v2/pkg/commands/discovery-server/cobra.go     |  2 +-
 v2/pkg/commands/p2p/client/receive/cobra.go   |  4 +-
 v2/pkg/commands/p2p/client/send/cobra.go      |  4 +-
 v2/pkg/commands/root/discoveryServer.go       |  9 +++--
 v2/pkg/commands/root/listenWebRTC.go          |  2 +-
 v2/pkg/commands/root/runServer.go             | 22 ++++++++++-
 v2/pkg/configuration/configuration.go         |  2 +-
 v2/pkg/configuration/discovery.go             |  4 +-
 v2/pkg/file/tar.go                            | 14 +++++--
 v2/pkg/net/network.go                         | 10 ++---
 .../signallingserver/discoveryServer.go       | 15 ++++---
 13 files changed, 103 insertions(+), 32 deletions(-)

diff --git a/v2/integration_testing/testing.go b/v2/integration_testing/testing.go
index 229a0c6..fc6e083 100644
--- a/v2/integration_testing/testing.go
+++ b/v2/integration_testing/testing.go
@@ -7,6 +7,7 @@ import (
 	"os"
 	"os/exec"
 	"path/filepath"
+	"runtime"
 	"strings"
 	"time"
 
@@ -144,8 +145,11 @@ func (suite *TestSuite) SetupSuite() {
 	suite.Require().NoError(err)
 
 	newCmdPath := filepath.Join(suite.TestDir, "oneshot.testing")
-	cpOPut, err := exec.Command("cp", cmdPath, newCmdPath).CombinedOutput()
-	suite.Require().NoError(err, "failed to copy oneshot binary: %s", string(cpOPut))
+	if runtime.GOOS == "darwin" {
+		newCmdPath = filepath.Join(filepath.Dir(suite.TestDir), "oneshot.testing")
+	}
+	cpOut, err := exec.Command("cp", cmdPath, newCmdPath).CombinedOutput()
+	suite.Require().NoError(err, "failed to copy oneshot binary: %s", string(cpOut))
 
 	err = os.Chdir(suite.TestDir)
 	suite.Require().NoError(err)
diff --git a/v2/pkg/commands/config/get/cobra.go b/v2/pkg/commands/config/get/cobra.go
index 05e417a..e73b26c 100644
--- a/v2/pkg/commands/config/get/cobra.go
+++ b/v2/pkg/commands/config/get/cobra.go
@@ -1,6 +1,7 @@
 package get
 
 import (
+	"fmt"
 	"os"
 
 	"github.com/forestnode-io/oneshot/v2/pkg/configuration"
@@ -31,7 +32,7 @@ func (c *Cmd) Cobra() *cobra.Command {
 		Short: "Get an individual value from a oneshot configuration file.",
 		Long:  "Get an individual value from a oneshot configuration file.",
 		RunE:  c.run,
-		Args:  cobra.ExactArgs(1),
+		Args:  cobra.MaximumNArgs(1),
 	}
 
 	c.cobraCommand.SetUsageTemplate(usageTemplate)
@@ -40,7 +41,41 @@ func (c *Cmd) Cobra() *cobra.Command {
 }
 
 func (c *Cmd) run(cmd *cobra.Command, args []string) error {
-	configData := viper.Get(args[0])
+	m := map[string]any{}
+	err := viper.Unmarshal(&m)
+	if err != nil {
+		return fmt.Errorf("failed to unmarshal configuration: %w", err)
+	}
+
+	dc, ok := m["discovery"].(map[string]any)
+	if !ok {
+		return fmt.Errorf("failed to get discovery configuration")
+	}
+
+	key, ok := dc["key"].(string)
+	if !ok {
+		return fmt.Errorf("failed to get discovery key")
+	}
+
+	if key != "" {
+		if 8 < len(key) {
+			key = key[:8] + "..."
+		} else {
+			key = "********"
+		}
+		viper.Set("discovery.key", key)
+	}
+	dc["key"] = key
+	m["discovery"] = dc
+
+	v := viper.New()
+	v.MergeConfigMap(m)
+
+	if len(args) == 0 || args[0] == "" {
+		return yaml.NewEncoder(os.Stdout).Encode(v.AllSettings())
+	}
+
+	configData := v.Get(args[0])
 	if configData == nil {
 		return output.UsageErrorF("no such key: %s", args[0])
 	}
diff --git a/v2/pkg/commands/discovery-server/cobra.go b/v2/pkg/commands/discovery-server/cobra.go
index 63cf66d..0829ed0 100644
--- a/v2/pkg/commands/discovery-server/cobra.go
+++ b/v2/pkg/commands/discovery-server/cobra.go
@@ -71,7 +71,7 @@ func (c *Cmd) run(cmd *cobra.Command, args []string) error {
 			if err != nil {
 				return fmt.Errorf("unable to get source ip: %w", err)
 			}
-			uaConf.Domain = sip
+			uaConf.Domain = sip.String()
 		}
 	}
 	if uaConf.Port == 0 {
diff --git a/v2/pkg/commands/p2p/client/receive/cobra.go b/v2/pkg/commands/p2p/client/receive/cobra.go
index 7dacd4c..7b4bb90 100644
--- a/v2/pkg/commands/p2p/client/receive/cobra.go
+++ b/v2/pkg/commands/p2p/client/receive/cobra.go
@@ -126,7 +126,7 @@ func (c *Cmd) receive(cmd *cobra.Command, args []string) error {
 
 		signaller, bat, err = signallers.NewFileClientSignaller(offerFilePath, answerFilePath)
 	} else {
-		corr, err := discovery.NegotiateOfferRequest(ctx, dsConfig.Host, baConfig.Username, baConfig.Password, http.DefaultClient)
+		corr, err := discovery.NegotiateOfferRequest(ctx, dsConfig.Addr, baConfig.Username, baConfig.Password, http.DefaultClient)
 		if err != nil {
 			return fmt.Errorf("failed to negotiate offer request: %w", err)
 		}
@@ -134,7 +134,7 @@ func (c *Cmd) receive(cmd *cobra.Command, args []string) error {
 		if err != nil {
 			return fmt.Errorf("failed to create transport: %w", err)
 		}
-		signaller, bat, err = signallers.NewServerClientSignaller(dsConfig.Host, corr.SessionID, corr.RTCSessionDescription, nil)
+		signaller, bat, err = signallers.NewServerClientSignaller(dsConfig.Addr, corr.SessionID, corr.RTCSessionDescription, nil)
 	}
 	if err != nil {
 		return fmt.Errorf("failed to create signaller: %w", err)
diff --git a/v2/pkg/commands/p2p/client/send/cobra.go b/v2/pkg/commands/p2p/client/send/cobra.go
index 00f90d3..e212aee 100644
--- a/v2/pkg/commands/p2p/client/send/cobra.go
+++ b/v2/pkg/commands/p2p/client/send/cobra.go
@@ -11,7 +11,6 @@ import (
 	"strings"
 	"time"
 
-	"github.com/moby/moby/pkg/namesgenerator"
 	"github.com/forestnode-io/oneshot/v2/pkg/commands/p2p/client/discovery"
 	"github.com/forestnode-io/oneshot/v2/pkg/commands/p2p/client/send/configuration"
 	rootconfig "github.com/forestnode-io/oneshot/v2/pkg/configuration"
@@ -22,6 +21,7 @@ import (
 	"github.com/forestnode-io/oneshot/v2/pkg/net/webrtc/sdp/signallers"
 	oneshotos "github.com/forestnode-io/oneshot/v2/pkg/os"
 	"github.com/forestnode-io/oneshot/v2/pkg/output"
+	"github.com/moby/moby/pkg/namesgenerator"
 	"github.com/pion/webrtc/v3"
 	"github.com/rs/zerolog"
 	"github.com/spf13/cobra"
@@ -69,7 +69,7 @@ func (c *Cmd) send(cmd *cobra.Command, args []string) error {
 
 		fileName            = config.Name
 		webRTCSignallingDir = p2pConfig.DiscoveryDir
-		webRTCSignallingURL = dsConfig.Host
+		webRTCSignallingURL = dsConfig.Addr
 
 		err error
 	)
diff --git a/v2/pkg/commands/root/discoveryServer.go b/v2/pkg/commands/root/discoveryServer.go
index 24836ff..6ddd3b3 100644
--- a/v2/pkg/commands/root/discoveryServer.go
+++ b/v2/pkg/commands/root/discoveryServer.go
@@ -18,7 +18,7 @@ func (r *rootCommand) sendArrivalToDiscoveryServer(ctx context.Context, cmd stri
 		dsConfig = config.Discovery
 	)
 
-	if dsConfig.Host == "" {
+	if dsConfig.Addr == "" {
 		return nil
 	}
 
@@ -29,11 +29,14 @@ func (r *rootCommand) sendArrivalToDiscoveryServer(ctx context.Context, cmd stri
 		Cmd:                cmd,
 	}
 
-	ipThatCanReachDiscoveryServer, err := oneshotnet.GetSourceIP(dsConfig.Host, 0)
+	ipThatCanReachDiscoveryServer, err := oneshotnet.GetSourceIP(dsConfig.Addr, 0)
 	if err != nil {
 		return fmt.Errorf("unable to reach the discovery server: %w", err)
 	}
-	arrival.Redirect = ipThatCanReachDiscoveryServer
+	arrival.Redirect = ipThatCanReachDiscoveryServer.String()
+	if ipThatCanReachDiscoveryServer.To16() != nil {
+		arrival.Redirect = "[" + arrival.Redirect + "]"
+	}
 
 	scheme := "http"
 	if config.Server.TLSCert != "" {
diff --git a/v2/pkg/commands/root/listenWebRTC.go b/v2/pkg/commands/root/listenWebRTC.go
index bd8fb6f..0fa6620 100644
--- a/v2/pkg/commands/root/listenWebRTC.go
+++ b/v2/pkg/commands/root/listenWebRTC.go
@@ -85,7 +85,7 @@ func getSignaller(ctx context.Context, config *configuration.Root, portMapAddr s
 	var (
 		dsConf              = config.Discovery
 		p2pConf             = config.NATTraversal.P2P
-		webRTCSignallingURL = dsConf.Host
+		webRTCSignallingURL = dsConf.Addr
 		webRTCSignallingDir = p2pConf.DiscoveryDir
 	)
 
diff --git a/v2/pkg/commands/root/runServer.go b/v2/pkg/commands/root/runServer.go
index b21fc21..ffd9003 100644
--- a/v2/pkg/commands/root/runServer.go
+++ b/v2/pkg/commands/root/runServer.go
@@ -41,6 +41,24 @@ func (r *rootCommand) init(cmd *cobra.Command, args []string) error {
 		return fmt.Errorf("failed to hydrate configuration: %w", err)
 	}
 
+	if r.config.Discovery.Addr != "" {
+		host, port, err := net.SplitHostPort(r.config.Discovery.Addr)
+		if err != nil {
+			if !strings.Contains(err.Error(), "missing port in address") {
+				return fmt.Errorf("failed to parse discovery server URL: %w", err)
+			}
+			host = r.config.Discovery.Addr
+		}
+		if port == "" {
+			if r.config.Discovery.Insecure {
+				port = "80"
+			} else {
+				port = "443"
+			}
+		}
+		r.config.Discovery.Addr = net.JoinHostPort(host, port)
+	}
+
 	if r.config.Output.Quiet {
 		output.Quiet(ctx)
 	} else {
@@ -140,7 +158,7 @@ func (r *rootCommand) runServer(cmd *cobra.Command, args []string) error {
 	dsConfig := r.config.Discovery
 	connConf := signallingserver.DiscoveryServerConfig{
 		Enabled:                   dsConfig.Enabled,
-		URL:                       dsConfig.Host,
+		URL:                       dsConfig.Addr,
 		Key:                       dsConfig.Key,
 		Insecure:                  dsConfig.Insecure,
 		SendReports:               dsConfig.Reports.Enabled,
@@ -152,7 +170,7 @@ func (r *rootCommand) runServer(cmd *cobra.Command, args []string) error {
 			APIVersion: version.APIVersion,
 		},
 	}
-	if err := signallingserver.ConnectToDiscoveryServer(ctx, connConf); err != nil {
+	if err := signallingserver.ConnectToDiscoveryServer(ctx, &connConf); err != nil {
 		log.Error().Err(err).
 			Msg("failed to connect to discovery server")
 	} else {
diff --git a/v2/pkg/configuration/configuration.go b/v2/pkg/configuration/configuration.go
index 17c258c..26c6573 100644
--- a/v2/pkg/configuration/configuration.go
+++ b/v2/pkg/configuration/configuration.go
@@ -188,7 +188,7 @@ func setDefault() {
 
 	// discovery
 	viper.SetDefault("discovery.enabled", true)
-	viper.SetDefault("discovery.host", "")
+	viper.SetDefault("discovery.addr", "")
 	viper.SetDefault("discovery.key", "")
 	viper.SetDefault("discovery.keypath", "")
 	viper.SetDefault("discovery.insecure", false)
diff --git a/v2/pkg/configuration/discovery.go b/v2/pkg/configuration/discovery.go
index 6dd5a87..a057096 100644
--- a/v2/pkg/configuration/discovery.go
+++ b/v2/pkg/configuration/discovery.go
@@ -22,7 +22,7 @@ type Reports struct {
 
 type Discovery struct {
 	Enabled      bool    `mapstructure:"enabled" yaml:"enabled"`
-	Host         string  `mapstructure:"host" yaml:"host"`
+	Addr         string  `mapstructure:"addr" yaml:"addr"`
 	Key          string  `mapstructure:"key" yaml:"key" json:"-"`
 	KeyPath      string  `mapstructure:"keyPath" yaml:"keyPath"`
 	Insecure     bool    `mapstructure:"insecure" yaml:"insecure"`
@@ -37,7 +37,7 @@ func setDiscoveryFlags(cmd *cobra.Command) {
 	defer cmd.PersistentFlags().AddFlagSet(fs)
 
 	flags.Bool(fs, "discovery.enabled", "discovery-enabled", "Enable discovery server.")
-	flags.String(fs, "discovery.url", "discovery-url", "URL of the discovery server to connect to.")
+	flags.String(fs, "discovery.addr", "discovery-addr", "Address (host:port) of the discovery server to connect to.")
 	flags.String(fs, "discovery.keypath", "discovery-key-path", "Path to the key to present to the discovery server.")
 	flags.String(fs, "discovery.key", "discovery-key", "Key to present to the discovery server.")
 	fs.Lookup("discovery-key").DefValue = ""
diff --git a/v2/pkg/file/tar.go b/v2/pkg/file/tar.go
index d35cecf..ba40383 100644
--- a/v2/pkg/file/tar.go
+++ b/v2/pkg/file/tar.go
@@ -30,7 +30,7 @@ func tarball(compress bool, paths []string, w io.Writer) error {
 		return name
 	}
 
-	walkFunc := func(path string) func(string, os.FileInfo, error) error {
+	walkFunc := func(path string, buf []byte) func(string, os.FileInfo, error) error {
 		dir := filepath.Dir(path)
 		return func(fp string, info os.FileInfo, err error) error {
 			if err != nil {
@@ -72,7 +72,7 @@ func tarball(compress bool, paths []string, w io.Writer) error {
 			}
 			defer fh.Close()
 
-			if _, err = io.Copy(tw, fh); err != nil {
+			if _, err = io.CopyBuffer(tw, fh, buf); err != nil {
 				return err
 			}
 
@@ -87,12 +87,18 @@ func tarball(compress bool, paths []string, w io.Writer) error {
 			return err
 		}
 		if info.IsDir() { // Archiving a directory; needs to be walked
-			err := filepath.Walk(path, walkFunc(path))
+			buf := make([]byte, 32*1024)
+			err := filepath.Walk(path, walkFunc(path, buf))
 			if err != nil {
 				return err
 			}
 		} else { // Archiving a single file or symlink
-			if err = walkFunc(path)(path, info, nil); err != nil {
+			size := info.Size()
+			if size == 0 {
+				size = 32 * 1024
+			}
+			buf := make([]byte, size)
+			if err = walkFunc(path, buf)(path, info, nil); err != nil {
 				return err
 			}
 		}
diff --git a/v2/pkg/net/network.go b/v2/pkg/net/network.go
index 3d34329..95820e1 100644
--- a/v2/pkg/net/network.go
+++ b/v2/pkg/net/network.go
@@ -6,8 +6,8 @@ import (
 	"strings"
 	"time"
 
-	"github.com/jackpal/gateway"
 	"github.com/forestnode-io/oneshot/v2/pkg/log"
+	"github.com/jackpal/gateway"
 )
 
 // HostAddresses returns all available ip addresses from all interfaces
@@ -45,11 +45,11 @@ func HostAddresses() ([]string, error) {
 // GetSourceIP returns the ip address used to access target:port
 // If target is the empty string then the default gateway ip is used.
 // If the port is 0, then 80 is used by default.
-func GetSourceIP(target string, port int) (string, error) {
+func GetSourceIP(target string, port int) (net.IP, error) {
 	if target == "" {
 		ip, err := gateway.DiscoverGateway()
 		if err != nil {
-			return "", err
+			return nil, err
 		}
 		target = ip.String()
 	}
@@ -64,13 +64,13 @@ func GetSourceIP(target string, port int) (string, error) {
 		conn, err = net.Dial("udp", target)
 	}
 	if err != nil {
-		return "", err
+		return nil, err
 	}
 	defer conn.Close()
 
 	localAddr := conn.LocalAddr().(*net.UDPAddr)
 
-	return localAddr.IP.String(), nil
+	return localAddr.IP, nil
 }
 
 func PreferNonPrivateIP(ips []string) (string, string) {
diff --git a/v2/pkg/net/webrtc/signallingserver/discoveryServer.go b/v2/pkg/net/webrtc/signallingserver/discoveryServer.go
index 2a07737..101163f 100644
--- a/v2/pkg/net/webrtc/signallingserver/discoveryServer.go
+++ b/v2/pkg/net/webrtc/signallingserver/discoveryServer.go
@@ -70,14 +70,19 @@ func WithDiscoveryServer(ctx context.Context) (context.Context, <-chan struct{})
 
 // ConnectToDiscoveryServer connects to the discovery server and sends a handshake.
 // The connection is kept open and stuffed into the context.
-func ConnectToDiscoveryServer(ctx context.Context, c DiscoveryServerConfig) error {
+func ConnectToDiscoveryServer(ctx context.Context, c *DiscoveryServerConfig) error {
 	dds, ok := ctx.Value(discoveryServerKey{}).(**DiscoveryServer)
 	if !ok || dds == nil {
 		return nil
 	}
 
+	defer func() {
+		if (*dds).config == nil {
+			(*dds).config = c
+		}
+	}()
+
 	if !c.Enabled {
-		(*dds).config = &c
 		return nil
 	}
 
@@ -85,17 +90,17 @@ func ConnectToDiscoveryServer(ctx context.Context, c DiscoveryServerConfig) erro
 	log := zerolog.Ctx(ctx)
 
 	log.Debug().Msg("connecting to discovery server")
-	conn, err := getConnectionToDiscoveryServer(ctx, &c)
+	conn, err := getConnectionToDiscoveryServer(ctx, c)
 	if err != nil {
 		return fmt.Errorf("failed to connect to discovery server: %w", err)
 	}
 
 	log.Debug().Msg("opening bidirectional stream to discovery server")
-	ds, err := newDiscoveryServer(ctx, &c, conn)
+	ds, err := newDiscoveryServer(ctx, c, conn)
 	if err != nil {
 		return fmt.Errorf("failed to create discovery server: %w", err)
 	}
-	ds.config = &c
+	ds.config = c
 	ds.doneFunc = (*dds).doneFunc
 	*dds = ds
 	return nil