From 774a9a08028ea87798a6b5111ce81411cd8c4aea Mon Sep 17 00:00:00 2001
From: firefart <firefart@gmail.com>
Date: Wed, 11 May 2022 21:21:20 +0200
Subject: [PATCH] we still need the protocol parameter

---
 internal/cmd/info.go | 17 ++++++-----------
 main.go              |  3 +++
 2 files changed, 9 insertions(+), 11 deletions(-)

diff --git a/internal/cmd/info.go b/internal/cmd/info.go
index a3d60358..ee6b8945 100644
--- a/internal/cmd/info.go
+++ b/internal/cmd/info.go
@@ -13,6 +13,7 @@ import (
 type InfoOpts struct {
 	TurnServer string
 	UseTLS     bool
+	Protocol   string
 	Timeout    time.Duration
 	Log        *logrus.Logger
 }
@@ -24,6 +25,9 @@ func (opts InfoOpts) Validate() error {
 	if !strings.Contains(opts.TurnServer, ":") {
 		return fmt.Errorf("turnserver needs a port")
 	}
+	if opts.Protocol != "tcp" && opts.Protocol != "udp" {
+		return fmt.Errorf("protocol needs to be either tcp or udp")
+	}
 	if opts.Log == nil {
 		return fmt.Errorf("please supply a valid logger")
 	}
@@ -64,7 +68,7 @@ func Info(opts InfoOpts) error {
 }
 
 func testStun(opts InfoOpts) ([]internal.Attribute, error) {
-	conn, err := internal.Connect("udp", opts.TurnServer, opts.UseTLS, opts.Timeout)
+	conn, err := internal.Connect(opts.Protocol, opts.TurnServer, opts.UseTLS, opts.Timeout)
 	if err != nil {
 		return nil, err
 	}
@@ -83,16 +87,7 @@ func testStun(opts InfoOpts) ([]internal.Attribute, error) {
 }
 
 func testTurn(opts InfoOpts, proto internal.RequestedTransport) ([]internal.Attribute, error) {
-	var protoString string
-	switch proto {
-	case internal.RequestedTransportTCP:
-		protoString = "tcp"
-	case internal.RequestedTransportUDP:
-		protoString = "udp"
-	default:
-		protoString = "udp"
-	}
-	conn, err := internal.Connect(protoString, opts.TurnServer, opts.UseTLS, opts.Timeout)
+	conn, err := internal.Connect(opts.Protocol, opts.TurnServer, opts.UseTLS, opts.Timeout)
 	if err != nil {
 		return nil, err
 	}
diff --git a/main.go b/main.go
index 25cbd9e3..9b39306d 100644
--- a/main.go
+++ b/main.go
@@ -54,6 +54,7 @@ func main() {
 					&cli.BoolFlag{Name: "debug", Aliases: []string{"d"}, Value: false, Usage: "enable debug output"},
 					&cli.StringFlag{Name: "turnserver", Aliases: []string{"s"}, Required: true, Usage: "turn server to connect to in the format host:port"},
 					&cli.BoolFlag{Name: "tls", Value: false, Usage: "Use TLS for connecting (false in most tests)"},
+					&cli.StringFlag{Name: "protocol", Value: "udp", Usage: "protocol to use when connecting to the TURN server. Supported values: tcp and udp"},
 					&cli.DurationFlag{Name: "timeout", Value: 1 * time.Second, Usage: "connect timeout to turn server"},
 				},
 				Before: func(ctx *cli.Context) error {
@@ -65,10 +66,12 @@ func main() {
 				Action: func(c *cli.Context) error {
 					turnServer := c.String("turnserver")
 					useTLS := c.Bool("tls")
+					protocol := c.String("protocol")
 					timeout := c.Duration("timeout")
 					return cmd.Info(cmd.InfoOpts{
 						TurnServer: turnServer,
 						UseTLS:     useTLS,
+						Protocol:   protocol,
 						Log:        log,
 						Timeout:    timeout,
 					})