feat(ssh): add throttle package to network connections

Author: henrybarreto

ssh/pkg/dialer/throttle.go

@@ -0,0 +1,247 @@
+package dialer
+
+import (
+	"context"
+	"errors"
+	"io"
+	"net"
+	"sync"
+	"time"
+
+	"golang.org/x/time/rate"
+)
+
+// ErrNegativeLimit is returned when attempting to set a negative limit.
+var ErrNegativeLimit = errors.New("negative throttle limit")
+
+// Option configures a Throttler.
+type Option func(*Throttler)
+
+// WithReadLimit sets the read bytes-per-second limit and burst.
+// If bps <= 0 => unlimited. If burst <=0 it defaults to bps.
+func WithReadLimit(bps int, burst int) Option {
+	return func(t *Throttler) {
+		t.setLimiter(&t.readMu, &t.readLimiter, bps, burst)
+	}
+}
+
+// WithWriteLimit sets the write bytes-per-second limit and burst.
+// If bps <= 0 => unlimited. If burst <=0 it defaults to bps.
+func WithWriteLimit(bps int, burst int) Option {
+	return func(t *Throttler) {
+		t.setLimiter(&t.writeMu, &t.writeLimiter, bps, burst)
+	}
+}
+
+// Throttler wraps an underlying io.Reader / io.Writer (optionally both) and
+// enforces directional byte-per-second limits using token buckets.
+// It is safe for concurrent use of Read and Write.
+type Throttler struct {
+	// Underlying read side (may be nil if only writing).
+	R io.Reader
+	// Underlying write side (may be nil if only reading).
+	W io.Writer
+
+	readMu      sync.RWMutex
+	readLimiter *rate.Limiter
+
+	writeMu      sync.RWMutex
+	writeLimiter *rate.Limiter
+}
+
+func NewThrottler(r io.Reader, w io.Writer, opts ...Option) *Throttler {
+	t := &Throttler{R: r, W: w}
+
+	for _, o := range opts {
+		o(t)
+	}
+
+	return t
+}
+
+// setLimiter (internal) creates or clears a limiter based on bps.
+func (t *Throttler) setLimiter(mu *sync.RWMutex, lim **rate.Limiter, bps int, burst int) {
+	mu.Lock()
+	defer mu.Unlock()
+
+	if bps <= 0 {
+		*lim = nil
+
+		return
+	}
+
+	if burst <= 0 {
+		burst = bps
+	}
+
+	*lim = rate.NewLimiter(rate.Limit(bps), burst)
+}
+
+// UpdateReadLimit dynamically changes the read limit.
+func (t *Throttler) UpdateReadLimit(bps int, burst int) error {
+	if bps < 0 || burst < 0 {
+		return ErrNegativeLimit
+	}
+
+	t.setLimiter(&t.readMu, &t.readLimiter, bps, burst)
+
+	return nil
+}
+
+// UpdateWriteLimit dynamically changes the write limit.
+func (t *Throttler) UpdateWriteLimit(bps int, burst int) error {
+	if bps < 0 || burst < 0 {
+		return ErrNegativeLimit
+	}
+
+	t.setLimiter(&t.writeMu, &t.writeLimiter, bps, burst)
+
+	return nil
+}
+
+// Read implements io.Reader with throttling.
+func (t *Throttler) Read(p []byte) (int, error) {
+	if t.R == nil {
+		return 0, errors.New("read not supported (nil underlying Reader)")
+	}
+
+	lim := t.getReadLimiter()
+
+	if lim == nil {
+		return t.R.Read(p)
+	}
+
+	maxChunk := lim.Burst()
+	if maxChunk <= 0 {
+		maxChunk = 32 * 1024
+	}
+
+	total := 0
+	for total < len(p) {
+		remaining := len(p) - total
+		chunk := min(remaining, maxChunk)
+
+		if err := lim.WaitN(context.Background(), chunk); err != nil {
+			if total > 0 {
+				return total, err
+			}
+
+			return 0, err
+		}
+
+		n, err := t.R.Read(p[total : total+chunk])
+		total += n
+		if err != nil || n == 0 {
+			return total, err
+		}
+
+		if n < chunk {
+			break
+		}
+	}
+
+	return total, nil
+}
+
+// Write implements io.Writer with throttling.
+func (t *Throttler) Write(p []byte) (int, error) {
+	if t.W == nil {
+		return 0, errors.New("write not supported (nil underlying Writer)")
+	}
+
+	lim := t.getWriteLimiter()
+
+	if lim == nil {
+		return t.W.Write(p)
+	}
+
+	maxChunk := lim.Burst()
+	if maxChunk <= 0 {
+		maxChunk = 32 * 1024
+	}
+
+	total := 0
+	for total < len(p) {
+		remaining := len(p) - total
+		chunk := min(remaining, maxChunk)
+
+		if err := lim.WaitN(context.Background(), chunk); err != nil {
+			if total > 0 {
+				return total, err
+			}
+
+			return 0, err
+		}
+
+		n, err := t.W.Write(p[total : total+chunk])
+		total += n
+		if err != nil || n == 0 {
+			return total, err
+		}
+
+		if n < chunk {
+			break
+		}
+	}
+
+	return total, nil
+}
+
+// Helper getters with read locks for concurrency.
+func (t *Throttler) getReadLimiter() *rate.Limiter {
+	t.readMu.RLock()
+	defer t.readMu.RUnlock()
+
+	return t.readLimiter
+}
+
+func (t *Throttler) getWriteLimiter() *rate.Limiter {
+	t.writeMu.RLock()
+	defer t.writeMu.RUnlock()
+
+	return t.writeLimiter
+}
+
+type ConnThrottler struct {
+	Conn      net.Conn
+	Throttler *Throttler
+}
+
+func (c *ConnThrottler) Close() error {
+	return c.Conn.Close()
+}
+
+func (c *ConnThrottler) LocalAddr() net.Addr {
+	return c.Conn.LocalAddr()
+}
+
+func (c *ConnThrottler) Read(b []byte) (n int, err error) {
+	return c.Throttler.Read(b)
+}
+
+func (c *ConnThrottler) RemoteAddr() net.Addr {
+	return c.Conn.RemoteAddr()
+}
+
+func (c *ConnThrottler) SetDeadline(t time.Time) error {
+	return c.Conn.SetDeadline(t)
+}
+
+func (c *ConnThrottler) SetReadDeadline(t time.Time) error {
+	return c.Conn.SetReadDeadline(t)
+}
+
+func (c *ConnThrottler) SetWriteDeadline(t time.Time) error {
+	return c.Conn.SetWriteDeadline(t)
+}
+
+func (c *ConnThrottler) Write(b []byte) (n int, err error) {
+	return c.Throttler.Write(b)
+}
+
+func NewConnThrottler(conn net.Conn, readBps, readBurst, writeBps, writeBurst int) net.Conn {
+	return &ConnThrottler{
+		Conn:      conn,
+		Throttler: NewThrottler(conn, conn, WithReadLimit(readBps, readBurst), WithWriteLimit(writeBps, writeBurst)),
+	}
+}

