mirror of
https://github.com/fatedier/frp.git
synced 2026-08-24 09:18:51 +00:00
limit: clamp bandwidth limiter burst (#5471)
This commit is contained in:
@@ -66,7 +66,7 @@ func NewProxy(
|
|||||||
var limiter *rate.Limiter
|
var limiter *rate.Limiter
|
||||||
limitBytes := pxyConf.GetBaseConfig().Transport.BandwidthLimit.Bytes()
|
limitBytes := pxyConf.GetBaseConfig().Transport.BandwidthLimit.Bytes()
|
||||||
if limitBytes > 0 && pxyConf.GetBaseConfig().Transport.BandwidthLimitMode == types.BandwidthLimitModeClient {
|
if limitBytes > 0 && pxyConf.GetBaseConfig().Transport.BandwidthLimitMode == types.BandwidthLimitModeClient {
|
||||||
limiter = rate.NewLimiter(rate.Limit(float64(limitBytes)), int(limitBytes))
|
limiter = limit.NewBandwidthLimiter(limitBytes)
|
||||||
}
|
}
|
||||||
|
|
||||||
baseProxy := BaseProxy{
|
baseProxy := BaseProxy{
|
||||||
|
|||||||
@@ -0,0 +1,37 @@
|
|||||||
|
// Copyright 2026 The frp Authors
|
||||||
|
//
|
||||||
|
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||||
|
// you may not use this file except in compliance with the License.
|
||||||
|
// You may obtain a copy of the License at
|
||||||
|
//
|
||||||
|
// http://www.apache.org/licenses/LICENSE-2.0
|
||||||
|
//
|
||||||
|
// Unless required by applicable law or agreed to in writing, software
|
||||||
|
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||||
|
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||||
|
// See the License for the specific language governing permissions and
|
||||||
|
// limitations under the License.
|
||||||
|
|
||||||
|
package limit
|
||||||
|
|
||||||
|
import (
|
||||||
|
"fmt"
|
||||||
|
|
||||||
|
"golang.org/x/time/rate"
|
||||||
|
)
|
||||||
|
|
||||||
|
// NewBandwidthLimiter creates a limiter whose rate preserves the configured
|
||||||
|
// byte limit while keeping the burst representable as an int on all targets.
|
||||||
|
func NewBandwidthLimiter(bytes int64) *rate.Limiter {
|
||||||
|
if bytes <= 0 {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
maxInt := int64(^uint(0) >> 1)
|
||||||
|
burst := min(bytes, maxInt)
|
||||||
|
return rate.NewLimiter(rate.Limit(float64(bytes)), int(burst))
|
||||||
|
}
|
||||||
|
|
||||||
|
func invalidBurstError(burst int) error {
|
||||||
|
return fmt.Errorf("invalid limiter burst: %d", burst)
|
||||||
|
}
|
||||||
@@ -0,0 +1,65 @@
|
|||||||
|
// Copyright 2026 The frp Authors
|
||||||
|
//
|
||||||
|
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||||
|
// you may not use this file except in compliance with the License.
|
||||||
|
// You may obtain a copy of the License at
|
||||||
|
//
|
||||||
|
// http://www.apache.org/licenses/LICENSE-2.0
|
||||||
|
//
|
||||||
|
// Unless required by applicable law or agreed to in writing, software
|
||||||
|
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||||
|
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||||
|
// See the License for the specific language governing permissions and
|
||||||
|
// limitations under the License.
|
||||||
|
|
||||||
|
package limit
|
||||||
|
|
||||||
|
import (
|
||||||
|
"bytes"
|
||||||
|
"strconv"
|
||||||
|
"strings"
|
||||||
|
"testing"
|
||||||
|
|
||||||
|
"github.com/stretchr/testify/require"
|
||||||
|
"golang.org/x/time/rate"
|
||||||
|
)
|
||||||
|
|
||||||
|
func TestNewBandwidthLimiterClampsBurstToTargetInt(t *testing.T) {
|
||||||
|
const bytesPerSecond = int64(1 << 31)
|
||||||
|
|
||||||
|
limiter := NewBandwidthLimiter(bytesPerSecond)
|
||||||
|
require.NotNil(t, limiter)
|
||||||
|
|
||||||
|
wantBurst := bytesPerSecond
|
||||||
|
maxInt := int64(^uint(0) >> 1)
|
||||||
|
if wantBurst > maxInt {
|
||||||
|
wantBurst = maxInt
|
||||||
|
}
|
||||||
|
require.Equal(t, int(wantBurst), limiter.Burst())
|
||||||
|
require.Equal(t, rate.Limit(float64(bytesPerSecond)), limiter.Limit())
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestNewBandwidthLimiterDisablesNonPositiveLimit(t *testing.T) {
|
||||||
|
require.Nil(t, NewBandwidthLimiter(0))
|
||||||
|
require.Nil(t, NewBandwidthLimiter(-1))
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestReaderAndWriterRejectInvalidBurst(t *testing.T) {
|
||||||
|
for _, burst := range []int{0, -1} {
|
||||||
|
t.Run("reader/"+strconv.Itoa(burst), func(t *testing.T) {
|
||||||
|
reader := NewReader(strings.NewReader("payload"), rate.NewLimiter(rate.Limit(1), burst))
|
||||||
|
n, err := reader.Read(make([]byte, 1))
|
||||||
|
require.Zero(t, n)
|
||||||
|
require.EqualError(t, err, "invalid limiter burst: "+strconv.Itoa(burst))
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("writer/"+strconv.Itoa(burst), func(t *testing.T) {
|
||||||
|
var dst bytes.Buffer
|
||||||
|
writer := NewWriter(&dst, rate.NewLimiter(rate.Limit(1), burst))
|
||||||
|
n, err := writer.Write([]byte("payload"))
|
||||||
|
require.Zero(t, n)
|
||||||
|
require.EqualError(t, err, "invalid limiter burst: "+strconv.Itoa(burst))
|
||||||
|
require.Empty(t, dst.Bytes())
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -35,6 +35,12 @@ func NewReader(r io.Reader, limiter *rate.Limiter) *Reader {
|
|||||||
|
|
||||||
func (r *Reader) Read(p []byte) (n int, err error) {
|
func (r *Reader) Read(p []byte) (n int, err error) {
|
||||||
b := r.limiter.Burst()
|
b := r.limiter.Burst()
|
||||||
|
if b <= 0 {
|
||||||
|
if len(p) == 0 {
|
||||||
|
return 0, nil
|
||||||
|
}
|
||||||
|
return 0, invalidBurstError(b)
|
||||||
|
}
|
||||||
if b < len(p) {
|
if b < len(p) {
|
||||||
p = p[:b]
|
p = p[:b]
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -34,8 +34,15 @@ func NewWriter(w io.Writer, limiter *rate.Limiter) *Writer {
|
|||||||
}
|
}
|
||||||
|
|
||||||
func (w *Writer) Write(p []byte) (n int, err error) {
|
func (w *Writer) Write(p []byte) (n int, err error) {
|
||||||
|
if len(p) == 0 {
|
||||||
|
return 0, nil
|
||||||
|
}
|
||||||
|
|
||||||
var nn int
|
var nn int
|
||||||
b := w.limiter.Burst()
|
b := w.limiter.Burst()
|
||||||
|
if b <= 0 {
|
||||||
|
return 0, invalidBurstError(b)
|
||||||
|
}
|
||||||
for {
|
for {
|
||||||
end := len(p)
|
end := len(p)
|
||||||
if end == 0 {
|
if end == 0 {
|
||||||
|
|||||||
@@ -536,7 +536,7 @@ func NewProxy(ctx context.Context, options *Options) (pxy Proxy, err error) {
|
|||||||
var limiter *rate.Limiter
|
var limiter *rate.Limiter
|
||||||
limitBytes := configurer.GetBaseConfig().Transport.BandwidthLimit.Bytes()
|
limitBytes := configurer.GetBaseConfig().Transport.BandwidthLimit.Bytes()
|
||||||
if limitBytes > 0 && configurer.GetBaseConfig().Transport.BandwidthLimitMode == types.BandwidthLimitModeServer {
|
if limitBytes > 0 && configurer.GetBaseConfig().Transport.BandwidthLimitMode == types.BandwidthLimitModeServer {
|
||||||
limiter = rate.NewLimiter(rate.Limit(float64(limitBytes)), int(limitBytes))
|
limiter = limit.NewBandwidthLimiter(limitBytes)
|
||||||
}
|
}
|
||||||
|
|
||||||
basePxy := BaseProxy{
|
basePxy := BaseProxy{
|
||||||
|
|||||||
Reference in New Issue
Block a user