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

@ -41,7 +41,7 @@ import (
plan: plan:
weirdness with admin and "default" account weirdness with admin and "default" account
v.Info() endpoint double selects from DB v.Info() endpoint double selects from DB
purge accounts that were not logged into in X purge accounts that were not logged int o in X
reset daily limits for users reset daily limits for users
Make sure account endpoints make sense for admins Make sure account endpoints make sense for admins
add logic to set "expires" column (this is gonna be dirty) add logic to set "expires" column (this is gonna be dirty)
@ -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