add hook for net forward

master
兔子 8 months ago
parent 02f79f2eb0
commit e80e5c5679

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

Loading…
Cancel
Save