status-go/whisper/rate_limiter_test.go

158 lines
3.6 KiB
Go

package whisper
import (
"bytes"
"testing"
"time"
"github.com/ethereum/go-ethereum/p2p/enode"
"github.com/ethereum/go-ethereum/p2p"
"github.com/stretchr/testify/require"
)
func TestPeerRateLimiterDecorator(t *testing.T) {
in, out := p2p.MsgPipe()
payload := []byte{0x01, 0x02, 0x03}
msg := p2p.Msg{
Code: 1,
Size: uint32(len(payload)),
Payload: bytes.NewReader(payload),
ReceivedAt: time.Now(),
}
go func() {
err := in.WriteMsg(msg)
require.NoError(t, err)
}()
messages := make(chan p2p.Msg, 1)
runLoop := func(p *Peer, rw p2p.MsgReadWriter) error {
msg, err := rw.ReadMsg()
if err != nil {
return err
}
messages <- msg
return nil
}
r := NewPeerRateLimiter(&mockRateLimiterHandler{}, nil)
err := r.decorate(nil, out, runLoop)
require.NoError(t, err)
receivedMsg := <-messages
receivedPayload := make([]byte, receivedMsg.Size)
_, err = receivedMsg.Payload.Read(receivedPayload)
require.NoError(t, err)
require.Equal(t, msg.Code, receivedMsg.Code)
require.Equal(t, payload, receivedPayload)
}
func TestPeerLimiterThrottlingWithZeroLimit(t *testing.T) {
r := NewPeerRateLimiter(&mockRateLimiterHandler{}, &PeerRateLimiterConfig{})
for i := 0; i < 1000; i++ {
throttle := r.throttleIP("<nil>")
require.False(t, throttle)
throttle = r.throttlePeer([]byte{0x01, 0x02, 0x03})
require.False(t, throttle)
}
}
func TestPeerLimiterHandler(t *testing.T) {
h := &mockRateLimiterHandler{}
r := NewPeerRateLimiter(h, nil)
p := &Peer{
peer: p2p.NewPeer(enode.ID{0xaa, 0xbb, 0xcc}, "test-peer", nil),
}
rw1, rw2 := p2p.MsgPipe()
count := 100
go func() {
err := echoMessages(r, p, rw2)
require.NoError(t, err)
}()
done := make(chan struct{})
go func() {
for i := 0; i < count; i++ {
msg, err := rw1.ReadMsg()
require.NoError(t, err)
require.EqualValues(t, 101, msg.Code)
}
close(done)
}()
for i := 0; i < count; i += 1 {
err := rw1.WriteMsg(p2p.Msg{Code: 101})
require.NoError(t, err)
}
<-done
require.EqualValues(t, 100-defaultPeerRateLimiterConfig.LimitPerSecIP, h.exceedIPLimit)
require.EqualValues(t, 100-defaultPeerRateLimiterConfig.LimitPerSecPeerID, h.exceedPeerLimit)
}
func TestPeerLimiterHandlerWithWhitelisting(t *testing.T) {
h := &mockRateLimiterHandler{}
r := NewPeerRateLimiter(h, &PeerRateLimiterConfig{
LimitPerSecIP: 1,
LimitPerSecPeerID: 1,
WhitelistedIPs: []string{"<nil>"}, // no IP is represented as <nil> string
WhitelistedPeerIDs: []enode.ID{enode.ID{0xaa, 0xbb, 0xcc}},
})
p := &Peer{
peer: p2p.NewPeer(enode.ID{0xaa, 0xbb, 0xcc}, "test-peer", nil),
}
rw1, rw2 := p2p.MsgPipe()
count := 100
go func() {
err := echoMessages(r, p, rw2)
require.NoError(t, err)
}()
done := make(chan struct{})
go func() {
for i := 0; i < count; i++ {
msg, err := rw1.ReadMsg()
require.NoError(t, err)
require.EqualValues(t, 101, msg.Code)
}
close(done)
}()
for i := 0; i < count; i += 1 {
err := rw1.WriteMsg(p2p.Msg{Code: 101})
require.NoError(t, err)
}
<-done
require.Equal(t, 0, h.exceedIPLimit)
require.Equal(t, 0, h.exceedPeerLimit)
}
func echoMessages(r *PeerRateLimiter, p *Peer, rw p2p.MsgReadWriter) error {
return r.decorate(p, rw, func(p *Peer, rw p2p.MsgReadWriter) error {
for {
msg, err := rw.ReadMsg()
if err != nil {
return err
}
err = rw.WriteMsg(msg)
if err != nil {
return err
}
}
})
}
type mockRateLimiterHandler struct {
exceedPeerLimit int
exceedIPLimit int
}
func (m *mockRateLimiterHandler) ExceedPeerLimit() { m.exceedPeerLimit += 1 }
func (m *mockRateLimiterHandler) ExceedIPLimit() { m.exceedIPLimit += 1 }