feat: add up and down speed limits (#54)

Co-authored-by: Jianglong <youjianglong@tianzongyouxi.com>
This commit is contained in:
Jianglong
2023-09-14 02:35:16 +08:00
committed by GitHub
parent 95becf6093
commit e79f31db71
5 changed files with 24 additions and 4 deletions
+13 -2
View File
@@ -10,6 +10,7 @@ import (
"time"
"github.com/fatih/color"
"github.com/juju/ratelimit"
"github.com/kevwan/tproxy/display"
"github.com/kevwan/tproxy/protocol"
)
@@ -58,7 +59,12 @@ func (c *PairedConnection) handleClientMessage() {
r, w := io.Pipe()
tee := io.MultiWriter(c.svrConn, w)
go protocol.CreateInterop(settings.Protocol).Dump(r, protocol.ClientSide, c.id, settings.Quiet)
c.copyData(tee, c.cliConn, protocol.ClientSide)
var src io.Reader = c.cliConn
if settings.UpLimit > 0 {
bucket := ratelimit.NewBucket(time.Second, settings.UpLimit)
src = ratelimit.Reader(src, bucket)
}
c.copyData(tee, src, protocol.ClientSide)
}
func (c *PairedConnection) handleServerMessage() {
@@ -68,7 +74,12 @@ func (c *PairedConnection) handleServerMessage() {
r, w := io.Pipe()
tee := io.MultiWriter(newDelayedWriter(c.cliConn, settings.Delay, c.stopChan), w)
go protocol.CreateInterop(settings.Protocol).Dump(r, protocol.ServerSide, c.id, settings.Quiet)
c.copyData(tee, c.svrConn, protocol.ServerSide)
var src io.Reader = c.svrConn
if settings.DownLimit > 0 {
bucket := ratelimit.NewBucket(time.Second, settings.DownLimit)
src = ratelimit.Reader(src, bucket)
}
c.copyData(tee, src, protocol.ServerSide)
}
func (c *PairedConnection) process() {
+1
View File
@@ -11,6 +11,7 @@ require (
)
require (
github.com/juju/ratelimit v1.0.2 // indirect
github.com/mattn/go-colorable v0.1.13 // indirect
github.com/mattn/go-isatty v0.0.19 // indirect
github.com/mattn/go-runewidth v0.0.14 // indirect
+2
View File
@@ -7,6 +7,8 @@ github.com/golang/snappy v0.0.1/go.mod h1:/XxbfmMg8lxefKM7IXC3fBNl/7bRcc72aCRzEW
github.com/google/go-cmp v0.5.2/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE=
github.com/google/go-cmp v0.5.5 h1:Khx7svrCpmxxtHBq5j2mp/xVjsi8hQMfNLvJFAlrGgU=
github.com/google/go-cmp v0.5.5/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE=
github.com/juju/ratelimit v1.0.2 h1:sRxmtRiajbvrcLQT7S+JbqU0ntsb9W2yhSdNN8tWfaI=
github.com/juju/ratelimit v1.0.2/go.mod h1:qapgC/Gy+xNh9UxzV13HGGl/6UXNN+ct+vwSgWNm/qk=
github.com/klauspost/compress v1.13.6/go.mod h1:/3/Vjq9QcHkK5uEr5lBEmyoZ1iFhe47etQ6QUkpK6sk=
github.com/mattn/go-colorable v0.1.13 h1:fFA4WZxdEF4tXPZVKMLwD8oUnCTTo08duU7wxecdEvA=
github.com/mattn/go-colorable v0.1.13/go.mod h1:7S9/ev0klgBDR4GtXTXX8a3vIGJpMovkB8vQcUbaXHg=
+5 -1
View File
@@ -10,10 +10,12 @@ type Settings struct {
Protocol string
Stat bool
Quiet bool
UpLimit int64
DownLimit int64
}
func saveSettings(localHost string, localPort int, remote string, delay time.Duration,
protocol string, stat, quiet bool) {
protocol string, stat, quiet bool, upLimit, downLimit int64) {
if localHost != "" {
settings.LocalHost = localHost
}
@@ -27,4 +29,6 @@ func saveSettings(localHost string, localPort int, remote string, delay time.Dur
settings.Protocol = protocol
settings.Stat = stat
settings.Quiet = quiet
settings.UpLimit = upLimit
settings.DownLimit = downLimit
}
+3 -1
View File
@@ -20,6 +20,8 @@ func main() {
stat = flag.Bool("s", false, "Enable statistics")
quiet = flag.Bool("q", false,
"Quiet mode, only prints connection open/close and stats, default false")
upLimit = flag.Int64("U", 0, "Upward speed limit(Bytes/second)")
downLimit = flag.Int64("D", 0, "Downward speed limit(Bytes/second)")
)
if len(os.Args) <= 1 {
@@ -28,7 +30,7 @@ func main() {
}
flag.Parse()
saveSettings(*localHost, *localPort, *remote, *delay, *protocol, *stat, *quiet)
saveSettings(*localHost, *localPort, *remote, *delay, *protocol, *stat, *quiet, *upLimit, *downLimit)
if len(settings.Remote) == 0 {
fmt.Fprintln(os.Stderr, color.HiRedString("[x] Remote target required"))