feat: add delay option

conn.go

@@ -1,10 +1,12 @@
 package main
 
 import (
+	"errors"
 	"fmt"
 	"io"
 	"net"
 	"sync"
+	"time"
 
 	"github.com/fatih/color"
 	"github.com/kevwan/tproxy/display"
@@ -17,21 +19,28 @@ const (
 	useOfClosedConn = "use of closed network connection"
 )
 
+var clientCanceled = errors.New("client canceled")
+
 type PairedConnection struct {
-	id      int
-	cliConn net.Conn
-	svrConn net.Conn
-	once    sync.Once
+	id       int
+	cliConn  net.Conn
+	svrConn  net.Conn
+	once     sync.Once
+	stopChan chan struct{}
 }
 
 func NewPairedConnection(id int, cliConn net.Conn) *PairedConnection {
 	return &PairedConnection{
-		id:      id,
-		cliConn: cliConn,
+		id:       id,
+		cliConn:  cliConn,
+		stopChan: make(chan struct{}),
 	}
 }
 
 func (c *PairedConnection) handleClientMessage() {
+	// client closed also trigger server close.
+	defer c.stop()
+
 	r, w := io.Pipe()
 	tee := io.MultiWriter(c.svrConn, w)
 
@@ -44,8 +53,11 @@ func (c *PairedConnection) handleClientMessage() {
 }
 
 func (c *PairedConnection) handleServerMessage() {
+	// server closed also trigger client close.
+	defer c.stop()
+
 	r, w := io.Pipe()
-	tee := io.MultiWriter(c.cliConn, w)
+	tee := io.MultiWriter(newDelayedWriter(c.cliConn, settings.Delay, c.stopChan), w)
 	go protocol.NewDumper(r, serverSide, c.id, settings.Silent, protocol.CreateInterop(settings.Protocol)).Dump()
 	_, e := io.Copy(tee, c.svrConn)
 	if e != nil && e != io.EOF {
@@ -54,13 +66,12 @@ func (c *PairedConnection) handleServerMessage() {
 			color.HiRed("handleServerMessage: io.Copy error: %v", e)
 		}
 	}
-
-	c.stop()
 }
 
 func (c *PairedConnection) process() {
-	conn, err := net.Dial("tcp", settings.RemoteHost)
 	defer c.stop()
+
+	conn, err := net.Dial("tcp", settings.RemoteHost)
 	if err != nil {
 		display.PrintlnWithTime(color.HiRedString("[x][%d] Couldn't connect to server: %v", c.id, err))
 		return
@@ -76,6 +87,7 @@ func (c *PairedConnection) process() {
 
 func (c *PairedConnection) stop() {
 	c.once.Do(func() {
+		close(c.stopChan)
 		if c.cliConn != nil {
 			display.PrintlnWithTime(color.HiBlueString("[%d] Client connection closed", c.id))
 			c.cliConn.Close()
@@ -110,3 +122,32 @@ func startListener() error {
 		go pconn.process()
 	}
 }
+
+type delayedWriter struct {
+	writer   io.Writer
+	delay    time.Duration
+	stopChan <-chan struct{}
+}
+
+func newDelayedWriter(writer io.Writer, delay time.Duration, stopChan <-chan struct{}) delayedWriter {
+	return delayedWriter{
+		writer:   writer,
+		delay:    delay,
+		stopChan: stopChan,
+	}
+}
+
+func (w delayedWriter) Write(p []byte) (int, error) {
+	if w.delay == 0 {
+		return w.writer.Write(p)
+	}
+
+	timer := time.NewTimer(w.delay)
+	defer timer.Stop()
+	select {
+	case <-timer.C:
+		return w.writer.Write(p)
+	case <-w.stopChan:
+		return 0, clientCanceled
+	}
+}

main.go

@@ -4,6 +4,8 @@ import (
 	"flag"
 	"fmt"
 	"os"
+
+	"github.com/fatih/color"
 )
 
 var settings Settings
@@ -13,21 +15,22 @@ func main() {
 		localPort = flag.Int("p", 0, "Local port to listen on")
 		localHost = flag.String("l", "localhost", "Local address to listen on")
 		remote    = flag.String("r", "", "Remote address (host:port) to connect")
+		delay     = flag.Duration("d", 0, "the delay to relay packets")
 		protocol  = flag.String("t", "", "The type of protocol, currently support grpc")
 		silent    = flag.Bool("silent", false, "Only prints connection open/close and stats, default false")
 	)
 
 	flag.Parse()
-	saveSettings(*localHost, *localPort, *remote, *protocol, *silent)
+	saveSettings(*localHost, *localPort, *remote, *delay, *protocol, *silent)
 
 	if settings.RemoteHost == "" {
-		fmt.Fprintln(os.Stderr, "[x] Remote host required")
+		fmt.Fprintln(os.Stderr, color.HiRedString("[x] Remote host required"))
 		flag.PrintDefaults()
 		os.Exit(1)
 	}
 
 	if err := startListener(); err != nil {
-		fmt.Fprintf(os.Stderr, "[x] Failed to start listener: %v\n", err)
+		fmt.Fprintln(os.Stderr, color.HiRedString("[x] Failed to start listener: %v", err))
 		os.Exit(1)
 	}
 }

readme-cn.md

@@ -37,16 +37,18 @@ $ docker run --rm -it -p <listen-port>:<listen-port> -p <remote-port>:<remote-po
 ```shell
 $ tproxy --help
 Usage of tproxy:
+  -d duration
+    	the delay to relay packets
   -l string
-        Local address to listen on (default "localhost")
+    	Local address to listen on (default "localhost")
   -p int
-        Local port to listen on
+    	Local port to listen on
   -r string
-        Remote address (host:port) to connect
+    	Remote address (host:port) to connect
   -silent
-        Only prints connection open/close and stats, default false
+    	Only prints connection open/close and stats, default false
   -t string
-        The type of protocol, currently support grpc
+    	The type of protocol, currently support grpc
 ```
 
 ## 示例

readme.md

@@ -37,16 +37,18 @@ $ docker run --rm -it -p <listen-port>:<listen-port> -p <remote-port>:<remote-po
 ```shell
 $ tproxy --help
 Usage of tproxy:
+  -d duration
+    	the delay to relay packets
   -l string
-        Local address to listen on (default "localhost")
+    	Local address to listen on (default "localhost")
   -p int
-        Local port to listen on
+    	Local port to listen on
   -r string
-        Remote address (host:port) to connect
+    	Remote address (host:port) to connect
   -silent
-        Only prints connection open/close and stats, default false
+    	Only prints connection open/close and stats, default false
   -t string
-        The type of protocol, currently support grpc
+    	The type of protocol, currently support grpc
 ```
 
 ## Examples

settings.go

@@ -1,14 +1,18 @@
 package main
 
+import "time"
+
 type Settings struct {
 	RemoteHost string
 	LocalHost  string
 	LocalPort  int
+	Delay      time.Duration
 	Protocol   string
 	Silent     bool
 }
 
-func saveSettings(localHost string, localPort int, remoteHost, protocol string, silent bool) {
+func saveSettings(localHost string, localPort int, remoteHost string, delay time.Duration,
+	protocol string, silent bool) {
 	if localHost != "" {
 		settings.LocalHost = localHost
 	}
@@ -18,6 +22,7 @@ func saveSettings(localHost string, localPort int, remoteHost, protocol string,
 	if remoteHost != "" {
 		settings.RemoteHost = remoteHost
 	}
+	settings.Delay = delay
 	settings.Protocol = protocol
 	settings.Silent = silent
 }