diff --git a/cmd/bench/cmd.go b/cmd/bench/cmd.go index 16e793a..93369bd 100644 --- a/cmd/bench/cmd.go +++ b/cmd/bench/cmd.go @@ -1,6 +1,7 @@ package bench import ( + "github.com/antonito/gfile/internal/utils" "github.com/antonito/gfile/pkg/session/bench" "github.com/antonito/gfile/pkg/session/common" log "github.com/sirupsen/logrus" @@ -10,13 +11,23 @@ import ( func handler(c *cli.Context) error { isMaster := c.Bool("master") - sess := bench.NewWith(bench.Config{ + conf := bench.Config{ Master: isMaster, Configuration: common.Configuration{ OnCompletion: func() { }, }, - }) + } + + customSTUN := c.String("stun") + if customSTUN != "" { + if err := utils.ParseSTUN(customSTUN); err != nil { + return err + } + conf.STUN = customSTUN + } + + sess := bench.NewWith(conf) return sess.Start() } @@ -33,6 +44,10 @@ func New() cli.Command { Name: "master, m", Usage: "Is creating the SDP offer?", }, + cli.StringFlag{ + Name: "stun", + Usage: "Use a specific STUN server (ex: --stun stun.l.google.com:19302)", + }, }, } } diff --git a/cmd/receive/cmd.go b/cmd/receive/cmd.go index 7093460..f39a64b 100644 --- a/cmd/receive/cmd.go +++ b/cmd/receive/cmd.go @@ -6,6 +6,7 @@ import ( log "github.com/sirupsen/logrus" + "github.com/antonito/gfile/internal/utils" "github.com/antonito/gfile/pkg/session/receiver" "gopkg.in/urfave/cli.v1" ) @@ -21,9 +22,19 @@ func handler(c *cli.Context) error { } defer f.Close() - sess := receiver.NewWith(receiver.Config{ + conf := receiver.Config{ Stream: f, - }) + } + + customSTUN := c.String("stun") + if customSTUN != "" { + if err := utils.ParseSTUN(customSTUN); err != nil { + return err + } + conf.STUN = customSTUN + } + + sess := receiver.NewWith(conf) return sess.Start() } @@ -40,6 +51,10 @@ func New() cli.Command { Name: "output, o", Usage: "Output", }, + cli.StringFlag{ + Name: "stun", + Usage: "Use a specific STUN server (ex: --stun stun.l.google.com:19302)", + }, }, } } diff --git a/cmd/send/cmd.go b/cmd/send/cmd.go index ceb76bb..0c82670 100644 --- a/cmd/send/cmd.go +++ b/cmd/send/cmd.go @@ -4,6 +4,7 @@ import ( "fmt" "os" + "github.com/antonito/gfile/internal/utils" "github.com/antonito/gfile/pkg/session/common" "github.com/antonito/gfile/pkg/session/sender" log "github.com/sirupsen/logrus" @@ -20,13 +21,23 @@ func handler(c *cli.Context) error { return err } defer f.Close() - sess := sender.NewWith(sender.Config{ + conf := sender.Config{ Stream: f, Configuration: common.Configuration{ OnCompletion: func() { }, }, - }) + } + + customSTUN := c.String("stun") + if customSTUN != "" { + if err := utils.ParseSTUN(customSTUN); err != nil { + return err + } + conf.STUN = customSTUN + } + + sess := sender.NewWith(conf) return sess.Start() } @@ -43,6 +54,10 @@ func New() cli.Command { Name: "file, f", Usage: "Send content of file `FILE`", }, + cli.StringFlag{ + Name: "stun", + Usage: "Use a specific STUN server (ex: --stun stun.l.google.com:19302)", + }, }, } } diff --git a/internal/session/session.go b/internal/session/session.go index d762739..c8209b2 100644 --- a/internal/session/session.go +++ b/internal/session/session.go @@ -21,22 +21,29 @@ type Session struct { sdpOutput io.Writer peerConnection *webrtc.PeerConnection onCompletion CompletionHandler + stunServers []string } // New creates a new Session -func New(sdpInput io.Reader, sdpOutput io.Writer) Session { - if sdpInput == nil { - sdpInput = os.Stdin - } - if sdpOutput == nil { - sdpOutput = os.Stdout - } - return Session{ +func New(sdpInput io.Reader, sdpOutput io.Writer, customSTUN string) Session { + sess := Session{ sdpInput: sdpInput, sdpOutput: sdpOutput, Done: make(chan struct{}), NetworkStats: stats.New(), + stunServers: []string{fmt.Sprintf("stun:%s", customSTUN)}, + } + + if sdpInput == nil { + sess.sdpInput = os.Stdin + } + if sdpOutput == nil { + sess.sdpOutput = os.Stdout + } + if customSTUN == "" { + sess.stunServers = []string{"stun:stun.l.google.com:19302"} } + return sess } // CreateConnection prepares a WebRTC connection @@ -44,7 +51,7 @@ func (s *Session) CreateConnection(onConnectionStateChange func(connectionState config := webrtc.Configuration{ ICEServers: []webrtc.ICEServer{ { - URLs: []string{"stun:stun.l.google.com:19302"}, + URLs: s.stunServers, }, }, } diff --git a/internal/utils/stun_arg.go b/internal/utils/stun_arg.go new file mode 100644 index 0000000..b732308 --- /dev/null +++ b/internal/utils/stun_arg.go @@ -0,0 +1,20 @@ +package utils + +import ( + "fmt" + "strconv" + "strings" +) + +// ParseSTUN checks if a STUN addr is valid +func ParseSTUN(stunAddr string) error { + arr := strings.Split(stunAddr, ":") + if len(arr) != 2 { + return fmt.Errorf("invalid stun adress") + } + port, err := strconv.Atoi(arr[1]) + if err != nil || (port <= 0 || port > 0xffff) { + return fmt.Errorf("invalid port %v", port) + } + return nil +} diff --git a/internal/utils/stun_arg_test.go b/internal/utils/stun_arg_test.go new file mode 100644 index 0000000..f395cf5 --- /dev/null +++ b/internal/utils/stun_arg_test.go @@ -0,0 +1,56 @@ +package utils + +import ( + "fmt" + "testing" + + "github.com/stretchr/testify/assert" +) + +func Test_STUN_Arg(t *testing.T) { + assert := assert.New(t) + + tests := []struct { + input string + err error + }{ + { + input: "", + err: fmt.Errorf("invalid stun adress"), + }, + { + input: "test", + err: fmt.Errorf("invalid stun adress"), + }, + { + input: "stun:lol:lol", + err: fmt.Errorf("invalid stun adress"), + }, + { + input: "test:wtf", + err: fmt.Errorf("invalid port 0"), + }, + { + input: "test:-2", + err: fmt.Errorf("invalid port -2"), + }, + { + // 0xffff + 1 + input: "test:65536", + err: fmt.Errorf("invalid port 65536"), + }, + { + input: "test:5432", + err: nil, + }, + { + input: "stun.l.google.com:19302", + err: nil, + }, + } + + for _, cur := range tests { + err := ParseSTUN(cur.input) + assert.Equal(cur.err, err) + } +} diff --git a/pkg/session/bench/benchmark.go b/pkg/session/bench/benchmark.go index 5c4854a..6f2c77c 100644 --- a/pkg/session/bench/benchmark.go +++ b/pkg/session/bench/benchmark.go @@ -57,5 +57,5 @@ type Config struct { // NewWith createa a new benchmark Session with custom configuration func NewWith(c Config) *Session { - return new(internalSess.New(c.SDPProvider, c.SDPOutput), c.Master) + return new(internalSess.New(c.SDPProvider, c.SDPOutput, c.STUN), c.Master) } diff --git a/pkg/session/common/config.go b/pkg/session/common/config.go index d2a2993..2c20aba 100644 --- a/pkg/session/common/config.go +++ b/pkg/session/common/config.go @@ -11,4 +11,5 @@ type Configuration struct { SDPProvider io.Reader // The SDP reader SDPOutput io.Writer // The SDP writer OnCompletion session.CompletionHandler // Handler to call on session completion + STUN string // Custom STUN server } diff --git a/pkg/session/receiver/receiver.go b/pkg/session/receiver/receiver.go index 045b273..3cccfbc 100644 --- a/pkg/session/receiver/receiver.go +++ b/pkg/session/receiver/receiver.go @@ -27,7 +27,7 @@ func new(s internalSess.Session, f io.Writer) *Session { // New creates a new receiver session func New(f io.Writer) *Session { - return new(internalSess.New(nil, nil), f) + return new(internalSess.New(nil, nil, ""), f) } // Config contains custom configuration for a session @@ -38,7 +38,7 @@ type Config struct { // NewWith createa a new receiver Session with custom configuration func NewWith(c Config) *Session { - return new(internalSess.New(c.SDPProvider, c.SDPOutput), c.Stream) + return new(internalSess.New(c.SDPProvider, c.SDPOutput, c.STUN), c.Stream) } // SetStream changes the stream, useful for WASM integration diff --git a/pkg/session/sender/sender.go b/pkg/session/sender/sender.go index 6dec882..db2f551 100644 --- a/pkg/session/sender/sender.go +++ b/pkg/session/sender/sender.go @@ -55,7 +55,7 @@ func new(s internalSess.Session, f io.Reader) *Session { // New creates a new receiver session func New(f io.Reader) *Session { - return new(internalSess.New(nil, nil), f) + return new(internalSess.New(nil, nil, ""), f) } // Config contains custom configuration for a session @@ -66,7 +66,7 @@ type Config struct { // NewWith createa a new sender Session with custom configuration func NewWith(c Config) *Session { - return new(internalSess.New(c.SDPProvider, c.SDPOutput), c.Stream) + return new(internalSess.New(c.SDPProvider, c.SDPOutput, c.STUN), c.Stream) } // SetStream changes the stream, useful for WASM integration