Tier based tests

This commit is contained in:
binwiederhier 2023-01-08 20:46:46 -05:00
parent 1f54adad71
commit d8032e1c9e
4 changed files with 35 additions and 9 deletions

View File

@ -55,8 +55,6 @@ import (
- figure out what settings are "web" or "phone" - figure out what settings are "web" or "phone"
Tests: Tests:
- visitor with/without user - visitor with/without user
- plan-based message expiry
- plan-based attachment expiry
Docs: Docs:
- "expires" field in message - "expires" field in message
Refactor: Refactor:
@ -65,7 +63,6 @@ import (
- Password reset - Password reset
- Pricing - Pricing
- change email - change email
*/ */
// Server is the main server, providing the UI and API for ntfy // Server is the main server, providing the UI and API for ntfy
@ -530,6 +527,9 @@ func (s *Server) handlePublishWithoutResponse(r *http.Request, v *visitor) (*mes
if err != nil { if err != nil {
return nil, err return nil, err
} }
if err := v.MessageAllowed(); err != nil {
return nil, errHTTPTooManyRequestsLimitRequests // FIXME make one for messages
}
body, err := util.Peek(r.Body, s.config.MessageLimit) body, err := util.Peek(r.Body, s.config.MessageLimit)
if err != nil { if err != nil {
return nil, err return nil, err

View File

@ -462,6 +462,7 @@ func TestAccount_Reservation_PublishByAnonymousFails(t *testing.T) {
require.Nil(t, s.userManager.CreateTier(&user.Tier{ require.Nil(t, s.userManager.CreateTier(&user.Tier{
Code: "pro", Code: "pro",
MessagesLimit: 20,
ReservationsLimit: 2, ReservationsLimit: 2,
})) }))
require.Nil(t, s.userManager.ChangeTier("phil", "pro")) require.Nil(t, s.userManager.ChangeTier("phil", "pro"))

View File

@ -1098,7 +1098,7 @@ func TestServer_PublishWithTierBasedMessageLimitAndExpiry(t *testing.T) {
require.Nil(t, s.userManager.CreateTier(&user.Tier{ require.Nil(t, s.userManager.CreateTier(&user.Tier{
Code: "test", Code: "test",
MessagesLimit: 5, MessagesLimit: 5,
MessagesExpiryDuration: 1, // Second MessagesExpiryDuration: -5, // Second, what a hack!
})) }))
require.Nil(t, s.userManager.AddUser("phil", "phil", user.RoleUser)) require.Nil(t, s.userManager.AddUser("phil", "phil", user.RoleUser))
require.Nil(t, s.userManager.ChangeTier("phil", "test")) require.Nil(t, s.userManager.ChangeTier("phil", "test"))
@ -1115,7 +1115,15 @@ func TestServer_PublishWithTierBasedMessageLimitAndExpiry(t *testing.T) {
response := request(t, s, "PUT", "/mytopic", "this is too much", map[string]string{ response := request(t, s, "PUT", "/mytopic", "this is too much", map[string]string{
"Authorization": util.BasicAuth("phil", "phil"), "Authorization": util.BasicAuth("phil", "phil"),
}) })
require.Equal(t, 413, response.Code) require.Equal(t, 429, response.Code)
// Run pruning and see if they are gone
s.execManager()
response = request(t, s, "GET", "/mytopic/json?poll=1", "", map[string]string{
"Authorization": util.BasicAuth("phil", "phil"),
})
require.Equal(t, 200, response.Code)
require.Empty(t, response.Body)
} }
func TestServer_PublishAttachment(t *testing.T) { func TestServer_PublishAttachment(t *testing.T) {
@ -1318,6 +1326,7 @@ func TestServer_PublishAttachmentWithTierBasedExpiry(t *testing.T) {
sevenDaysInSeconds := int64(604800) sevenDaysInSeconds := int64(604800)
require.Nil(t, s.userManager.CreateTier(&user.Tier{ require.Nil(t, s.userManager.CreateTier(&user.Tier{
Code: "test", Code: "test",
MessagesLimit: 10,
MessagesExpiryDuration: sevenDaysInSeconds, MessagesExpiryDuration: sevenDaysInSeconds,
AttachmentFileSizeLimit: 50_000, AttachmentFileSizeLimit: 50_000,
AttachmentTotalSizeLimit: 200_000, AttachmentTotalSizeLimit: 200_000,
@ -1362,8 +1371,10 @@ func TestServer_PublishAttachmentWithTierBasedLimits(t *testing.T) {
// Create tier with certain limits // Create tier with certain limits
require.Nil(t, s.userManager.CreateTier(&user.Tier{ require.Nil(t, s.userManager.CreateTier(&user.Tier{
Code: "test", Code: "test",
MessagesLimit: 100,
AttachmentFileSizeLimit: 50_000, AttachmentFileSizeLimit: 50_000,
AttachmentTotalSizeLimit: 200_000, AttachmentTotalSizeLimit: 200_000,
AttachmentExpiryDuration: 30,
})) }))
require.Nil(t, s.userManager.AddUser("phil", "phil", user.RoleUser)) require.Nil(t, s.userManager.AddUser("phil", "phil", user.RoleUser))
require.Nil(t, s.userManager.ChangeTier("phil", "test")) require.Nil(t, s.userManager.ChangeTier("phil", "test"))
@ -1377,12 +1388,14 @@ func TestServer_PublishAttachmentWithTierBasedLimits(t *testing.T) {
// Publish large file as anonymous // Publish large file as anonymous
response = request(t, s, "PUT", "/mytopic", largeFile, nil) response = request(t, s, "PUT", "/mytopic", largeFile, nil)
require.Equal(t, 413, response.Code) require.Equal(t, 413, response.Code)
require.Equal(t, 41301, toHTTPError(t, response.Body.String()).Code)
// Publish too large file as phil // Publish too large file as phil
response = request(t, s, "PUT", "/mytopic", largeFile+" a few more bytes", map[string]string{ response = request(t, s, "PUT", "/mytopic", largeFile+" a few more bytes", map[string]string{
"Authorization": util.BasicAuth("phil", "phil"), "Authorization": util.BasicAuth("phil", "phil"),
}) })
require.Equal(t, 413, response.Code) require.Equal(t, 413, response.Code)
require.Equal(t, 41301, toHTTPError(t, response.Body.String()).Code)
// Publish large file as phil (4x) // Publish large file as phil (4x)
for i := 0; i < 4; i++ { for i := 0; i < 4; i++ {
@ -1398,6 +1411,7 @@ func TestServer_PublishAttachmentWithTierBasedLimits(t *testing.T) {
"Authorization": util.BasicAuth("phil", "phil"), "Authorization": util.BasicAuth("phil", "phil"),
}) })
require.Equal(t, 413, response.Code) require.Equal(t, 413, response.Code)
require.Equal(t, 41301, toHTTPError(t, response.Body.String()).Code)
} }
func TestServer_PublishAttachmentBandwidthLimit(t *testing.T) { func TestServer_PublishAttachmentBandwidthLimit(t *testing.T) {

View File

@ -29,12 +29,13 @@ type visitor struct {
userManager *user.Manager // May be nil! userManager *user.Manager // May be nil!
ip netip.Addr ip netip.Addr
user *user.User user *user.User
messages int64 // Number of messages sent messages int64 // Number of messages sent, reset every day
emails int64 // Number of emails sent emails int64 // Number of emails sent, reset every day
requestLimiter *rate.Limiter // Rate limiter for (almost) all requests (including messages) requestLimiter *rate.Limiter // Rate limiter for (almost) all requests (including messages)
messagesLimiter util.Limiter // Rate limiter for messages, may be nil
emailsLimiter *rate.Limiter // Rate limiter for emails emailsLimiter *rate.Limiter // Rate limiter for emails
subscriptionLimiter util.Limiter // Fixed limiter for active subscriptions (ongoing connections) subscriptionLimiter util.Limiter // Fixed limiter for active subscriptions (ongoing connections)
bandwidthLimiter util.Limiter bandwidthLimiter util.Limiter // Limiter for attachment bandwidth downloads
accountLimiter *rate.Limiter // Rate limiter for account creation accountLimiter *rate.Limiter // Rate limiter for account creation
firebase time.Time // Next allowed Firebase message firebase time.Time // Next allowed Firebase message
seen time.Time seen time.Time
@ -61,6 +62,7 @@ type visitorInfo struct {
} }
func newVisitor(conf *Config, messageCache *messageCache, userManager *user.Manager, ip netip.Addr, user *user.User) *visitor { func newVisitor(conf *Config, messageCache *messageCache, userManager *user.Manager, ip netip.Addr, user *user.User) *visitor {
var messagesLimiter util.Limiter
var requestLimiter, emailsLimiter, accountLimiter *rate.Limiter var requestLimiter, emailsLimiter, accountLimiter *rate.Limiter
var messages, emails int64 var messages, emails int64
if user != nil { if user != nil {
@ -71,6 +73,7 @@ func newVisitor(conf *Config, messageCache *messageCache, userManager *user.Mana
} }
if user != nil && user.Tier != nil { if user != nil && user.Tier != nil {
requestLimiter = rate.NewLimiter(dailyLimitToRate(user.Tier.MessagesLimit), conf.VisitorRequestLimitBurst) requestLimiter = rate.NewLimiter(dailyLimitToRate(user.Tier.MessagesLimit), conf.VisitorRequestLimitBurst)
messagesLimiter = util.NewFixedLimiter(user.Tier.MessagesLimit)
emailsLimiter = rate.NewLimiter(dailyLimitToRate(user.Tier.EmailsLimit), conf.VisitorEmailLimitBurst) emailsLimiter = rate.NewLimiter(dailyLimitToRate(user.Tier.EmailsLimit), conf.VisitorEmailLimitBurst)
} else { } else {
requestLimiter = rate.NewLimiter(rate.Every(conf.VisitorRequestLimitReplenish), conf.VisitorRequestLimitBurst) requestLimiter = rate.NewLimiter(rate.Every(conf.VisitorRequestLimitReplenish), conf.VisitorRequestLimitBurst)
@ -85,6 +88,7 @@ func newVisitor(conf *Config, messageCache *messageCache, userManager *user.Mana
messages: messages, messages: messages,
emails: emails, emails: emails,
requestLimiter: requestLimiter, requestLimiter: requestLimiter,
messagesLimiter: messagesLimiter,
emailsLimiter: emailsLimiter, emailsLimiter: emailsLimiter,
subscriptionLimiter: util.NewFixedLimiter(int64(conf.VisitorSubscriptionLimit)), subscriptionLimiter: util.NewFixedLimiter(int64(conf.VisitorSubscriptionLimit)),
bandwidthLimiter: util.NewBytesLimiter(conf.VisitorAttachmentDailyBandwidthLimit, 24*time.Hour), bandwidthLimiter: util.NewBytesLimiter(conf.VisitorAttachmentDailyBandwidthLimit, 24*time.Hour),
@ -116,6 +120,13 @@ func (v *visitor) FirebaseTemporarilyDeny() {
v.firebase = time.Now().Add(v.config.FirebaseQuotaExceededPenaltyDuration) v.firebase = time.Now().Add(v.config.FirebaseQuotaExceededPenaltyDuration)
} }
func (v *visitor) MessageAllowed() error {
if v.messagesLimiter != nil && v.messagesLimiter.Allow(1) != nil {
return errVisitorLimitReached
}
return nil
}
func (v *visitor) EmailAllowed() error { func (v *visitor) EmailAllowed() error {
if !v.emailsLimiter.Allow() { if !v.emailsLimiter.Allow() {
return errVisitorLimitReached return errVisitorLimitReached