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

---
 netforward/cmd.go     |   2 +
 netforward/forward.go | 109 ++++++++++++++++++++++++++++++++++++++----
 2 files changed, 102 insertions(+), 9 deletions(-)

diff --git a/netforward/cmd.go b/netforward/cmd.go
index 92b7dd3..32d54f9 100644
--- a/netforward/cmd.go
+++ b/netforward/cmd.go
@@ -14,6 +14,8 @@ var f = new(NetForward)
 var dialTimeout, udpTimeout int64
 
 func init() {
+	CmdNetforward.Flags().BoolVarP(&f.StdinMode, "stdin", "s", false, "enable stdin mode")
+	CmdNetforward.Flags().IntVarP(&f.DelayMilSec, "delay", "D", 0, "delay milliseconds")
 	CmdNetforward.Flags().StringVarP(&f.LocalAddr, "local", "l", "0.0.0.0", "bind address")
 	CmdNetforward.Flags().IntVarP(&f.LocalPort, "port", "p", 11270, "local listen port")
 	CmdNetforward.Flags().BoolVarP(&f.EnableTCP, "enable-tcp-forward", "t", true, "enable tcp forward mode")
diff --git a/netforward/forward.go b/netforward/forward.go
index e8c3bb5..e8aabc8 100644
--- a/netforward/forward.go
+++ b/netforward/forward.go
@@ -1,12 +1,14 @@
 package netforward
 
 import (
+	"b612.me/stario"
 	"b612.me/starlog"
 	"context"
 	"errors"
 	"fmt"
-	"io"
 	"net"
+	"strconv"
+	"strings"
 	"sync"
 	"sync/atomic"
 	"time"
@@ -18,6 +20,8 @@ type NetForward struct {
 	RemoteURI   string
 	EnableTCP   bool
 	EnableUDP   bool
+	DelayMilSec int
+	StdinMode   bool
 	DialTimeout time.Duration
 	UDPTimeout  time.Duration
 	stopCtx     context.Context
@@ -42,6 +46,55 @@ func (n *NetForward) Run() error {
 	if n.DialTimeout == 0 {
 		n.DialTimeout = time.Second * 5
 	}
+	if n.StdinMode {
+	breakfor:
+		for {
+			cmd := strings.TrimSpace(stario.MessageBoxRaw("", "").MustString())
+			for strings.Contains(cmd, "  ") {
+				cmd = strings.Replace(cmd, "  ", " ", -1)
+			}
+			cmds := strings.Split(cmd, " ")
+			if len(cmds) < 3 {
+				starlog.Errorln("Invalid Command", cmd)
+				continue
+			}
+			switch cmds[0] + cmds[1] {
+			case "setremote":
+				n.RemoteURI = cmds[2]
+				starlog.Noticef("Remote URI Set to %s\n", n.RemoteURI)
+			case "setdelay":
+				tmp, err := strconv.Atoi(cmds[2])
+				if err != nil {
+					starlog.Errorln("Invalid Delay Value", cmds[2])
+					continue
+				}
+				n.DelayMilSec = tmp
+				starlog.Noticef("Delay Set to %d\n", n.DelayMilSec)
+			case "setdialtimeout":
+				tmp, err := strconv.Atoi(cmds[2])
+				if err != nil {
+					starlog.Errorln("Invalid Dial Timeout Value", cmds[2])
+					continue
+				}
+				n.DialTimeout = time.Millisecond * time.Duration(tmp)
+				starlog.Noticef("Dial Timeout Set to %d\n", n.DialTimeout)
+			case "setudptimeout":
+				tmp, err := strconv.Atoi(cmds[2])
+				if err != nil {
+					starlog.Errorln("Invalid UDP Timeout Value", cmds[2])
+					continue
+				}
+				n.UDPTimeout = time.Millisecond * time.Duration(tmp)
+				starlog.Noticef("UDP Timeout Set to %d\n", n.UDPTimeout)
+			case "setstdin":
+				if cmds[2] == "off" {
+					n.StdinMode = false
+					starlog.Noticef("Stdin Mode Off\n")
+					break breakfor
+				}
+			}
+		}
+	}
 	if n.EnableTCP {
 		go n.runTCP()
 	}
@@ -85,7 +138,7 @@ func (n *NetForward) runTCP() error {
 				return
 			}
 			log.Infof("TCP Connect %s <==> %s\n", conn.RemoteAddr().String(), rmt.RemoteAddr().String())
-			Copy(rmt, conn)
+			Copy(rmt, conn, n.DelayMilSec)
 			log.Noticef("TCP Connection Closed  %s <==> %s", conn.RemoteAddr().String(), n.RemoteURI)
 		}(conn)
 	}
@@ -108,9 +161,12 @@ func (u UDPConn) Read(p []byte) (n int, err error) {
 	return u.Conn.Read(p)
 }
 
-func (u UDPConn) Work() {
+func (u UDPConn) Work(delay int) {
 	buf := make([]byte, 8192)
 	for {
+		if delay > 0 {
+			time.Sleep(time.Millisecond * time.Duration(delay))
+		}
 		count, err := u.Read(buf)
 		if err != nil {
 			u.Close()
@@ -190,10 +246,13 @@ func (n *NetForward) runUDP() error {
 					lastbeat:   time.Now().Unix(),
 				}
 				udpMap[rmt.String()] = addr
-				go addr.Work()
+				go addr.Work(n.DelayMilSec)
 				log.Infof("UDP Connect %s <==> %s\n", rmt.String(), n.RemoteURI)
 			}
 			mu.Unlock()
+			if n.DelayMilSec > 0 {
+				time.Sleep(time.Millisecond * time.Duration(n.DelayMilSec))
+			}
 			_, err := addr.Write(data)
 			if err != nil {
 				mu.Lock()
@@ -206,18 +265,50 @@ func (n *NetForward) runUDP() error {
 	}
 }
 
-func Copy(dst, src net.Conn) {
+func Copy(dst, src net.Conn, delay int) {
 	var wg sync.WaitGroup
 	wg.Add(2)
 	go func() {
 		defer wg.Done()
-		io.Copy(dst, src)
+		bufsize := make([]byte, 32*1024)
+		for {
+			if delay > 0 {
+				time.Sleep(time.Millisecond * time.Duration(delay))
+			}
+			n, err := src.Read(bufsize)
+			if err != nil {
+				dst.Close()
+				src.Close()
+				return
+			}
+			_, err = dst.Write(bufsize[:n])
+			if err != nil {
+				src.Close()
+				dst.Close()
+				return
+			}
+		}
 	}()
 	go func() {
 		defer wg.Done()
-		io.Copy(src, dst)
+		bufsize := make([]byte, 32*1024)
+		for {
+			if delay > 0 {
+				time.Sleep(time.Millisecond * time.Duration(delay))
+			}
+			n, err := dst.Read(bufsize)
+			if err != nil {
+				src.Close()
+				dst.Close()
+				return
+			}
+			_, err = src.Write(bufsize[:n])
+			if err != nil {
+				src.Close()
+				dst.Close()
+				return
+			}
+		}
 	}()
 	wg.Wait()
-	dst.Close()
-	src.Close()
 }