Refactor mailserver tests (#970)

mailserver/cleaner_test.go

@@ -8,6 +8,7 @@ import (
 	"github.com/ethereum/go-ethereum/common"
 	"github.com/ethereum/go-ethereum/rlp"
 	whisper "github.com/ethereum/go-ethereum/whisper/whisperv6"
+	"github.com/stretchr/testify/require"
 	"github.com/syndtr/goleveldb/leveldb"
 	"github.com/syndtr/goleveldb/leveldb/storage"
 	"github.com/syndtr/goleveldb/leveldb/util"
@@ -73,7 +74,8 @@ func setupTestServer(t *testing.T) *WMailServer {
 }
 
 func archiveEnvelope(t *testing.T, sentTime time.Time, server *WMailServer) *whisper.Envelope {
-	env := generateEnvelope(t, sentTime)
+	env, err := generateEnvelope(sentTime)
+	require.NoError(t, err)
 	server.Archive(env)
 
 	return env
@@ -82,14 +84,15 @@ func archiveEnvelope(t *testing.T, sentTime time.Time, server *WMailServer) *whi
 func testPrune(t *testing.T, u time.Time, expected int, c *Cleaner, s *WMailServer) {
 	upper := uint32(u.Unix())
 	_, err := c.Prune(0, upper)
-	assert(err == nil, "", t)
+	require.NoError(t, err)
+
 	count := countMessages(t, s.db)
-	assert(count == expected, fmt.Sprintf("expected %d message, got: %d", expected, count), t)
+	require.Equal(t, expected, count, fmt.Sprintf("expected %d message, got: %d", expected, count))
 }
 
 func testMessagesCount(t *testing.T, expected int, s *WMailServer) {
 	count := countMessages(t, s.db)
-	assert(count == expected, fmt.Sprintf("expected %d message, got: %d", expected, count), t)
+	require.Equal(t, expected, count, fmt.Sprintf("expected %d message, got: %d", expected, count))
 }
 
 func countMessages(t *testing.T, db *leveldb.DB) int {

mailserver/limiter.go

@@ -0,0 +1,50 @@
+package mailserver
+
+import (
+	"sync"
+	"time"
+)
+
+type limiter struct {
+	mu sync.RWMutex
+
+	timeout time.Duration
+	db      map[string]time.Time
+}
+
+func newLimiter(timeout time.Duration) *limiter {
+	return &limiter{
+		timeout: timeout,
+		db:      make(map[string]time.Time),
+	}
+}
+
+func (l *limiter) add(id string) {
+	l.mu.Lock()
+	defer l.mu.Unlock()
+
+	l.db[id] = time.Now()
+}
+
+func (l *limiter) isAllowed(id string) bool {
+	l.mu.RLock()
+	defer l.mu.RUnlock()
+
+	if lastRequestTime, ok := l.db[id]; ok {
+		return lastRequestTime.Add(l.timeout).Before(time.Now())
+	}
+
+	return true
+}
+
+func (l *limiter) deleteExpired() {
+	l.mu.Lock()
+	defer l.mu.Unlock()
+
+	now := time.Now()
+	for id, lastRequestTime := range l.db {
+		if lastRequestTime.Add(l.timeout).Before(now) {
+			delete(l.db, id)
+		}
+	}
+}

mailserver/limiter_test.go

@@ -0,0 +1,93 @@
+package mailserver
+
+import (
+	"fmt"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestIsAllowed(t *testing.T) {
+	peerID := "peerID"
+	testCases := []struct {
+		t               time.Duration
+		shouldBeAllowed bool
+		db              func() map[string]time.Time
+		errMsg          string
+		info            string
+	}{
+		{
+			t:               5 * time.Millisecond,
+			shouldBeAllowed: true,
+			db: func() map[string]time.Time {
+				return make(map[string]time.Time)
+			},
+			errMsg: "Expected limiter not to allow with empty db",
+			info:   "Expecting limiter.isAllowed to not allow with an empty db",
+		},
+		{
+			t:               5 * time.Millisecond,
+			shouldBeAllowed: true,
+			db: func() map[string]time.Time {
+				db := make(map[string]time.Time)
+				db[peerID] = time.Now().Add(time.Duration(-10) * time.Millisecond)
+				return db
+			},
+			errMsg: "Expected limiter to allow with peer on its db",
+			info:   "Expecting limiter.isAllowed to allow with an expired peer on its db",
+		},
+		{
+			t:               5 * time.Millisecond,
+			shouldBeAllowed: false,
+			db: func() map[string]time.Time {
+				db := make(map[string]time.Time)
+				db[peerID] = time.Now().Add(time.Duration(-1) * time.Millisecond)
+				return db
+			},
+			errMsg: "Expected limiter to not allow with peer on its db",
+			info:   "Expecting limiter.isAllowed to not allow with a non expired peer on its db",
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.info, func(*testing.T) {
+			l := newLimiter(tc.t)
+			l.db = tc.db()
+			assert.Equal(t, tc.shouldBeAllowed, l.isAllowed(peerID), tc.errMsg)
+		})
+	}
+}
+
+func TestRemoveExpiredRateLimits(t *testing.T) {
+	peer := "peer"
+	l := newLimiter(time.Duration(5) * time.Second)
+	for i := 0; i < 10; i++ {
+		peerID := fmt.Sprintf("%s%d", peer, i)
+		l.db[peerID] = time.Now().Add(time.Duration(i*(-2)) * time.Second)
+	}
+
+	l.deleteExpired()
+	assert.Equal(t, 3, len(l.db))
+
+	for i := 0; i < 3; i++ {
+		peerID := fmt.Sprintf("%s%d", peer, i)
+		_, ok := l.db[peerID]
+		assert.True(t, ok, fmt.Sprintf("Non expired peer '%s' should exist, but it doesn't", peerID))
+	}
+	for i := 3; i < 10; i++ {
+		peerID := fmt.Sprintf("%s%d", peer, i)
+		_, ok := l.db[peerID]
+		assert.False(t, ok, fmt.Sprintf("Expired peer '%s' should not exist, but it does", peerID))
+	}
+}
+
+func TestAddingLimts(t *testing.T) {
+	peerID := "peerAdding"
+	l := newLimiter(time.Duration(5) * time.Second)
+	pre := time.Now()
+	l.add(peerID)
+	post := time.Now()
+	assert.True(t, l.db[peerID].After(pre))
+	assert.True(t, l.db[peerID].Before(post))
+}

mailserver/mailserver.go

@@ -18,8 +18,9 @@ package mailserver
 
 import (
 	"encoding/binary"
+	"errors"
 	"fmt"
-	"sync"
+
 	"time"
 
 	"github.com/ethereum/go-ethereum/common"
@@ -36,7 +37,12 @@ const (
 	maxQueryRange = 24 * time.Hour
 )
 
-// WMailServer whisper mailserver
+var (
+	errDirectoryNotProvided = errors.New("data directory not provided")
+	errPasswordNotProvided  = errors.New("password is not specified")
+)
+
+// WMailServer whisper mailserver.
 type WMailServer struct {
 	db    *leveldb.DB
 	w     *whisper.Whisper
@@ -46,14 +52,14 @@ type WMailServer struct {
 	tick  *ticker
 }
 
-// DBKey key to be stored on db
+// DBKey key to be stored on db.
 type DBKey struct {
 	timestamp uint32
 	hash      common.Hash
 	raw       []byte
 }
 
-// NewDbKey creates a new DBKey with the given values
+// NewDbKey creates a new DBKey with the given values.
 func NewDbKey(t uint32, h common.Hash) *DBKey {
 	const sz = common.HashLength + 4
 	var k DBKey
@@ -65,16 +71,16 @@ func NewDbKey(t uint32, h common.Hash) *DBKey {
 	return &k
 }
 
-// Init initializes mailServer
+// Init initializes mailServer.
 func (s *WMailServer) Init(shh *whisper.Whisper, config *params.WhisperConfig) error {
 	var err error
 
 	if len(config.DataDir) == 0 {
-		return fmt.Errorf("data directory not provided")
+		return errDirectoryNotProvided
 	}
 
 	if len(config.Password) == 0 {
-		return fmt.Errorf("password is not specified")
+		return errPasswordNotProvided
 	}
 
 	s.db, err = leveldb.OpenFile(config.DataDir, nil)
@@ -85,34 +91,50 @@ func (s *WMailServer) Init(shh *whisper.Whisper, config *params.WhisperConfig) e
 	s.w = shh
 	s.pow = config.MinimumPoW
 
+	if err := s.setupWhisperIdentity(config); err != nil {
+		return err
+	}
+	s.setupLimiter(time.Duration(config.MailServerRateLimit) * time.Second)
+
+	return nil
+}
+
+// setupLimiter in case limit is bigger than 0 it will setup an automated
+// limit db cleanup.
+func (s *WMailServer) setupLimiter(rateLimit time.Duration) {
+	limit := rateLimit * time.Second
+	if limit > 0 {
+		s.limit = newLimiter(limit)
+		s.setupMailServerCleanup(limit)
+	}
+}
+
+// setupWhisperIdentity setup the whisper identity (symkey) for the current mail
+// server.
+func (s *WMailServer) setupWhisperIdentity(config *params.WhisperConfig) error {
 	MailServerKeyID, err := s.w.AddSymKeyFromPassword(config.Password)
 	if err != nil {
 		return fmt.Errorf("create symmetric key: %s", err)
 	}
+
 	s.key, err = s.w.GetSymKey(MailServerKeyID)
 	if err != nil {
 		return fmt.Errorf("save symmetric key: %s", err)
 	}
-	limit := time.Duration(config.MailServerRateLimit) * time.Second
-	if limit > 0 {
-		s.limit = newLimiter(limit)
-		s.setupMailServerCleanup(limit)
-	}
 
 	return nil
 }
 
+// setupMailServerCleanup periodically runs an expired entries deleteion for
+// stored limits.
 func (s *WMailServer) setupMailServerCleanup(period time.Duration) {
-	if period <= 0 {
-		return
-	}
 	if s.tick == nil {
 		s.tick = &ticker{}
 	}
 	go s.tick.run(period, s.limit.deleteExpired)
 }
 
-// Close the mailserver and its associated db connection
+// Close the mailserver and its associated db connection.
 func (s *WMailServer) Close() {
 	if s.db != nil {
 		if err := s.db.Close(); err != nil {
@@ -124,41 +146,47 @@ func (s *WMailServer) Close() {
 	}
 }
 
-// Archive a whisper envelope
+// Archive a whisper envelope.
 func (s *WMailServer) Archive(env *whisper.Envelope) {
 	key := NewDbKey(env.Expiry-env.TTL, env.Hash())
 	rawEnvelope, err := rlp.EncodeToBytes(env)
 	if err != nil {
 		log.Error(fmt.Sprintf("rlp.EncodeToBytes failed: %s", err))
 	} else {
-		err = s.db.Put(key.raw, rawEnvelope, nil)
+		if err = s.db.Put(key.raw, rawEnvelope, nil); err != nil {
-		if err != nil {
 			log.Error(fmt.Sprintf("Writing to DB failed: %s", err))
 		}
 	}
 }
 
-// DeliverMail sends mail to specified whisper peer
+// DeliverMail sends mail to specified whisper peer.
 func (s *WMailServer) DeliverMail(peer *whisper.Peer, request *whisper.Envelope) {
 	if peer == nil {
 		log.Error("Whisper peer is nil")
 		return
 	}
+	s.managePeerLimits(peer.ID())
+
+	if ok, lower, upper, bloom := s.validateRequest(peer.ID(), request); ok {
+		s.processRequest(peer, lower, upper, bloom)
+	}
+}
+
+// managePeerLimits in case limit its been setup on the current server and limit
+// allows the query, it will store/update new query time for the current peer.
+func (s *WMailServer) managePeerLimits(peer []byte) {
 	if s.limit != nil {
-		peerID := string(peer.ID())
+		peerID := string(peer)
 		if !s.limit.isAllowed(peerID) {
 			log.Info("peerID exceeded the number of requests per second")
 			return
 		}
 		s.limit.add(peerID)
 	}
-
-	ok, lower, upper, bloom := s.validateRequest(peer.ID(), request)
-	if ok {
-		s.processRequest(peer, lower, upper, bloom)
-	}
 }
 
+// processRequest processes the current request and re-sends all stored messages
+// accomplishing lower and upper limits.
 func (s *WMailServer) processRequest(peer *whisper.Peer, lower, upper uint32, bloom []byte) []*whisper.Envelope {
 	ret := make([]*whisper.Envelope, 0)
 	var err error
@@ -197,6 +225,7 @@ func (s *WMailServer) processRequest(peer *whisper.Peer, lower, upper uint32, bl
 	return ret
 }
 
+// validateRequest runs different validations on the current request.
 func (s *WMailServer) validateRequest(peerID []byte, request *whisper.Envelope) (bool, uint32, uint32, []byte) {
 	if s.pow > 0.0 && request.PoW() < s.pow {
 		return false, 0, 0, nil
@@ -209,27 +238,14 @@ func (s *WMailServer) validateRequest(peerID []byte, request *whisper.Envelope)
 		return false, 0, 0, nil
 	}
 
-	src := crypto.FromECDSAPub(decrypted.Src)
+	if err := s.checkMsgSignature(decrypted, peerID); err != nil {
-	if len(src)-len(peerID) == 1 {
+		log.Warn(err.Error())
-		src = src[1:]
-	}
-
-	// if you want to check the signature, you can do it here. e.g.:
-	// if !bytes.Equal(peerID, src) {
-	if src == nil {
-		log.Warn(fmt.Sprintf("Wrong signature of p2p request"))
 		return false, 0, 0, nil
 	}
 
-	payloadSize := len(decrypted.Payload)
+	bloom, err := s.bloomFromReceivedMessage(decrypted)
-	bloom := decrypted.Payload[8 : 8+whisper.BloomFilterSize]
+	if err != nil {
-	if payloadSize < 8 {
+		log.Warn(err.Error())
-		log.Warn(fmt.Sprintf("Undersized p2p request"))
-		return false, 0, 0, nil
-	} else if payloadSize == 8 {
-		bloom = whisper.MakeFullNodeBloom()
-	} else if payloadSize < 8+whisper.BloomFilterSize {
-		log.Warn(fmt.Sprintf("Undersized bloom filter in p2p request"))
 		return false, 0, 0, nil
 	}
 
@@ -246,46 +262,34 @@ func (s *WMailServer) validateRequest(peerID []byte, request *whisper.Envelope)
 	return true, lower, upper, bloom
 }
 
-type limiter struct {
+// checkMsgSignature returns an error in case the message is not correcly signed
-	mu sync.RWMutex
+func (s *WMailServer) checkMsgSignature(msg *whisper.ReceivedMessage, id []byte) error {
-
+	src := crypto.FromECDSAPub(msg.Src)
-	timeout time.Duration
+	if len(src)-len(id) == 1 {
-	db      map[string]time.Time
+		src = src[1:]
-}
-
-func newLimiter(timeout time.Duration) *limiter {
-	return &limiter{
-		timeout: timeout,
-		db:      make(map[string]time.Time),
-	}
-}
-
-func (l *limiter) add(id string) {
-	l.mu.Lock()
-	defer l.mu.Unlock()
-
-	l.db[id] = time.Now()
-}
-
-func (l *limiter) isAllowed(id string) bool {
-	l.mu.RLock()
-	defer l.mu.RUnlock()
-
-	if lastRequestTime, ok := l.db[id]; ok {
-		return lastRequestTime.Add(l.timeout).Before(time.Now())
 	}
 
-	return true
+	// if you want to check the signature, you can do it here. e.g.:
-}
+	// if !bytes.Equal(peerID, src) {
-
+	if src == nil {
-func (l *limiter) deleteExpired() {
+		return errors.New("Wrong signature of p2p request")
-	l.mu.Lock()
-	defer l.mu.Unlock()
-
-	now := time.Now()
-	for id, lastRequestTime := range l.db {
-		if lastRequestTime.Add(l.timeout).Before(now) {
-			delete(l.db, id)
-		}
 	}
+
+	return nil
+}
+
+// bloomFromReceivedMessage gor a given whisper.ReceivedMessage it extracts the
+// used bloom filter
+func (s *WMailServer) bloomFromReceivedMessage(msg *whisper.ReceivedMessage) ([]byte, error) {
+	payloadSize := len(msg.Payload)
+
+	if payloadSize < 8 {
+		return nil, errors.New("Undersized p2p request")
+	} else if payloadSize == 8 {
+		return whisper.MakeFullNodeBloom(), nil
+	} else if payloadSize < 8+whisper.BloomFilterSize {
+		return nil, errors.New("Undersized bloom filter in p2p request")
+	}
+
+	return msg.Payload[8 : 8+whisper.BloomFilterSize], nil
 }

mailserver/mailserver_test.go

@@ -20,21 +20,23 @@ import (
 	"bytes"
 	"crypto/ecdsa"
 	"encoding/binary"
+	"errors"
+	"fmt"
 	"io/ioutil"
 	"testing"
 	"time"
 
 	"github.com/ethereum/go-ethereum/common"
 	"github.com/ethereum/go-ethereum/crypto"
+	"github.com/ethereum/go-ethereum/rlp"
 	whisper "github.com/ethereum/go-ethereum/whisper/whisperv6"
 	"github.com/status-im/status-go/geth/params"
+	"github.com/stretchr/testify/suite"
 )
 
 const powRequirement = 0.00001
-const peerID = "peerID"
 
 var keyID string
-var shh *whisper.Whisper
 var seed = time.Now().Unix()
 
 type ServerTestParams struct {
@@ -45,83 +47,312 @@ type ServerTestParams struct {
 	key   *ecdsa.PrivateKey
 }
 
-func assert(statement bool, text string, t *testing.T) {
+func TestMailserverSuite(t *testing.T) {
-	if !statement {
+	suite.Run(t, new(MailserverSuite))
-		t.Fatal(text)
+}
+
+type MailserverSuite struct {
+	suite.Suite
+	server *WMailServer
+	shh    *whisper.Whisper
+	config *params.WhisperConfig
+}
+
+func (s *MailserverSuite) SetupTest() {
+	s.server = &WMailServer{}
+	s.shh = whisper.New(&whisper.DefaultConfig)
+	s.shh.RegisterServer(s.server)
+	s.config = &params.WhisperConfig{
+		DataDir:             "/tmp/",
+		Password:            "pwd",
+		MailServerRateLimit: 5,
 	}
 }
 
-func TestDBKey(t *testing.T) {
+func (s *MailserverSuite) TestInit() {
+	testCases := []struct {
+		config        params.WhisperConfig
+		expectedError error
+		limiterActive bool
+		info          string
+	}{
+		{
+			config:        params.WhisperConfig{DataDir: ""},
+			expectedError: errDirectoryNotProvided,
+			limiterActive: false,
+			info:          "Initializing a mail server with a config with empty DataDir",
+		},
+		{
+			config:        params.WhisperConfig{DataDir: "/tmp/", Password: ""},
+			expectedError: errPasswordNotProvided,
+			limiterActive: false,
+			info:          "Initializing a mail server with a config with an empty password",
+		},
+		{
+			config:        params.WhisperConfig{DataDir: "/invalid-path", Password: "pwd"},
+			expectedError: errors.New("open DB: mkdir /invalid-path: permission denied"),
+			limiterActive: false,
+			info:          "Initializing a mail server with a config with an unexisting DataDir",
+		},
+		{
+			config:        *s.config,
+			expectedError: nil,
+			limiterActive: true,
+			info:          "Initializing a mail server with a config with correct config and active limiter",
+		},
+		{
+			config: params.WhisperConfig{
+				DataDir:  "/tmp/",
+				Password: "pwd",
+			},
+			expectedError: nil,
+			limiterActive: false,
+			info:          "Initializing a mail server with a config with empty DataDir and inactive limiter",
+		},
+	}
+
+	for _, tc := range testCases {
+		s.T().Run(tc.info, func(*testing.T) {
+			s.server.limit = nil
+			err := s.server.Init(s.shh, &tc.config)
+			s.server.tick = nil
+			s.server.Close()
+			s.Equal(tc.expectedError, err)
+			s.Equal(tc.limiterActive, (s.server.limit != nil))
+		})
+	}
+}
+
+func (s *MailserverSuite) TestArchive() {
+	err := s.server.Init(s.shh, s.config)
+	s.server.tick = nil
+	s.NoError(err)
+	defer s.server.Close()
+
+	env, err := generateEnvelope(time.Now())
+	s.NoError(err)
+	rawEnvelope, err := rlp.EncodeToBytes(env)
+	s.NoError(err)
+
+	s.server.Archive(env)
+	key := NewDbKey(env.Expiry-env.TTL, env.Hash())
+	archivedEnvelope, err := s.server.db.Get(key.raw, nil)
+	s.NoError(err)
+
+	s.Equal(rawEnvelope, archivedEnvelope)
+}
+
+func (s *MailserverSuite) TestManageLimits() {
+	s.server.limit = newLimiter(time.Duration(5) * time.Millisecond)
+	s.server.managePeerLimits([]byte("peerID"))
+	s.Equal(1, len(s.server.limit.db))
+	firstSaved := s.server.limit.db["peerID"]
+
+	// second call when limit is not accomplished does not store a new limit
+	s.server.managePeerLimits([]byte("peerID"))
+	s.Equal(1, len(s.server.limit.db))
+	s.Equal(firstSaved, s.server.limit.db["peerID"])
+}
+
+func (s *MailserverSuite) TestDBKey() {
 	var h common.Hash
 	i := uint32(time.Now().Unix())
 	k := NewDbKey(i, h)
-	assert(len(k.raw) == common.HashLength+4, "wrong DB key length", t)
+	s.Equal(len(k.raw), common.HashLength+4, "wrong DB key length")
-	assert(byte(i%0x100) == k.raw[3], "raw representation should be big endian", t)
+	s.Equal(byte(i%0x100), k.raw[3], "raw representation should be big endian")
-	assert(byte(i/0x1000000) == k.raw[0], "big endian expected", t)
+	s.Equal(byte(i/0x1000000), k.raw[0], "big endian expected")
 }
 
-func TestMailServer(t *testing.T) {
+func (s *MailserverSuite) TestMailServer() {
 	var server WMailServer
 
-	setupServer(t, &server)
+	s.setupServer(&server)
 	defer server.Close()
 
-	env := generateEnvelope(t, time.Now())
+	env, err := generateEnvelope(time.Now())
+	s.NoError(err)
+
 	server.Archive(env)
-	deliverTest(t, &server, env)
+	testCases := []struct {
+		params      *ServerTestParams
+		emptyLow    bool
+		lowModifier int32
+		uppModifier int32
+		topic       byte
+		expect      bool
+		shouldFail  bool
+		info        string
+	}{
+		{
+			params:      s.defaultServerParams(env),
+			lowModifier: 0,
+			uppModifier: 0,
+			expect:      true,
+			shouldFail:  false,
+			info:        "Processing a request where from and to are equals to an existing register, should provide results",
+		},
+		{
+			params:      s.defaultServerParams(env),
+			lowModifier: 1,
+			uppModifier: 1,
+			expect:      false,
+			shouldFail:  false,
+			info:        "Processing a request where from and to are great than any existing register, should not provide results",
+		},
+		{
+			params:      s.defaultServerParams(env),
+			lowModifier: 0,
+			uppModifier: 1,
+			topic:       0xFF,
+			expect:      false,
+			shouldFail:  false,
+			info:        "Processing a request where to is grat than any existing register and with a specific topic, should not provide results",
+		},
+		{
+			params:      s.defaultServerParams(env),
+			emptyLow:    true,
+			lowModifier: 4,
+			uppModifier: -1,
+			shouldFail:  true,
+			info:        "Processing a request where to is lower than from should fail",
+		},
+		{
+			params:      s.defaultServerParams(env),
+			emptyLow:    true,
+			lowModifier: 0,
+			uppModifier: 24,
+			shouldFail:  true,
+			info:        "Processing a request where difference between from and to is > 24 should fail",
+		},
+	}
+	for _, tc := range testCases {
+		s.T().Run(tc.info, func(*testing.T) {
+			if tc.lowModifier != 0 {
+				tc.params.low = tc.params.birth + uint32(tc.lowModifier)
+			}
+			if tc.uppModifier != 0 {
+				tc.params.upp = tc.params.birth + uint32(tc.uppModifier)
+			}
+			if tc.emptyLow {
+				tc.params.low = 0
+			}
+			if tc.topic == 0xFF {
+				tc.params.topic[0] = tc.topic
+			}
+
+			request := s.createRequest(tc.params)
+			src := crypto.FromECDSAPub(&tc.params.key.PublicKey)
+			ok, lower, upper, bloom := server.validateRequest(src, request)
+			if tc.shouldFail {
+				if ok {
+					s.T().Fatal(err)
+				}
+				return
+			}
+			if !ok {
+				s.T().Fatalf("request validation failed, seed: %d.", seed)
+			}
+			if lower != tc.params.low {
+				s.T().Fatalf("request validation failed (lower bound), seed: %d.", seed)
+			}
+			if upper != tc.params.upp {
+				s.T().Fatalf("request validation failed (upper bound), seed: %d.", seed)
+			}
+			expectedBloom := whisper.TopicToBloom(tc.params.topic)
+			if !bytes.Equal(bloom, expectedBloom) {
+				s.T().Fatalf("request validation failed (topic), seed: %d.", seed)
+			}
+
+			var exist bool
+			mail := server.processRequest(nil, tc.params.low, tc.params.upp, bloom)
+			for _, msg := range mail {
+				if msg.Hash() == env.Hash() {
+					exist = true
+					break
+				}
+			}
+
+			if exist != tc.expect {
+				s.T().Fatalf("error: exist = %v, seed: %d.", exist, seed)
+			}
+
+			src[0]++
+			ok, lower, upper, _ = server.validateRequest(src, request)
+			if !ok {
+				// request should be valid regardless of signature
+				s.T().Fatalf("request validation false negative, seed: %d (lower: %d, upper: %d).", seed, lower, upper)
+			}
+		})
+	}
 }
 
-func TestRateLimits(t *testing.T) {
+func (s *MailserverSuite) TestBloomFromReceivedMessage() {
-	l := newLimiter(time.Duration(5 * time.Millisecond))
+	testCases := []struct {
-	assert(l.isAllowed(peerID), "Expected limiter not to allow with empty db", t)
+		msg           whisper.ReceivedMessage
+		expectedBloom []byte
+		expectedErr   error
+		info          string
+	}{
+		{
+			msg:           whisper.ReceivedMessage{},
+			expectedBloom: []byte(nil),
+			expectedErr:   errors.New("Undersized p2p request"),
+			info:          "getting bloom filter for an empty whisper message should produce an error",
+		},
+		{
+			msg:           whisper.ReceivedMessage{Payload: []byte("hohohohoho")},
+			expectedBloom: []byte(nil),
+			expectedErr:   errors.New("Undersized bloom filter in p2p request"),
+			info:          "getting bloom filter for a malformed whisper message should produce an error",
+		},
+		{
+			msg:           whisper.ReceivedMessage{Payload: []byte("12345678")},
+			expectedBloom: whisper.MakeFullNodeBloom(),
+			expectedErr:   nil,
+			info:          "getting bloom filter for a valid whisper message should be successful",
+		},
+	}
 
-	l.db[peerID] = time.Now().Add(time.Duration(-10 * time.Millisecond))
+	for _, tc := range testCases {
-	assert(l.isAllowed(peerID), "Expected limiter to allow with peer on its db", t)
+		s.T().Run(tc.info, func(*testing.T) {
-
+			bloom, err := s.server.bloomFromReceivedMessage(&tc.msg)
-	l.db[peerID] = time.Now().Add(time.Duration(-1 * time.Millisecond))
+			s.Equal(tc.expectedErr, err)
-	assert(!l.isAllowed(peerID), "Expected limiter to not allow with peer on its db", t)
+			s.Equal(tc.expectedBloom, bloom)
+		})
+	}
 }
 
-func TestRemoveExpiredRateLimits(t *testing.T) {
+func (s *MailserverSuite) setupServer(server *WMailServer) {
-	l := newLimiter(time.Duration(10) * time.Second)
+	const password = "password_for_this_test"
-	l.db[peerID] = time.Now().Add(time.Duration(-10) * time.Second)
+	const dbPath = "whisper-server-test"
-	l.db[peerID+"A"] = time.Now().Add(time.Duration(10) * time.Second)
+
-	l.deleteExpired()
+	dir, err := ioutil.TempDir("", dbPath)
-	_, ok := l.db[peerID]
+	if err != nil {
-	assert(!ok, "Expired peer should not exist, but it does ", t)
+		s.T().Fatal(err)
-	_, ok = l.db[peerID+"A"]
+	}
-	assert(ok, "Non expired peer should exist, but it doesn't", t)
+
+	s.shh = whisper.New(&whisper.DefaultConfig)
+	s.shh.RegisterServer(server)
+
+	err = server.Init(s.shh, &params.WhisperConfig{DataDir: dir, Password: password, MinimumPoW: powRequirement})
+	if err != nil {
+		s.T().Fatal(err)
+	}
+
+	keyID, err = s.shh.AddSymKeyFromPassword(password)
+	if err != nil {
+		s.T().Fatalf("failed to create symmetric key for mail request: %s", err)
+	}
 }
 
-func generateEnvelope(t *testing.T, now time.Time) *whisper.Envelope {
+func (s *MailserverSuite) defaultServerParams(env *whisper.Envelope) *ServerTestParams {
-	h := crypto.Keccak256Hash([]byte("test sample data"))
+	id, err := s.shh.NewKeyPair()
-	params := &whisper.MessageParams{
-		KeySym:   h[:],
-		Topic:    whisper.TopicType{0x1F, 0x7E, 0xA1, 0x7F},
-		Payload:  []byte("test payload"),
-		PoW:      powRequirement,
-		WorkTime: 2,
-	}
-
-	msg, err := whisper.NewSentMessage(params)
 	if err != nil {
-		t.Fatalf("failed to create new message with seed %d: %s.", seed, err)
+		s.T().Fatalf("failed to generate new key pair with seed %d: %s.", seed, err)
 	}
-	env, err := msg.Wrap(params, now)
+	testPeerID, err := s.shh.GetPrivateKey(id)
 	if err != nil {
-		t.Fatalf("failed to wrap with seed %d: %s.", seed, err)
+		s.T().Fatalf("failed to retrieve new key pair with seed %d: %s.", seed, err)
-	}
-	return env
-}
-
-func serverParams(t *testing.T, env *whisper.Envelope) *ServerTestParams {
-	id, err := shh.NewKeyPair()
-	if err != nil {
-		t.Fatalf("failed to generate new key pair with seed %d: %s.", seed, err)
-	}
-	testPeerID, err := shh.GetPrivateKey(id)
-	if err != nil {
-		t.Fatalf("failed to retrieve new key pair with seed %d: %s.", seed, err)
 	}
 	birth := env.Expiry - env.TTL
 
@@ -133,86 +364,17 @@ func serverParams(t *testing.T, env *whisper.Envelope) *ServerTestParams {
 		key:   testPeerID,
 	}
 }
-func deliverTest(t *testing.T, server *WMailServer, env *whisper.Envelope) {
-	p := serverParams(t, env)
-	singleRequest(t, server, env, p, true)
 
-	p.low = p.birth + 1
+func (s *MailserverSuite) createRequest(p *ServerTestParams) *whisper.Envelope {
-	p.upp = p.birth + 1
-	singleRequest(t, server, env, p, false)
-
-	p.low = p.birth
-	p.upp = p.birth + 1
-	p.topic[0] = 0xFF
-	singleRequest(t, server, env, p, false)
-
-	p.low = 0
-	p.upp = p.birth - 1
-	failRequest(t, server, p, "validation should fail due to negative query time range")
-
-	p.low = 0
-	p.upp = p.birth + 24
-	failRequest(t, server, p, "validation should fail due to query big time range")
-}
-
-func failRequest(t *testing.T, server *WMailServer, p *ServerTestParams, err string) {
-	request := createRequest(t, p)
-	src := crypto.FromECDSAPub(&p.key.PublicKey)
-	ok, _, _, _ := server.validateRequest(src, request)
-	if ok {
-		t.Fatalf(err)
-	}
-}
-
-func singleRequest(t *testing.T, server *WMailServer, env *whisper.Envelope, p *ServerTestParams, expect bool) {
-	request := createRequest(t, p)
-	src := crypto.FromECDSAPub(&p.key.PublicKey)
-	ok, lower, upper, bloom := server.validateRequest(src, request)
-	if !ok {
-		t.Fatalf("request validation failed, seed: %d.", seed)
-	}
-	if lower != p.low {
-		t.Fatalf("request validation failed (lower bound), seed: %d.", seed)
-	}
-	if upper != p.upp {
-		t.Fatalf("request validation failed (upper bound), seed: %d.", seed)
-	}
-	expectedBloom := whisper.TopicToBloom(p.topic)
-	if !bytes.Equal(bloom, expectedBloom) {
-		t.Fatalf("request validation failed (topic), seed: %d.", seed)
-	}
-
-	var exist bool
-	mail := server.processRequest(nil, p.low, p.upp, bloom)
-	for _, msg := range mail {
-		if msg.Hash() == env.Hash() {
-			exist = true
-			break
-		}
-	}
-
-	if exist != expect {
-		t.Fatalf("error: exist = %v, seed: %d.", exist, seed)
-	}
-
-	src[0]++
-	ok, lower, upper, _ = server.validateRequest(src, request)
-	if !ok {
-		// request should be valid regardless of signature
-		t.Fatalf("request validation false negative, seed: %d (lower: %d, upper: %d).", seed, lower, upper)
-	}
-}
-
-func createRequest(t *testing.T, p *ServerTestParams) *whisper.Envelope {
 	bloom := whisper.TopicToBloom(p.topic)
 	data := make([]byte, 8)
 	binary.BigEndian.PutUint32(data, p.low)
 	binary.BigEndian.PutUint32(data[4:], p.upp)
 	data = append(data, bloom...)
 
-	key, err := shh.GetSymKey(keyID)
+	key, err := s.shh.GetSymKey(keyID)
 	if err != nil {
-		t.Fatalf("failed to retrieve sym key with seed %d: %s.", seed, err)
+		s.T().Fatalf("failed to retrieve sym key with seed %d: %s.", seed, err)
 	}
 
 	params := &whisper.MessageParams{
@@ -226,34 +388,33 @@ func createRequest(t *testing.T, p *ServerTestParams) *whisper.Envelope {
 
 	msg, err := whisper.NewSentMessage(params)
 	if err != nil {
-		t.Fatalf("failed to create new message with seed %d: %s.", seed, err)
+		s.T().Fatalf("failed to create new message with seed %d: %s.", seed, err)
 	}
 	env, err := msg.Wrap(params, time.Now())
 	if err != nil {
-		t.Fatalf("failed to wrap with seed %d: %s.", seed, err)
+		s.T().Fatalf("failed to wrap with seed %d: %s.", seed, err)
 	}
 	return env
 }
 
-func setupServer(t *testing.T, server *WMailServer) {
+func generateEnvelope(sentTime time.Time) (*whisper.Envelope, error) {
-	const password = "password_for_this_test"
+	h := crypto.Keccak256Hash([]byte("test sample data"))
-	const dbPath = "whisper-server-test"
+	params := &whisper.MessageParams{
-
+		KeySym:   h[:],
-	dir, err := ioutil.TempDir("", dbPath)
+		Topic:    whisper.TopicType{0x1F, 0x7E, 0xA1, 0x7F},
-	if err != nil {
+		Payload:  []byte("test payload"),
-		t.Fatal(err)
+		PoW:      powRequirement,
+		WorkTime: 2,
 	}
 
-	shh = whisper.New(&whisper.DefaultConfig)
+	msg, err := whisper.NewSentMessage(params)
-	shh.RegisterServer(server)
-
-	err = server.Init(shh, &params.WhisperConfig{DataDir: dir, Password: password, MinimumPoW: powRequirement})
 	if err != nil {
-		t.Fatal(err)
+		return nil, fmt.Errorf("failed to create new message with seed %d: %s", seed, err)
+	}
+	env, err := msg.Wrap(params, sentTime)
+	if err != nil {
+		return nil, fmt.Errorf("failed to wrap with seed %d: %s", seed, err)
 	}
 
-	keyID, err = shh.AddSymKeyFromPassword(password)
+	return env, nil
-	if err != nil {
-		t.Fatalf("Failed to create symmetric key for mail request: %s", err)
-	}
 }