feat: udp bandwidth limit support

Author: cyunrei

README.md

@@ -6,7 +6,7 @@ Portbridge is a user-space port-forwarding tool with cross-platform support.
 
 - Cross-Platform Support (Linux / Windows / Darwin)
 - TCP and UDP Forward Support
-- Bandwidth Limit Support for TCP
+- TCP and UDP Bandwidth Limit Support
 - Batch Port Forwarding Rules Support
 
 # Usage
@@ -17,15 +17,15 @@ Portbridge Options:
   -s, --source=             Source address and port to bind locally
   -d, --destination=        Destination address and port to connect remotely
   -p, --protocol=           Specify the source protocol type
-  -b, --bandwidth-limit=    TCP Bandwidth limit in KiB (default: 0)
+  -b, --bandwidth-limit=    Bandwidth limit in KiB (default: 0)
       --udp-buffer-size=    UDP data forwarding buffer size in bytes (default: 1024)
       --udp-timeout-second= UDP data forwarding time out in second (default: 5)
   -f, --rule-file=          Batch port forwarding file path
   -h, --help                Show help message
   -v, --version             Print the version number
 ```
 
-Example:
+Examples:
 
 - Access the Cloudflare DNS (ipv6) via 127.0.0.2:53 with 100 udp buffer size
 

cmd/options/options.go

@@ -4,7 +4,7 @@ type Options struct {
 	SourceAddr       string `short:"s" long:"source" description:"Source address and port to bind locally" required:"true"`
 	DestinationAddr  string `short:"d" long:"destination" description:"Destination address and port to connect remotely" required:"true"`
 	Protocol         string `short:"p" long:"protocol" description:"Specify the source protocol type" required:"true"`
-	BandwidthLimit   uint64 `short:"b" long:"bandwidth-limit" description:"TCP Bandwidth limit in KiB" default:"0"`
+	BandwidthLimit   uint64 `short:"b" long:"bandwidth-limit" description:"Bandwidth limit in KiB" default:"0"`
 	UDPBufferSize    uint64 `long:"udp-buffer-size" description:"UDP data forwarding buffer size in bytes" default:"1024"`
 	UDPTimeoutSecond uint64 `long:"udp-timeout-second" description:"UDP data forwarding time out in second" default:"5"`
 	RuleFile         string `short:"f" long:"rule-file" description:"Batch port forwarding file path"`

pkg/forward/udp.go

@@ -1,7 +1,8 @@
 package forward
 
 import (
-	"errors"
+	"context"
+	"golang.org/x/time/rate"
 	"net"
 	"time"
 )
@@ -21,7 +22,11 @@ func NewUDPDataForwarder() *UDPDataForwarder {
 }
 
 func (f *UDPDataForwarder) Forward(sourceConn, destinationConn net.Conn) error {
-	return f.ForwardWithNormal(sourceConn, destinationConn)
+	if f.BandwidthLimit != DefaultBandwidthLimit {
+		return f.ForwardWithTrafficControl(sourceConn, destinationConn)
+	} else {
+		return f.ForwardWithNormal(sourceConn, destinationConn)
+	}
 }
 
 func (f *UDPDataForwarder) ForwardWithNormal(sourceConn, destinationConn net.Conn) error {
@@ -47,8 +52,7 @@ func (f *UDPDataForwarder) ForwardWithNormal(sourceConn, destinationConn net.Con
 			destinationConnBuffer := make([]byte, f.BufferSize)
 			destinationConn.SetReadDeadline(time.Now().Add(f.DeadlineSecond * time.Second))
 			m, _, err := destinationUDPConn.ReadFromUDP(destinationConnBuffer)
-			var netErr net.Error
-			if errors.As(err, &netErr) && netErr.Timeout() {
+			if err != nil {
 				return
 			}
 
@@ -62,7 +66,52 @@ func (f *UDPDataForwarder) ForwardWithNormal(sourceConn, destinationConn net.Con
 }
 
 func (f *UDPDataForwarder) ForwardWithTrafficControl(sourceConn, destinationConn net.Conn) error {
-	return f.ForwardWithNormal(sourceConn, destinationConn)
+	sourceUDPConn, _ := sourceConn.(*net.UDPConn)
+	destinationUDPConn, _ := destinationConn.(*net.UDPConn)
+
+	limiter := rate.NewLimiter(rate.Limit(f.BandwidthLimit*1024/8), int(f.BandwidthLimit*1024/8))
+
+	sourceConnBuffer := make([]byte, f.BufferSize)
+	for {
+		sourceConn.SetReadDeadline(time.Now().Add(f.DeadlineSecond * time.Second))
+		n, sourceConnAddr, err := sourceUDPConn.ReadFromUDP(sourceConnBuffer)
+		if err != nil {
+			continue
+		}
+
+		data := make([]byte, n)
+		copy(data, sourceConnBuffer[:n])
+
+		go func(data []byte, sourceConnAddr *net.UDPAddr) {
+			err := limiter.WaitN(context.Background(), n)
+			if err != nil {
+				return
+			}
+
+			_, err = destinationConn.Write(data)
+			if err != nil {
+				return
+			}
+
+			destinationConnBuffer := make([]byte, f.BufferSize)
+			destinationConn.SetReadDeadline(time.Now().Add(f.DeadlineSecond * time.Second))
+			m, _, err := destinationUDPConn.ReadFromUDP(destinationConnBuffer)
+			if err != nil {
+				return
+			}
+
+			err = limiter.WaitN(context.Background(), m)
+			if err != nil {
+				return
+			}
+
+			_, err = sourceUDPConn.WriteToUDP(destinationConnBuffer[:m], sourceConnAddr)
+			if err != nil {
+				return
+			}
+
+		}(data, sourceConnAddr)
+	}
 }
 
 func (f *UDPDataForwarder) SetBandwidthLimit(bandwidthLimit uint64) *UDPDataForwarder {