@@ -18,6 +18,7 @@ import (
1818 "github.com/go-chi/chi/v5"
1919 "github.com/go-chi/chi/v5/middleware"
2020 "github.com/go-chi/render"
21+ R "github.com/juju/ratelimit"
2122 "github.com/sagernet/fswatch"
2223 L "github.com/sagernet/sing-box/log"
2324 "github.com/sagernet/sing/common"
4445 runningPort int
4546 domainListPath string
4647 blacklistPath string
48+ bandwidthLimit int
4749)
4850
51+ var BandwidthLimiter * R.Bucket
52+
4953var Blacklist []RepoInfo
5054
5155var AcceptDomain = []string {
@@ -70,6 +74,7 @@ func init() {
7074 command .PersistentFlags ().IntVarP (& runningPort , "running-port" , "p" , 30000 , "disable color output" )
7175 command .PersistentFlags ().StringVarP (& domainListPath , "domain-list-path" , "d" , "domainlist.txt" , "set accept domain" )
7276 command .PersistentFlags ().StringVarP (& blacklistPath , "blacklist-path" , "b" , "blacklist.txt" , "set repository blacklist" )
77+ command .PersistentFlags ().IntVarP (& bandwidthLimit , "bandwidth-limit" , "l" , 0 , "set total bandwidth limit (MB/s), 0 as no limit" )
7378}
7479
7580func main () {
@@ -95,6 +100,10 @@ func newError(msg string) *HTTPError {
95100}
96101
97102func run (* cobra.Command , []string ) {
103+ if bandwidthLimit > 0 {
104+ BandwidthLimiter = R .NewBucketWithRate (float64 (bandwidthLimit * 1024 * 1024 ), int64 (bandwidthLimit * 1024 * 1024 ))
105+ log .Info ("Bandwidth limit is set as " , bandwidthLimit , "MB/s" )
106+ }
98107 if watcher , err := loadDomainList (); err == nil {
99108 err = watcher .Start ()
100109 if err == nil {
@@ -460,6 +469,34 @@ func responseWithRedirect(URL *url.URL) http.Handler {
460469 })
461470}
462471
472+ var _ io.Reader = (* LimitReader )(nil )
473+
474+ type LimitReader struct {
475+ reader io.Reader
476+ bucket * R.Bucket
477+ }
478+
479+ func NewLimitReader (reader io.Reader , bucket * R.Bucket ) * LimitReader {
480+ return & LimitReader {
481+ reader : reader ,
482+ bucket : bucket ,
483+ }
484+ }
485+
486+ func (lr * LimitReader ) Read (p []byte ) (int , error ) {
487+ sliceLen := int64 (len (p ))
488+ available := lr .bucket .TakeAvailable (sliceLen )
489+ if available == 0 {
490+ return 0 , nil
491+ }
492+ if available == sliceLen {
493+ return lr .reader .Read (p )
494+ }
495+ temp := make ([]byte , available )
496+ defer copy (p , temp )
497+ return lr .reader .Read (temp )
498+ }
499+
463500func sendRequestWithURL (URL * url.URL ) http.Handler {
464501 return http .HandlerFunc (func (w http.ResponseWriter , r * http.Request ) {
465502 ctx := r .Context ()
@@ -495,7 +532,11 @@ func sendRequestWithURL(URL *url.URL) http.Handler {
495532 }
496533 }
497534 w .WriteHeader (response .StatusCode )
498- io .Copy (w , response .Body )
535+ if BandwidthLimiter != nil {
536+ io .Copy (w , NewLimitReader (response .Body , BandwidthLimiter ))
537+ } else {
538+ io .Copy (w , response .Body )
539+ }
499540 log .InfoContext (ctx , "Success proxy request: " , URL , " , method: " , request .Method , " , status: " , response .StatusCode )
500541 })
501542}
0 commit comments