b612/star
1
0
Fork 0
Code Issues Pull Requests Projects Releases 3 Wiki Activity

add hook for net forward

Browse Source
This commit is contained in:
兔子 2024-03-22 14:37:42 +08:00
parent 68f1eef7a6
commit 02f79f2eb0
2 changed files with 102 additions and 9 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
netforward/cmd.go
View File

@ -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")

109
netforward/forward.go
View File

@ -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()
}