瀏覽代碼

Tier based tests

binwiederhier 3 年之前
父節點
當前提交
d8032e1c9e
共有 4 個文件被更改,包括 35 次插入9 次删除
  1. 4 4
      server/server.go
  2. 1 0
      server/server_account_test.go
  3. 16 2
      server/server_test.go
  4. 14 3
      server/visitor.go

+ 4 - 4
server/server.go

@@ -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

+ 1 - 0
server/server_account_test.go

@@ -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"))

+ 16 - 2
server/server_test.go

@@ -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) {

+ 14 - 3
server/visitor.go

@@ -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