From d5c980c4f044d90026642fd41406d460e4545635 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=90=B4=E6=B4=B2=E6=B4=8B?= Date: Tue, 19 Jan 2021 20:56:50 +0800 Subject: [PATCH] =?UTF-8?q?=E4=BF=AE=E6=94=B9io.copy?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- demo/gateway/main.go | 2 +- demo/tunnel/client/tcp/build.sh | 3 +++ demo/tunnel/client/tcp/main.go | 39 ++++++++++++++++++++++++++++++--- demo/tunnel/server/build.sh | 3 +++ demo/tunnel/server/main.go | 33 ++++++++++++++++++++++++++-- 5 files changed, 74 insertions(+), 6 deletions(-) create mode 100755 demo/tunnel/client/tcp/build.sh create mode 100755 demo/tunnel/server/build.sh diff --git a/demo/gateway/main.go b/demo/gateway/main.go index 72af98b5..0a63c589 100644 --- a/demo/gateway/main.go +++ b/demo/gateway/main.go @@ -13,7 +13,7 @@ var local = flag.String("local", "", "please enter monitor ip, example: 192.168. var target = flag.String("target", "", "please enter target ip, example : 192.168.1.105:8806") func main() { - + flag.Parse() l, err := net.Listen("tcp", *local) if err != nil { fmt.Println(err, err.Error()) diff --git a/demo/tunnel/client/tcp/build.sh b/demo/tunnel/client/tcp/build.sh new file mode 100755 index 00000000..f0f48dd6 --- /dev/null +++ b/demo/tunnel/client/tcp/build.sh @@ -0,0 +1,3 @@ +#!/bin/zsh + +GOARCH=amd64 GOOS=linux go build -ldflags '-w -s' -o tcp-tunnle-client \ No newline at end of file diff --git a/demo/tunnel/client/tcp/main.go b/demo/tunnel/client/tcp/main.go index 9d5bfb0d..a588442c 100644 --- a/demo/tunnel/client/tcp/main.go +++ b/demo/tunnel/client/tcp/main.go @@ -22,12 +22,13 @@ const ( ) var local = flag.String("local", "127.0.0.1:1080", "please enter local proxy ip") + var target = flag.String("target", "192.168.1.105:8806", "please enter target server ip") var server = flag.String("server", "192.168.0.104:8805", "please enter tcp tunnel server ip") func main() { flag.Parse() - tgt := ParseAddr(*target) + //tgt := ParseAddr(*target) l, err := net.Listen("tcp", *local) if err != nil { log.Fatalf("failed to listen on %s: %v", *local, err) @@ -42,12 +43,14 @@ func main() { for { c, err := l.Accept() + if err != nil { continue } go func() { + tgt := ParseAddr(*target) rc, err := net.Dial("tcp", *server) if err != nil { log.Fatal("can't connect " + *server) @@ -65,6 +68,7 @@ func main() { if err = relay(rc, c); err != nil { log.Fatal(err.Error() + "\n") } + }() } @@ -78,10 +82,10 @@ func relay(left, right net.Conn) error { wg.Add(1) go func() { defer wg.Done() - _, err1 = io.Copy(right, left) + _, err1 = SecureCopy(right, left) right.SetReadDeadline(time.Now().Add(wait)) // unblock read on right }() - _, err = io.Copy(left, right) + _, err = SecureCopy(left, right) left.SetReadDeadline(time.Now().Add(wait)) // unblock read on left wg.Wait() if err1 != nil && !errors.Is(err1, os.ErrDeadlineExceeded) { // requires Go 1.15+ @@ -129,3 +133,32 @@ func ParseAddr(s string) Addr { return addr } + +func SecureCopy(src io.ReadWriteCloser, dst io.ReadWriteCloser) (written int64, err error) { + size := 1024 + buf := make([]byte, size) + for { + nr, er := src.Read(buf) + if nr > 0 { + nw, ew := dst.Write(buf[0:nr]) + if nw > 0 { + written += int64(nw) + } + if ew != nil { + err = ew + break + } + if nr != nw { + err = io.ErrShortWrite + break + } + } + if er != nil { + if er != io.EOF { + err = er + } + break + } + } + return written, err +} diff --git a/demo/tunnel/server/build.sh b/demo/tunnel/server/build.sh new file mode 100755 index 00000000..d5f164ad --- /dev/null +++ b/demo/tunnel/server/build.sh @@ -0,0 +1,3 @@ +#!/bin/zsh + +GOARCH=amd64 GOOS=linux go build -ldflags '-w -s' -o tcp-tunnle-server diff --git a/demo/tunnel/server/main.go b/demo/tunnel/server/main.go index 1d911148..484027ac 100644 --- a/demo/tunnel/server/main.go +++ b/demo/tunnel/server/main.go @@ -93,10 +93,10 @@ func relay(left, right net.Conn) error { wg.Add(1) go func() { defer wg.Done() - _, err1 = io.Copy(right, left) + _, err1 = SecureCopy(right, left) right.SetReadDeadline(time.Now().Add(wait)) // unblock read on right }() - _, err = io.Copy(left, right) + _, err = SecureCopy(left, right) left.SetReadDeadline(time.Now().Add(wait)) // unblock read on left wg.Wait() if err1 != nil && !errors.Is(err1, os.ErrDeadlineExceeded) { // requires Go 1.15+ @@ -140,3 +140,32 @@ func readAddr(r io.Reader, b []byte) (Addr, error) { return nil, err } + +func SecureCopy(src io.ReadWriteCloser, dst io.ReadWriteCloser) (written int64, err error) { + size := 1024 + buf := make([]byte, size) + for { + nr, er := src.Read(buf) + if nr > 0 { + nw, ew := dst.Write(buf[0:nr]) + if nw > 0 { + written += int64(nw) + } + if ew != nil { + err = ew + break + } + if nr != nw { + err = io.ErrShortWrite + break + } + } + if er != nil { + if er != io.EOF { + err = er + } + break + } + } + return written, err +}