mirror of
https://github.com/kevwan/tproxy.git
synced 2026-05-23 07:40:35 +00:00
feat: add up and down speed limits (#54)
Co-authored-by: Jianglong <youjianglong@tianzongyouxi.com>
This commit is contained in:
@@ -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() {
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
@@ -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
|
||||
}
|
||||
|
||||
@@ -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"))
|
||||
|
||||
Reference in New Issue
Block a user