@@ -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"
5051 runningPort int
5152 domainListPath string
5253 blacklistPath string
54+ bandwidthLimit int
5355)
5456
57+ var BandwidthLimiter * R.Bucket
58+
5559var Blacklist []RepoInfo
5660
5761var AcceptDomain = []string {
@@ -76,6 +80,7 @@ func init() {
7680 command .PersistentFlags ().IntVarP (& runningPort , "running-port" , "p" , 30000 , "disable color output" )
7781 command .PersistentFlags ().StringVarP (& domainListPath , "domain-list-path" , "d" , "domainlist.txt" , "set accept domain" )
7882 command .PersistentFlags ().StringVarP (& blacklistPath , "blacklist-path" , "b" , "blacklist.txt" , "set repository blacklist" )
83+ command .PersistentFlags ().IntVarP (& bandwidthLimit , "bandwidth-limit" , "l" , 0 , "set total bandwidth limit (MB/s), 0 as no limit" )
7984}
8085
8186func main () {
@@ -101,6 +106,10 @@ func newError(msg string) *HTTPError {
101106}
102107
103108func run (* cobra.Command , []string ) {
109+ if bandwidthLimit > 0 {
110+ BandwidthLimiter = R .NewBucketWithRate (float64 (bandwidthLimit * 1024 * 1024 ), int64 (bandwidthLimit * 1024 * 1024 ))
111+ log .Info ("Bandwidth limit is set as " , bandwidthLimit , "MB/s" )
112+ }
104113 if watcher , err := loadDomainList (); err == nil {
105114 err = watcher .Start ()
106115 if err == nil {
@@ -466,6 +475,34 @@ func responseWithRedirect(URL *url.URL) http.Handler {
466475 })
467476}
468477
478+ var _ io.Reader = (* LimitReader )(nil )
479+
480+ type LimitReader struct {
481+ reader io.Reader
482+ bucket * R.Bucket
483+ }
484+
485+ func NewLimitReader (reader io.Reader , bucket * R.Bucket ) * LimitReader {
486+ return & LimitReader {
487+ reader : reader ,
488+ bucket : bucket ,
489+ }
490+ }
491+
492+ func (lr * LimitReader ) Read (p []byte ) (int , error ) {
493+ sliceLen := int64 (len (p ))
494+ available := lr .bucket .TakeAvailable (sliceLen )
495+ if available == 0 {
496+ return 0 , nil
497+ }
498+ if available == sliceLen {
499+ return lr .reader .Read (p )
500+ }
501+ temp := make ([]byte , available )
502+ defer copy (p , temp )
503+ return lr .reader .Read (temp )
504+ }
505+
469506func sendRequestWithURL (URL * url.URL ) http.Handler {
470507 return http .HandlerFunc (func (w http.ResponseWriter , r * http.Request ) {
471508 ctx := r .Context ()
@@ -511,7 +548,11 @@ func sendRequestWithURL(URL *url.URL) http.Handler {
511548 }
512549 }
513550 w .WriteHeader (response .StatusCode )
514- io .Copy (w , response .Body )
551+ if BandwidthLimiter != nil {
552+ io .Copy (w , NewLimitReader (response .Body , BandwidthLimiter ))
553+ } else {
554+ io .Copy (w , response .Body )
555+ }
515556 log .InfoContext (ctx , "Success proxy request: " , URL , " , method: " , request .Method , " , status: " , response .StatusCode )
516557 })
517558}
0 commit comments