Tier based tests

pull/584/head
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:
weirdness with admin and "default" account
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
Make sure account endpoints make sense for admins
add logic to set "expires" column (this is gonna be dirty)
@ -55,8 +55,6 @@ import (
- figure out what settings are "web" or "phone"
Tests:
- visitor with/without user
- plan-based message expiry
- plan-based attachment expiry
Docs:
- "expires" field in message
Refactor:
@ -65,7 +63,6 @@ import (
- Password reset
- Pricing
- change email
*/
// 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 {
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)
if err != nil {
return nil, err

View File

@ -462,6 +462,7 @@ func TestAccount_Reservation_PublishByAnonymousFails(t *testing.T) {
require.Nil(t, s.userManager.CreateTier(&user.Tier{
Code: "pro",
MessagesLimit: 20,
ReservationsLimit: 2,
}))
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{
Code: "test",
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.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{
"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) {
@ -1318,6 +1326,7 @@ func TestServer_PublishAttachmentWithTierBasedExpiry(t *testing.T) {
sevenDaysInSeconds := int64(604800)
require.Nil(t, s.userManager.CreateTier(&user.Tier{
Code: "test",
MessagesLimit: 10,
MessagesExpiryDuration: sevenDaysInSeconds,
AttachmentFileSizeLimit: 50_000,
AttachmentTotalSizeLimit: 200_000,
@ -1362,8 +1371,10 @@ func TestServer_PublishAttachmentWithTierBasedLimits(t *testing.T) {
// Create tier with certain limits
require.Nil(t, s.userManager.CreateTier(&user.Tier{
Code: "test",
MessagesLimit: 100,
AttachmentFileSizeLimit: 50_000,
AttachmentTotalSizeLimit: 200_000,
AttachmentExpiryDuration: 30,
}))
require.Nil(t, s.userManager.AddUser("phil", "phil", user.RoleUser))
require.Nil(t, s.userManager.ChangeTier("phil", "test"))
@ -1377,12 +1388,14 @@ func TestServer_PublishAttachmentWithTierBasedLimits(t *testing.T) {
// Publish large file as anonymous
response = request(t, s, "PUT", "/mytopic", largeFile, nil)
require.Equal(t, 413, response.Code)
require.Equal(t, 41301, toHTTPError(t, response.Body.String()).Code)
// Publish too large file as phil
response = request(t, s, "PUT", "/mytopic", largeFile+" a few more bytes", map[string]string{
"Authorization": util.BasicAuth("phil", "phil"),
})
require.Equal(t, 413, response.Code)
require.Equal(t, 41301, toHTTPError(t, response.Body.String()).Code)
// Publish large file as phil (4x)
for i := 0; i < 4; i++ {
@ -1398,6 +1411,7 @@ func TestServer_PublishAttachmentWithTierBasedLimits(t *testing.T) {
"Authorization": util.BasicAuth("phil", "phil"),
})
require.Equal(t, 413, response.Code)
require.Equal(t, 41301, toHTTPError(t, response.Body.String()).Code)
}
func TestServer_PublishAttachmentBandwidthLimit(t *testing.T) {

View File

@ -29,12 +29,13 @@ type visitor struct {
userManager *user.Manager // May be nil!
ip netip.Addr
user *user.User
messages int64 // Number of messages sent
emails int64 // Number of emails sent
messages int64 // Number of messages sent, reset every day
emails int64 // Number of emails sent, reset every day
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
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
firebase time.Time // Next allowed Firebase message
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 {
var messagesLimiter util.Limiter
var requestLimiter, emailsLimiter, accountLimiter *rate.Limiter
var messages, emails int64
if user != nil {
@ -71,6 +73,7 @@ func newVisitor(conf *Config, messageCache *messageCache, userManager *user.Mana
}
if user != nil && user.Tier != nil {
requestLimiter = rate.NewLimiter(dailyLimitToRate(user.Tier.MessagesLimit), conf.VisitorRequestLimitBurst)
messagesLimiter = util.NewFixedLimiter(user.Tier.MessagesLimit)
emailsLimiter = rate.NewLimiter(dailyLimitToRate(user.Tier.EmailsLimit), conf.VisitorEmailLimitBurst)
} else {
requestLimiter = rate.NewLimiter(rate.Every(conf.VisitorRequestLimitReplenish), conf.VisitorRequestLimitBurst)
@ -85,6 +88,7 @@ func newVisitor(conf *Config, messageCache *messageCache, userManager *user.Mana
messages: messages,
emails: emails,
requestLimiter: requestLimiter,
messagesLimiter: messagesLimiter,
emailsLimiter: emailsLimiter,
subscriptionLimiter: util.NewFixedLimiter(int64(conf.VisitorSubscriptionLimit)),
bandwidthLimiter: util.NewBytesLimiter(conf.VisitorAttachmentDailyBandwidthLimit, 24*time.Hour),
@ -116,6 +120,13 @@ func (v *visitor) FirebaseTemporarilyDeny() {
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 {
if !v.emailsLimiter.Allow() {
return errVisitorLimitReached