Clean up

Author: wzshiming

all_test.go

@@ -22,9 +22,10 @@ func TestProxyCommandUnit(t *testing.T) {
 				return nil, err
 			}
 			conn1, conn2 := net.Pipe()
+			var buf1, buf2 [32 * 1024]byte
 			go func() {
 				ctx, cancel := context.WithCancel(ctx)
-				err := Tunnel(ctx, &connWithCancel{conn, cancel}, conn1)
+				err := Tunnel(ctx, &connWithCancel{conn, cancel}, conn1, buf1[:], buf2[:])
 				if err != nil {
 					t.Error(err)
 					return

cmd/commandproxy/main.go

@@ -27,8 +27,8 @@ func main() {
 		log.Fatalln(err)
 		return
 	}
-
-	err = commandproxy.Tunnel(context.Background(), commandproxy.Stdio, conn)
+	var buf1, buf2 [32 * 1024]byte
+	err = commandproxy.Tunnel(context.Background(), commandproxy.Stdio, conn, buf1[:], buf2[:])
 	if err != nil {
 		log.Fatalln(err)
 		return

dial.go

@@ -7,10 +7,12 @@ import (
 
 type DialProxyCommand []string
 
-func (p *DialProxyCommand) Format(network string, address string) ([]string, error) {
+func (p *DialProxyCommand) Format(network string, address string) []string {
 	host, port, err := net.SplitHostPort(address)
 	if err != nil {
-		return nil, err
+		host = address
+	} else if host == "" {
+		host = "0.0.0.0"
 	}
 	m := map[byte]string{
 		'n': network,
@@ -22,14 +24,11 @@ func (p *DialProxyCommand) Format(network string, address string) ([]string, err
 	for i := range proxy {
 		proxy[i] = ReplaceEscape(proxy[i], m)
 	}
-	return proxy, nil
+	return proxy
 }
 
 func (p *DialProxyCommand) DialContext(ctx context.Context, network string, address string) (net.Conn, error) {
-	proxy, err := p.Format(network, address)
-	if err != nil {
-		return nil, err
-	}
+	proxy := p.Format(network, address)
 	cp := ProxyCommand(ctx, proxy[0], proxy[1:]...)
 	return cp.Stdio()
 }

dial_test.go

@@ -3,6 +3,7 @@ package commandproxy
 import (
 	"context"
 	"net"
+	"reflect"
 	"testing"
 )
 
@@ -26,3 +27,48 @@ func TestDialProxyCommand(t *testing.T) {
 		t.Fatal(err)
 	}
 }
+
+func TestDialProxyCommand_Format(t *testing.T) {
+	type args struct {
+		network string
+		address string
+	}
+	tests := []struct {
+		name string
+		p    DialProxyCommand
+		args args
+		want []string
+	}{
+		{
+			p: DialProxyCommand{"nc", "%h", "%p"},
+			args: args{
+				network: "tcp",
+				address: ":1",
+			},
+			want: []string{"nc", "0.0.0.0", "1"},
+		},
+		{
+			p: DialProxyCommand{"nc", "%h", "%p"},
+			args: args{
+				network: "tcp",
+				address: "[::]:1",
+			},
+			want: []string{"nc", "::", "1"},
+		},
+		{
+			p: DialProxyCommand{"nc", "-U", "%h"},
+			args: args{
+				network: "unix",
+				address: "./x.socks",
+			},
+			want: []string{"nc", "-U", "./x.socks"},
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			if got := tt.p.Format(tt.args.network, tt.args.address); !reflect.DeepEqual(got, tt.want) {
+				t.Errorf("Format() = %v, want %v", got, tt.want)
+			}
+		})
+	}
+}

tunnel.go

@@ -5,17 +5,15 @@ import (
 	"io"
 )
 
-func Tunnel(ctx context.Context, c1, c2 io.ReadWriteCloser) error {
+func Tunnel(ctx context.Context, c1, c2 io.ReadWriteCloser, buf1, buf2 []byte) error {
 	ctx, cancel := context.WithCancel(ctx)
 	var errs tunnelErr
 	go func() {
-		var buf [32 * 1024]byte
-		_, errs[0] = io.CopyBuffer(c1, c2, buf[:])
+		_, errs[0] = io.CopyBuffer(c1, c2, buf1)
 		cancel()
 	}()
 	go func() {
-		var buf [32 * 1024]byte
-		_, errs[1] = io.CopyBuffer(c2, c1, buf[:])
+		_, errs[1] = io.CopyBuffer(c2, c1, buf2)
 		cancel()
 	}()
 	<-ctx.Done()

util.go

@@ -8,13 +8,13 @@ import (
 
 var (
 	errQuoteNotClose = fmt.Errorf("quote is not closed")
-	errEmptyCmmand   = fmt.Errorf("cmd is empty")
+	errEmptyCommand  = fmt.Errorf("cmd is empty")
 )
 
 func SplitCommand(cmd string) ([]string, error) {
 	cmd = strings.TrimSpace(cmd)
 	if cmd == "" {
-		return nil, errEmptyCmmand
+		return nil, errEmptyCommand
 	}
 	var ss []string