@@ -14,6 +14,7 @@ import (
1414 "syscall"
1515 "time"
1616
17+ "github.com/pkg/errors"
1718 "github.com/sirupsen/logrus"
1819 "github.com/spf13/cobra"
1920)
@@ -27,7 +28,9 @@ var rootCmd = &cobra.Command{
2728func init () {
2829 rootCmd .Flags ().BoolP ("quiet" , "q" , false , "Quiet mode" )
2930 rootCmd .Flags ().StringP ("source" , "s" , "" , "Source address" )
31+ rootCmd .MarkFlagRequired ("source" )
3032 rootCmd .Flags ().StringP ("destination" , "d" , "" , "Destination address" )
33+ rootCmd .MarkFlagRequired ("destination" )
3134}
3235
3336func Execute () {
@@ -47,14 +50,18 @@ var (
4750 SIGTERM = syscall .Signal (0xf )
4851)
4952
50- func listen (addr string ) (net.Listener , error ) {
51- //listen
52- network := "tcp"
53- if strings .HasPrefix (addr , "unix:" ) {
54- network = "unix"
55- addr = strings .TrimPrefix (addr , "unix:" )
53+ func listen (url string ) (net.Listener , error ) {
54+ parts := strings .SplitN (url , "://" , 2 )
55+ if len (parts ) != 2 {
56+ return nil , errors .Errorf ("invalid url: %s" , url )
5657 }
57- return net .Listen (network , addr )
58+ proto := parts [0 ]
59+ addr := parts [1 ]
60+ listener , err := net .Listen (proto , addr )
61+ if err != nil {
62+ return nil , errors .WithStack (err )
63+ }
64+ return listener , err
5865}
5966
6067func runAction (cmd * cobra.Command , args []string ) error {
@@ -120,13 +127,14 @@ var pool = sync.Pool{
120127 },
121128}
122129
123- func dial (destination string ) (net.Conn , error ) {
124- network := "tcp"
125- if strings .HasPrefix (destination , "unix:" ) {
126- network = "unix"
127- destination = strings .TrimPrefix (destination , "unix:" )
130+ func dial (url string ) (net.Conn , error ) {
131+ parts := strings .SplitN (url , "://" , 2 )
132+ if len (parts ) != 2 {
133+ return nil , errors .Errorf ("invalid url: %s" , url )
128134 }
129- return net .Dial (network , destination )
135+ proto := parts [0 ]
136+ addr := parts [1 ]
137+ return net .Dial (proto , addr )
130138}
131139
132140func fwd (uconn net.Conn , destination string , quiet bool ) {
0 commit comments