ssh/pkg/dialer/throttle_test.go

@@ -0,0 +1,147 @@
+package dialer
+
+import (
+	"bytes"
+	"io"
+	"net"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func expectedMinDuration(total, bps, burst int) time.Duration {
+	if bps <= 0 {
+		return 0
+	}
+
+	remaining := total - burst
+	if remaining <= 0 {
+		return 0
+	}
+
+	secs := float64(remaining) / float64(bps)
+
+	return time.Duration(secs * float64(time.Second))
+}
+
+func TestThrottler_TableDriven(t *testing.T) {
+	cases := []struct {
+		name string
+		run  func(t *testing.T)
+	}{
+		{
+			name: "UnlimitedReadFast",
+			run: func(t *testing.T) {
+				data := bytes.Repeat([]byte("x"), 1024)
+				r := bytes.NewReader(data)
+
+				th := NewThrottler(r, nil) // no limits
+
+				buf := make([]byte, len(data))
+				start := time.Now()
+				n, err := th.Read(buf)
+				dur := time.Since(start)
+
+				assert.Truef(t, err == nil || err == io.EOF, "unexpected read error: %v", err)
+				assert.Equal(t, len(data), n, "read bytes mismatch")
+				assert.LessOrEqual(t, dur, 100*time.Millisecond, "unlimited read took too long")
+			},
+		},
+		{
+			name: "NegativeLimitValidation",
+			run: func(t *testing.T) {
+				th := NewThrottler(nil, nil)
+				err := th.UpdateReadLimit(-1, 1)
+				assert.Equal(t, ErrNegativeLimit, err)
+				err = th.UpdateWriteLimit(-1, 1)
+				assert.Equal(t, ErrNegativeLimit, err)
+			},
+		},
+		{
+			name: "ReadRateEnforced",
+			run: func(t *testing.T) {
+				total := 200
+				bps := 50
+				burst := 10
+
+				data := bytes.Repeat([]byte("r"), total)
+				r := bytes.NewReader(data)
+				th := NewThrottler(r, nil, WithReadLimit(bps, burst))
+
+				buf := make([]byte, total)
+				start := time.Now()
+				n, err := th.Read(buf)
+				dur := time.Since(start)
+
+				assert.Truef(t, err == nil || err == io.EOF, "unexpected read error: %v", err)
+				assert.Equal(t, total, n, "read bytes mismatch")
+
+				expect := expectedMinDuration(total, bps, burst)
+				// allow 20% timing slack for scheduler and test flakiness
+				slack := expect / 5
+				assert.Truef(t, dur+slack >= expect, "read duration = %v; want at least ~%v (with slack %v)", dur, expect, slack)
+			},
+		},
+		{
+			name: "WriteRateEnforced",
+			run: func(t *testing.T) {
+				total := 200
+				bps := 50
+				burst := 10
+
+				var bufOut bytes.Buffer
+				th := NewThrottler(nil, &bufOut, WithWriteLimit(bps, burst))
+
+				data := bytes.Repeat([]byte("w"), total)
+				start := time.Now()
+				n, err := th.Write(data)
+				dur := time.Since(start)
+
+				assert.NoError(t, err, "unexpected write error")
+				assert.Equal(t, total, n, "written bytes mismatch")
+
+				expect := expectedMinDuration(total, bps, burst)
+				slack := expect / 5
+				assert.Truef(t, dur+slack >= expect, "write duration = %v; want at least ~%v (with slack %v)", dur, expect, slack)
+			},
+		},
+		{
+			name: "ConnThrottlerPassthrough",
+			run: func(t *testing.T) {
+				c1, c2 := net.Pipe()
+				t.Cleanup(func() { c1.Close(); c2.Close() })
+
+				// Wrap c2 with unlimited throttler
+				thrConn := NewConnThrottler(c2, 0, 0, 0, 0)
+
+				// write from c1, read from thrConn
+				msg := []byte("hello-throttle")
+
+				done := make(chan error, 1)
+				go func() {
+					defer c1.Close()
+					_, err := c1.Write(msg)
+					done <- err
+				}()
+
+				// read on wrapped conn
+				got := make([]byte, len(msg))
+				n, err := thrConn.Read(got)
+				assert.Truef(t, err == nil || err == io.EOF, "conn read error: %v", err)
+				assert.Equal(t, len(msg), n, "conn read bytes mismatch")
+				assert.Equal(t, msg, got, "conn read data mismatch")
+
+				// ensure writer had no error
+				err = <-done
+				assert.NoError(t, err, "writer error")
+			},
+		},
+	}
+
+	for _, tc := range cases {
+		t.Run(tc.name, func(t *testing.T) {
+			tc.run(t)
+		})
+	}
+}