BW limit parameters by JS script

Author: Snawoot

forward/bwlimit.go

@@ -3,18 +3,86 @@ package forward
 import (
 	"context"
 	"errors"
+	"fmt"
 	"io"
 	"sync"
 	"time"
 
 	"github.com/Snawoot/secache"
 	"github.com/Snawoot/secache/randmap"
 
+	clog "github.com/SenseUnit/dumbproxy/log"
 	"github.com/SenseUnit/dumbproxy/rate"
 )
 
 const copyChunkSize = 128 * 1024
 
+type LimitKind int
+
+const (
+	LimitKindNone LimitKind = iota
+	LimitKindStatic
+	LimitKindJS
+)
+
+type LimitSpec struct {
+	Kind LimitKind
+	Spec any
+}
+
+type StaticLimitSpec struct {
+	BPS      uint64
+	Burst    int64
+	Separate bool
+}
+
+type JSLimitSpec struct {
+	Filename  string
+	Instances int
+}
+
+type LimitParameters struct {
+	UploadBPS     float64 `json:"uploadBPS"`
+	UploadBurst   int64   `json:"uploadBurst"`
+	DownloadBPS   float64 `json:"downloadBPS"`
+	DownloadBurst int64   `json:"downloadBurst"`
+	GroupKey      *string `json:"groupKey"`
+	Separate      bool    `json:"separate"`
+}
+
+type LimitProvider = func(context.Context, string, string, string) (*LimitParameters, error)
+
+func ProviderFromSpec(spec LimitSpec, logger *clog.CondLogger) (LimitProvider, error) {
+	switch spec.Kind {
+	case LimitKindStatic:
+		staticSpec, ok := spec.Spec.(StaticLimitSpec)
+		if !ok {
+			return nil, fmt.Errorf("incorrect payload type in BW limit spec: %T", spec)
+		}
+		return func(_ context.Context, username, _, _ string) (*LimitParameters, error) {
+			return &LimitParameters{
+				UploadBPS:     float64(staticSpec.BPS),
+				UploadBurst:   staticSpec.Burst,
+				DownloadBPS:   float64(staticSpec.BPS),
+				DownloadBurst: staticSpec.Burst,
+				GroupKey:      &username,
+				Separate:      staticSpec.Separate,
+			}, nil
+		}, nil
+	case LimitKindJS:
+		jsSpec, ok := spec.Spec.(JSLimitSpec)
+		if !ok {
+			return nil, fmt.Errorf("incorrect payload type in BW limit spec: %T", spec)
+		}
+		j, err := NewJSLimitProvider(jsSpec.Filename, jsSpec.Instances, logger)
+		if err != nil {
+			return nil, err
+		}
+		return j.Parameters, nil
+	}
+	return nil, fmt.Errorf("unsupported BW limit kind %d", int(spec.Kind))
+}
+
 type cacheItem struct {
 	mux sync.RWMutex
 	ul  *rate.Limiter
@@ -38,17 +106,13 @@ func (i *cacheItem) unlock() {
 }
 
 type BWLimit struct {
-	bps      float64
-	burst    int64
-	separate bool
-	cache    secache.Cache[string, *cacheItem]
+	paramFn LimitProvider
+	cache   secache.Cache[string, *cacheItem]
 }
 
-func NewBWLimit(bytesPerSecond float64, burst int64, separate bool) *BWLimit {
+func NewBWLimit(p LimitProvider) *BWLimit {
 	return &BWLimit{
-		bps:      bytesPerSecond,
-		burst:    burst,
-		separate: separate,
+		paramFn: p,
 		cache: *(secache.New[string, *cacheItem](3, func(_ string, item *cacheItem) bool {
 			if item.tryLock() {
 				if item.ul.Tokens() >= float64(item.ul.Burst()) && item.dl.Tokens() >= float64(item.dl.Burst()) {
@@ -120,35 +184,46 @@ func (l *BWLimit) futureCopyAndCloseWrite(ctx context.Context, c chan<- error, r
 	close(c)
 }
 
-func (l *BWLimit) getRatelimiters(username string) (res *cacheItem) {
+func (l *BWLimit) getRatelimiters(ctx context.Context, username, network, address string) (*cacheItem, error) {
+	params, err := l.paramFn(ctx, username, network, address)
+	if err != nil {
+		return nil, err
+	}
+	groupKey := username
+	if params.GroupKey != nil {
+		groupKey = *params.GroupKey
+	}
+	var res *cacheItem
 	l.cache.Do(func(m *randmap.RandMap[string, *cacheItem]) {
 		var ok bool
-		res, ok = m.Get(username)
+		res, ok = m.Get(groupKey)
 		if ok {
 			res.rLock()
 		} else {
-			ul := rate.NewLimiter(rate.Limit(l.bps), max(copyChunkSize, l.burst))
+			ul := rate.NewLimiter(rate.Limit(params.UploadBPS), max(copyChunkSize, params.UploadBurst))
 			dl := ul
-			if l.separate {
-				dl = rate.NewLimiter(rate.Limit(l.bps), max(copyChunkSize, l.burst))
+			if params.Separate {
+				dl = rate.NewLimiter(rate.Limit(params.DownloadBPS), max(copyChunkSize, params.DownloadBurst))
 			}
 			res = &cacheItem{
 				ul: ul,
 				dl: dl,
 			}
 			res.rLock()
-			l.cache.SetLocked(m, username, res)
+			l.cache.SetLocked(m, groupKey, res)
 		}
 		return
 	})
-	return
+	return res, nil
 }
 
-func (l *BWLimit) PairConnections(ctx context.Context, username string, incoming, outgoing io.ReadWriteCloser) error {
-	ci := l.getRatelimiters(username)
+func (l *BWLimit) PairConnections(ctx context.Context, username string, incoming, outgoing io.ReadWriteCloser, network, address string) error {
+	ci, err := l.getRatelimiters(ctx, username, network, address)
+	if err != nil {
+		return fmt.Errorf("ratelimit parameter computarion failed for user %q: %w", username, err)
+	}
 	defer ci.rUnlock()
 
-	var err error
 	i2oErr := make(chan error, 1)
 	o2iErr := make(chan error, 1)
 	ctxErr := ctx.Done()

forward/direct.go

@@ -22,7 +22,7 @@ func futureCopyAndCloseWrite(c chan<- error, dst io.WriteCloser, src io.ReadClos
 	close(c)
 }
 
-func PairConnections(ctx context.Context, username string, incoming, outgoing io.ReadWriteCloser) error {
+func PairConnections(ctx context.Context, username string, incoming, outgoing io.ReadWriteCloser, _, _ string) error {
 	var err error
 	i2oErr := make(chan error, 1)
 	o2iErr := make(chan error, 1)

forward/jslimit.go

@@ -0,0 +1,107 @@
+package forward
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"os"
+
+	"github.com/dop251/goja"
+	"golang.org/x/sync/errgroup"
+
+	"github.com/SenseUnit/dumbproxy/dialer/dto"
+	"github.com/SenseUnit/dumbproxy/jsext"
+	clog "github.com/SenseUnit/dumbproxy/log"
+)
+
+type JSLimitFunc = func(req *jsext.JSRequestInfo, dst *jsext.JSDstInfo, username string) (*LimitParameters, error)
+
+type JSLimitProvider struct {
+	funcPool chan JSLimitFunc
+	logger   *clog.CondLogger
+}
+
+func NewJSLimitProvider(filename string, instances int, logger *clog.CondLogger) (*JSLimitProvider, error) {
+	script, err := os.ReadFile(filename)
+	if err != nil {
+		return nil, fmt.Errorf("unable to load JS script file %q: %w", filename, err)
+	}
+
+	instances = max(1, instances)
+	pool := make(chan JSLimitFunc, instances)
+	initGroup, _ := errgroup.WithContext(context.Background())
+
+	for i := 0; i < instances; i++ {
+		initGroup.Go(func() error {
+			vm := goja.New()
+			err := jsext.AddPrinter(vm, logger)
+			if err != nil {
+				return fmt.Errorf("can't add print function to runtime: %w", err)
+			}
+			err = jsext.ConfigureRuntime(vm)
+			if err != nil {
+				return fmt.Errorf("can't configure runtime: %w", err)
+			}
+			_, err = vm.RunString(string(script))
+			if err != nil {
+				return fmt.Errorf("script run failed: %w", err)
+			}
+
+			var f JSLimitFunc
+			var limitFnJSVal goja.Value
+			if ex := vm.Try(func() {
+				limitFnJSVal = vm.Get("bwLimit")
+			}); ex != nil {
+				return fmt.Errorf("\"bwLimit\" function cannot be located in VM context: %w", err)
+			}
+			if limitFnJSVal == nil {
+				return errors.New("\"bwLimit\" function is not defined")
+			}
+			err = vm.ExportTo(limitFnJSVal, &f)
+			if err != nil {
+				return fmt.Errorf("can't export \"bwLimit\" function from JS VM: %w", err)
+			}
+
+			pool <- f
+			return nil
+		})
+	}
+
+	err = initGroup.Wait()
+	if err != nil {
+		return nil, err
+	}
+
+	return &JSLimitProvider{
+		funcPool: pool,
+		logger:   logger,
+	}, nil
+}
+
+func (j *JSLimitProvider) Parameters(ctx context.Context, username, network, address string) (res *LimitParameters, err error) {
+	defer func() {
+		if err != nil {
+			j.logger.Error("%v", err)
+		}
+	}()
+	req, _ := dto.FilterParamsFromContext(ctx)
+	ri := jsext.JSRequestInfoFromRequest(req)
+	di, err := jsext.JSDstInfoFromContext(ctx, network, address)
+	if err != nil {
+		return nil, fmt.Errorf("unable to construct dst info: %w", err)
+	}
+	func() {
+		f := <-j.funcPool
+		defer func(pool chan JSLimitFunc, f JSLimitFunc) {
+			pool <- f
+		}(j.funcPool, f)
+		res, err = f(ri, di, username)
+	}()
+	if err != nil {
+		return nil, fmt.Errorf("JS limit script exception: %w", err)
+	}
+	if res == nil {
+		return nil, fmt.Errorf("JS limit script returned null object")
+	}
+	return res, nil
+}

handler/config.go

@@ -1,9 +1,6 @@
 package handler
 
 import (
-	"context"
-	"io"
-
 	"github.com/SenseUnit/dumbproxy/auth"
 	clog "github.com/SenseUnit/dumbproxy/log"
 )
@@ -19,7 +16,7 @@ type Config struct {
 	Logger *clog.CondLogger
 	// Forward optionally specifies custom connection pairing function
 	// which does actual data forwarding.
-	Forward func(ctx context.Context, username string, incoming, outgoing io.ReadWriteCloser) error
+	Forward ForwardFunc
 	// UserIPHints specifies whether allow IP hints set by user or not
 	UserIPHints bool
 }

handler/handler.go

@@ -10,7 +10,9 @@ import (
 	"math/rand/v2"
 	"net"
 	"net/http"
+	"net/url"
 	"strconv"
+	"strings"
 	"sync"
 
 	"github.com/SenseUnit/dumbproxy/auth"
@@ -27,7 +29,7 @@ type HandlerDialer interface {
 	DialContext(ctx context.Context, net, address string) (net.Conn, error)
 }
 
-type ForwardFunc = func(ctx context.Context, username string, incoming, outgoing io.ReadWriteCloser) error
+type ForwardFunc = func(ctx context.Context, username string, incoming, outgoing io.ReadWriteCloser, network, address string) error
 
 type ProxyHandler struct {
 	auth          auth.Auth
@@ -114,6 +116,8 @@ func (s *ProxyHandler) HandleTunnel(wr http.ResponseWriter, req *http.Request, u
 				username,
 				wrapH1ReqBody(io.NopCloser(io.LimitReader(rw.Reader, int64(buffered)))),
 				wrapH1RespWriter(conn),
+				"tcp",
+				req.RequestURI,
 			)
 			s.forward(
 				req.Context(),
@@ -123,31 +127,48 @@ func (s *ProxyHandler) HandleTunnel(wr http.ResponseWriter, req *http.Request, u
 					localconn,
 				),
 				conn,
+				"tcp",
+				req.RequestURI,
 			)
 		} else {
 			s.logger.Debug("not rescuing remaining data in bufio.ReadWriter")
 			fmt.Fprintf(localconn, "HTTP/%d.%d 200 OK\r\n\r\n", req.ProtoMajor, req.ProtoMinor)
-			s.forward(req.Context(), username, localconn, conn)
+			s.forward(req.Context(), username, localconn, conn, "tcp", req.RequestURI)
 		}
 	} else if req.ProtoMajor == 2 {
 		wr.Header()["Date"] = nil
 		wr.WriteHeader(http.StatusOK)
 		flush(wr)
-		s.forward(req.Context(), username, wrapH2(req.Body, wr), conn)
+		s.forward(req.Context(), username, wrapH2(req.Body, wr), conn, "tcp", req.RequestURI)
 	} else {
 		s.logger.Error("Unsupported protocol version: %s", req.Proto)
 		http.Error(wr, "Unsupported protocol version.", http.StatusBadRequest)
 		return
 	}
 }
 
+func addressFromURL(u *url.URL) string {
+	host := u.Hostname()
+	port := u.Port()
+	if port == "" {
+		switch strings.ToLower(u.Scheme) {
+		case "http":
+			port = "80"
+		case "https":
+			port = "443"
+		}
+	}
+	return net.JoinHostPort(host, port)
+}
+
 func (s *ProxyHandler) HandleRequest(wr http.ResponseWriter, req *http.Request, username string) {
 	req.RequestURI = ""
 	forwardReqBody := newH1ReqBodyPipe()
 	origBody := req.Body
 	req.Body = forwardReqBody.Body()
+	address := addressFromURL(req.URL)
 	go func() {
-		s.forward(req.Context(), username, wrapH1ReqBody(origBody), forwardReqBody)
+		s.forward(req.Context(), username, wrapH1ReqBody(origBody), forwardReqBody, "tcp", address)
 	}()
 	if req.ProtoMajor == 2 {
 		req.URL.Scheme = "http" // We can't access :scheme pseudo-header, so assume http
@@ -171,7 +192,7 @@ func (s *ProxyHandler) HandleRequest(wr http.ResponseWriter, req *http.Request,
 	copyHeader(wr.Header(), resp.Header)
 	wr.WriteHeader(resp.StatusCode)
 	flush(wr)
-	s.forward(req.Context(), username, wrapH1RespWriter(wr), wrapH1ReqBody(resp.Body))
+	s.forward(req.Context(), username, wrapH1RespWriter(wr), wrapH1ReqBody(resp.Body), "tcp", address)
 }
 
 func (s *ProxyHandler) HandleGetRandom(wr http.ResponseWriter, req *http.Request, username string) {

handler/socks.go

@@ -70,7 +70,7 @@ func SOCKSHandler(dialer HandlerDialer, logger *clog.CondLogger, forward Forward
 			return fmt.Errorf("failed to send reply, %v", err)
 		}
 
-		return forward(ctx, username, wrapSOCKS(request.Reader, writer), target)
+		return forward(ctx, username, wrapSOCKS(request.Reader, writer), target, "tcp", request.DestAddr.String())
 	}
 }