From e80e5c5679c9341ba680fcff3fd445a6fda8be51 Mon Sep 17 00:00:00 2001
From: Starainrt <i@b612.me>
Date: Fri, 22 Mar 2024 14:45:00 +0800
Subject: [PATCH] add hook for net forward

---
 netforward/forward.go | 23 ++++++++++++++++-------
 1 file changed, 16 insertions(+), 7 deletions(-)

diff --git a/netforward/forward.go b/netforward/forward.go
index e8aabc8..b02e361 100644
--- a/netforward/forward.go
+++ b/netforward/forward.go
@@ -21,6 +21,7 @@ type NetForward struct {
 	EnableTCP   bool
 	EnableUDP   bool
 	DelayMilSec int
+	DelayToward int
 	StdinMode   bool
 	DialTimeout time.Duration
 	UDPTimeout  time.Duration
@@ -62,6 +63,14 @@ func (n *NetForward) Run() error {
 			case "setremote":
 				n.RemoteURI = cmds[2]
 				starlog.Noticef("Remote URI Set to %s\n", n.RemoteURI)
+			case "setdelaytoward":
+				tmp, err := strconv.Atoi(cmds[2])
+				if err != nil {
+					starlog.Errorln("Invalid Delay Toward Value", cmds[2])
+					continue
+				}
+				n.DelayToward = tmp
+				starlog.Noticef("Delay Toward Set to %d\n", n.DelayToward)
 			case "setdelay":
 				tmp, err := strconv.Atoi(cmds[2])
 				if err != nil {
@@ -138,7 +147,7 @@ func (n *NetForward) runTCP() error {
 				return
 			}
 			log.Infof("TCP Connect %s <==> %s\n", conn.RemoteAddr().String(), rmt.RemoteAddr().String())
-			Copy(rmt, conn, n.DelayMilSec)
+			n.copy(rmt, conn)
 			log.Noticef("TCP Connection Closed  %s <==> %s", conn.RemoteAddr().String(), n.RemoteURI)
 		}(conn)
 	}
@@ -250,7 +259,7 @@ func (n *NetForward) runUDP() error {
 				log.Infof("UDP Connect %s <==> %s\n", rmt.String(), n.RemoteURI)
 			}
 			mu.Unlock()
-			if n.DelayMilSec > 0 {
+			if n.DelayMilSec > 0 || (n.DelayToward == 0 || n.DelayToward == 1) {
 				time.Sleep(time.Millisecond * time.Duration(n.DelayMilSec))
 			}
 			_, err := addr.Write(data)
@@ -265,15 +274,15 @@ func (n *NetForward) runUDP() error {
 	}
 }
 
-func Copy(dst, src net.Conn, delay int) {
+func (n *NetForward) copy(dst, src net.Conn) {
 	var wg sync.WaitGroup
 	wg.Add(2)
 	go func() {
 		defer wg.Done()
 		bufsize := make([]byte, 32*1024)
 		for {
-			if delay > 0 {
-				time.Sleep(time.Millisecond * time.Duration(delay))
+			if n.DelayMilSec > 0 && (n.DelayToward == 0 || n.DelayToward == 1) {
+				time.Sleep(time.Millisecond * time.Duration(n.DelayMilSec))
 			}
 			n, err := src.Read(bufsize)
 			if err != nil {
@@ -293,8 +302,8 @@ func Copy(dst, src net.Conn, delay int) {
 		defer wg.Done()
 		bufsize := make([]byte, 32*1024)
 		for {
-			if delay > 0 {
-				time.Sleep(time.Millisecond * time.Duration(delay))
+			if n.DelayMilSec > 0 && (n.DelayToward == 0 || n.DelayToward == 2) {
+				time.Sleep(time.Millisecond * time.Duration(n.DelayMilSec))
 			}
 			n, err := dst.Read(bufsize)
 			if err != nil {