From f666d97b64d6902ce15dcd329087cb1aadaf76d3 Mon Sep 17 00:00:00 2001 From: fatedier Date: Sun, 9 Aug 2026 17:30:39 +0800 Subject: [PATCH] limit: clamp bandwidth limiter burst (#5471) --- client/proxy/proxy.go | 2 +- pkg/util/limit/limiter.go | 37 +++++++++++++++++++ pkg/util/limit/limiter_test.go | 65 ++++++++++++++++++++++++++++++++++ pkg/util/limit/reader.go | 6 ++++ pkg/util/limit/writer.go | 7 ++++ server/proxy/proxy.go | 2 +- 6 files changed, 117 insertions(+), 2 deletions(-) create mode 100644 pkg/util/limit/limiter.go create mode 100644 pkg/util/limit/limiter_test.go diff --git a/client/proxy/proxy.go b/client/proxy/proxy.go index 5af56599..ece2f935 100644 --- a/client/proxy/proxy.go +++ b/client/proxy/proxy.go @@ -66,7 +66,7 @@ func NewProxy( var limiter *rate.Limiter limitBytes := pxyConf.GetBaseConfig().Transport.BandwidthLimit.Bytes() if limitBytes > 0 && pxyConf.GetBaseConfig().Transport.BandwidthLimitMode == types.BandwidthLimitModeClient { - limiter = rate.NewLimiter(rate.Limit(float64(limitBytes)), int(limitBytes)) + limiter = limit.NewBandwidthLimiter(limitBytes) } baseProxy := BaseProxy{ diff --git a/pkg/util/limit/limiter.go b/pkg/util/limit/limiter.go new file mode 100644 index 00000000..43ba7434 --- /dev/null +++ b/pkg/util/limit/limiter.go @@ -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) +} diff --git a/pkg/util/limit/limiter_test.go b/pkg/util/limit/limiter_test.go new file mode 100644 index 00000000..f245bcfc --- /dev/null +++ b/pkg/util/limit/limiter_test.go @@ -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()) + }) + } +} diff --git a/pkg/util/limit/reader.go b/pkg/util/limit/reader.go index efa828f4..eccca5a5 100644 --- a/pkg/util/limit/reader.go +++ b/pkg/util/limit/reader.go @@ -35,6 +35,12 @@ func NewReader(r io.Reader, limiter *rate.Limiter) *Reader { func (r *Reader) Read(p []byte) (n int, err error) { b := r.limiter.Burst() + if b <= 0 { + if len(p) == 0 { + return 0, nil + } + return 0, invalidBurstError(b) + } if b < len(p) { p = p[:b] } diff --git a/pkg/util/limit/writer.go b/pkg/util/limit/writer.go index 5256d1e2..56357228 100644 --- a/pkg/util/limit/writer.go +++ b/pkg/util/limit/writer.go @@ -34,8 +34,15 @@ func NewWriter(w io.Writer, limiter *rate.Limiter) *Writer { } func (w *Writer) Write(p []byte) (n int, err error) { + if len(p) == 0 { + return 0, nil + } + var nn int b := w.limiter.Burst() + if b <= 0 { + return 0, invalidBurstError(b) + } for { end := len(p) if end == 0 { diff --git a/server/proxy/proxy.go b/server/proxy/proxy.go index e0dbb7f2..f4e75f0a 100644 --- a/server/proxy/proxy.go +++ b/server/proxy/proxy.go @@ -536,7 +536,7 @@ func NewProxy(ctx context.Context, options *Options) (pxy Proxy, err error) { var limiter *rate.Limiter limitBytes := configurer.GetBaseConfig().Transport.BandwidthLimit.Bytes() if limitBytes > 0 && configurer.GetBaseConfig().Transport.BandwidthLimitMode == types.BandwidthLimitModeServer { - limiter = rate.NewLimiter(rate.Limit(float64(limitBytes)), int(limitBytes)) + limiter = limit.NewBandwidthLimiter(limitBytes) } basePxy := BaseProxy{