Generic TCP Corking

zy410692 · Aug 16, 2020 · d551e24 · d551e24
1 parent 09d8cb9
commit d551e24
Show file tree

Hide file tree

Showing 4 changed files with 41 additions and 39 deletions.
diff --git a/tcp.go b/tcp.go
@@ -1,6 +1,7 @@
 package main
 
 import (
+	"bufio"
 	"io"
 	"io/ioutil"
 	"net"
@@ -71,9 +72,8 @@ func tcpLocal(addr, server string, shadow func(net.Conn) net.Conn, getAddr func(
 				return
 			}
 			defer rc.Close()
-			tc := rc.(*net.TCPConn)
 			if config.TCPCork {
-				timedCork(tc, 10*time.Millisecond)
+				rc = timedCork(rc, 10*time.Millisecond, 1280)
 			}
 			rc = shadow(rc)
 
@@ -166,3 +166,42 @@ func relay(left, right net.Conn) error {
 	}
 	return err
 }
+
+type corkedConn struct {
+	net.Conn
+	bufw   *bufio.Writer
+	corked bool
+	delay  time.Duration
+	err    error
+	lock   sync.Mutex
+	once   sync.Once
+}
+
+func timedCork(c net.Conn, d time.Duration, bufSize int) net.Conn {
+	return &corkedConn{
+		Conn:   c,
+		bufw:   bufio.NewWriterSize(c, bufSize),
+		corked: true,
+		delay:  d,
+	}
+}
+
+func (w *corkedConn) Write(p []byte) (int, error) {
+	w.lock.Lock()
+	defer w.lock.Unlock()
+	if w.err != nil {
+		return 0, w.err
+	}
+	if w.corked {
+		w.once.Do(func() {
+			time.AfterFunc(w.delay, func() {
+				w.lock.Lock()
+				defer w.lock.Unlock()
+				w.corked = false
+				w.err = w.bufw.Flush()
+			})
+		})
+		return w.bufw.Write(p)
+	}
+	return w.Conn.Write(p)
+}
diff --git a/tcp_darwin.go b/tcp_darwin.go
@@ -2,8 +2,6 @@ package main
 
 import (
 	"net"
-	"syscall"
-	"time"
 
 	"github.com/shadowsocks/go-shadowsocks2/pfutil"
 	"github.com/shadowsocks/go-shadowsocks2/socks"
@@ -24,18 +22,3 @@ func natLookup(c net.Conn) (socks.Addr, error) {
 	}
 	panic("not TCP connection")
 }
-
-func timedCork(c *net.TCPConn, d time.Duration) error {
-	rc, err := c.SyscallConn()
-	if err != nil {
-		return err
-	}
-	rc.Control(func(fd uintptr) { err = syscall.SetsockoptInt(int(fd), syscall.IPPROTO_TCP, syscall.TCP_NOPUSH, 1) })
-	if err != nil {
-		return err
-	}
-	time.AfterFunc(d, func() {
-		rc.Control(func(fd uintptr) { syscall.SetsockoptInt(int(fd), syscall.IPPROTO_TCP, syscall.TCP_NOPUSH, 0) })
-	})
-	return nil
-}
diff --git a/tcp_linux.go b/tcp_linux.go
@@ -2,8 +2,6 @@ package main
 
 import (
 	"net"
-	"syscall"
-	"time"
 
 	"github.com/shadowsocks/go-shadowsocks2/nfutil"
 	"github.com/shadowsocks/go-shadowsocks2/socks"
@@ -28,18 +26,3 @@ func redir6Local(addr, server string, shadow func(net.Conn) net.Conn) {
 	logf("TCP6 redirect %s <-> %s", addr, server)
 	tcpLocal(addr, server, shadow, func(c net.Conn) (socks.Addr, error) { return getOrigDst(c, true) })
 }
-
-func timedCork(c *net.TCPConn, d time.Duration) error {
-	rc, err := c.SyscallConn()
-	if err != nil {
-		return err
-	}
-	rc.Control(func(fd uintptr) { err = syscall.SetsockoptInt(int(fd), syscall.IPPROTO_TCP, syscall.TCP_CORK, 1) })
-	if err != nil {
-		return err
-	}
-	time.AfterFunc(d, func() {
-		rc.Control(func(fd uintptr) { syscall.SetsockoptInt(int(fd), syscall.IPPROTO_TCP, syscall.TCP_CORK, 0) })
-	})
-	return nil
-}
diff --git a/tcp_other.go b/tcp_other.go
@@ -4,7 +4,6 @@ package main
 
 import (
 	"net"
-	"time"
 )
 
 func redirLocal(addr, server string, shadow func(net.Conn) net.Conn) {
@@ -14,5 +13,3 @@ func redirLocal(addr, server string, shadow func(net.Conn) net.Conn) {
 func redir6Local(addr, server string, shadow func(net.Conn) net.Conn) {
 	logf("TCP6 redirect not supported")
 }
-
-func timedCork(c *net.TCPConn, d time.Duration) error { return nil }