feat(GODT-1970): Soft limits

builder.go

@@ -2,6 +2,7 @@ package gluon
 
 import (
 	"crypto/tls"
+	"github.com/ProtonMail/gluon/limits"
 	"io"
 	"os"
 	"path/filepath"
@@ -29,6 +30,7 @@ type serverBuilder struct {
 	storeBuilder       store.Builder
 	reporter           reporter.Reporter
 	disableParallelism bool
+	imapLimits         limits.IMAP
 }
 
 func newBuilder() (*serverBuilder, error) {
@@ -37,7 +39,8 @@ func newBuilder() (*serverBuilder, error) {
 		cmdExecProfBuilder: &profiling.NullCmdExecProfilerBuilder{},
 		storeBuilder:       &store.OnDiskStoreBuilder{},
 		reporter:           &reporter.NullReporter{},
-		idleBulkTime:       time.Duration(500 * time.Millisecond),
+		idleBulkTime:       500 * time.Millisecond,
+		imapLimits:         limits.DefaultLimits(),
 	}, nil
 }
 
@@ -60,6 +63,7 @@ func (builder *serverBuilder) build() (*Server, error) {
 		builder.storeBuilder,
 		builder.delim,
 		builder.loginJailTime,
+		builder.imapLimits,
 	)
 	if err != nil {
 		return nil, err

internal/backend/backend.go

@@ -3,6 +3,7 @@ package backend
 import (
 	"context"
 	"fmt"
+	"github.com/ProtonMail/gluon/limits"
 	"path/filepath"
 	"sync"
 	"sync/atomic"
@@ -44,15 +45,18 @@ type Backend struct {
 	loginErrorCount int32
 	loginLock       sync.Mutex
 	loginWG         sync.WaitGroup
+
+	imapLimits limits.IMAP
 }
 
-func New(dir string, storeBuilder store.Builder, delim string, loginJailTime time.Duration) (*Backend, error) {
+func New(dir string, storeBuilder store.Builder, delim string, loginJailTime time.Duration, imapLimits limits.IMAP) (*Backend, error) {
 	return &Backend{
 		dir:           dir,
 		delim:         delim,
 		users:         make(map[string]*user),
 		storeBuilder:  storeBuilder,
 		loginJailTime: loginJailTime,
+		imapLimits:    imapLimits,
 	}, nil
 }
 
@@ -82,7 +86,7 @@ func (b *Backend) AddUser(ctx context.Context, userID string, conn connector.Con
 		return err
 	}
 
-	user, err := newUser(ctx, userID, db, conn, storeBuilder, b.delim)
+	user, err := newUser(ctx, userID, db, conn, storeBuilder, b.delim, b.imapLimits)
 	if err != nil {
 		return err
 	}

internal/backend/connector_updates.go

@@ -71,7 +71,17 @@ func (user *user) applyMailboxCreated(ctx context.Context, update *imap.MailboxC
 		return fmt.Errorf("attempting to create protected mailbox (recovery)")
 	}
 
+	if err := user.imapLimits.CheckUIDValidity(user.globalUIDValidity); err != nil {
+		return err
+	}
+
 	if exists, err := db.ReadResult(ctx, user.db, func(ctx context.Context, client *ent.Client) (bool, error) {
+		if mailboxCount, err := db.GetMailboxCount(ctx, client); err != nil {
+			return false, err
+		} else if err := user.imapLimits.CheckMailBoxCount(mailboxCount); err != nil {
+			return false, err
+		}
+
 		return db.MailboxExistsWithRemoteID(ctx, client, update.Mailbox.ID)
 	}); err != nil {
 		return err
@@ -426,7 +436,7 @@ func (user *user) setMessageMailboxes(ctx context.Context, tx *ent.Tx, messageID
 
 // applyMessagesAddedToMailbox adds the messages to the given mailbox.
 func (user *user) applyMessagesAddedToMailbox(ctx context.Context, tx *ent.Tx, mboxID imap.InternalMailboxID, messageIDs []imap.InternalMessageID) ([]db.UIDWithFlags, error) {
-	messageUIDs, update, err := state.AddMessagesToMailbox(ctx, tx, mboxID, messageIDs, nil)
+	messageUIDs, update, err := state.AddMessagesToMailbox(ctx, tx, mboxID, messageIDs, nil, user.imapLimits)
 	if err != nil {
 		return nil, err
 	}
@@ -636,7 +646,7 @@ func (user *user) applyMessageUpdated(ctx context.Context, update *imap.MessageU
 						return err
 					}
 
-					_, update, err := state.AddMessagesToMailbox(ctx, tx, internalMBoxID, []imap.InternalMessageID{newInternalID}, nil)
+					_, update, err := state.AddMessagesToMailbox(ctx, tx, internalMBoxID, []imap.InternalMessageID{newInternalID}, nil, user.imapLimits)
 					if err != nil {
 						return err
 					}

internal/backend/user.go

@@ -3,6 +3,7 @@ package backend
 import (
 	"context"
 	"fmt"
+	"github.com/ProtonMail/gluon/limits"
 	"sync"
 
 	"github.com/ProtonMail/gluon/connector"
@@ -39,6 +40,8 @@ type user struct {
 	globalUIDValidity imap.UID
 
 	recoveryMailboxID imap.InternalMailboxID
+
+	imapLimits limits.IMAP
 }
 
 func newUser(
@@ -48,6 +51,7 @@ func newUser(
 	conn connector.Connector,
 	store store.Store,
 	delimiter string,
+	imapLimits limits.IMAP,
 ) (*user, error) {
 	if err := database.Init(ctx); err != nil {
 		return nil, err
@@ -85,6 +89,8 @@ func newUser(
 		globalUIDValidity: conn.GetUIDValidity(),
 
 		recoveryMailboxID: recoveryMBox.ID,
+
+		imapLimits: imapLimits,
 	}
 
 	if err := user.deleteAllMessagesMarkedDeleted(ctx); err != nil {
@@ -209,6 +215,7 @@ func (user *user) newState() (*state.State, error) {
 	newState := state.NewState(
 		newStateUserInterfaceImpl(user, newStateConnectorImpl(user)),
 		user.delimiter,
+		user.imapLimits,
 	)
 
 	user.states[newState.StateID] = newState

internal/db/mailbox.go

@@ -427,3 +427,30 @@ func IsMessageInMailbox(ctx context.Context, client *ent.Client, mboxID imap.Int
 		s.Select(uid.MessageColumn)
 	}).Exist(ctx)
 }
+
+func GetMailboxCount(ctx context.Context, client *ent.Client) (int, error) {
+	return client.Mailbox.Query().Count(ctx)
+}
+
+func GetMailboxUID(ctx context.Context, client *ent.Client, mboxID imap.InternalMailboxID) (imap.UID, error) {
+	mbox, err := client.Mailbox.Query().Where(mailbox.ID(mboxID)).Select(mailbox.FieldUIDNext).Only(ctx)
+	if err != nil {
+		return 0, err
+	}
+
+	return mbox.UIDNext, err
+}
+
+func GetMailboxMessageCountAndUID(ctx context.Context, client *ent.Client, mboxID imap.InternalMailboxID) (int, imap.UID, error) {
+	messageCount, err := GetMailboxMessageCount(ctx, client, mboxID)
+	if err != nil {
+		return 0, 0, err
+	}
+
+	uid, err := GetMailboxUID(ctx, client, mboxID)
+	if err != nil {
+		return 0, 0, err
+	}
+
+	return messageCount, uid, err
+}

internal/session/session.go

@@ -18,6 +18,7 @@ import (
 	"github.com/ProtonMail/gluon/internal/backend"
 	"github.com/ProtonMail/gluon/internal/response"
 	"github.com/ProtonMail/gluon/internal/state"
+	"github.com/ProtonMail/gluon/limits"
 	"github.com/ProtonMail/gluon/liner"
 	"github.com/ProtonMail/gluon/profiling"
 	"github.com/ProtonMail/gluon/reporter"
@@ -82,6 +83,8 @@ type Session struct {
 
 	/// errorCount error counter
 	errorCount int
+
+	imapLimits limits.IMAP
 }
 
 func New(

internal/state/actions.go

@@ -262,7 +262,7 @@ func (state *State) actionAddMessagesToMailbox(
 		st = state
 	}
 
-	messageUIDs, update, err := AddMessagesToMailbox(ctx, tx, mboxID.InternalID, internalIDs, st)
+	messageUIDs, update, err := AddMessagesToMailbox(ctx, tx, mboxID.InternalID, internalIDs, st, state.imapLimits)
 	if err != nil {
 		return nil, err
 	}
@@ -286,7 +286,7 @@ func (state *State) actionAddRecoveredMessagesToMailbox(
 		return nil, nil, err
 	}
 
-	return AddMessagesToMailbox(ctx, tx, mboxID.InternalID, internalIDs, state)
+	return AddMessagesToMailbox(ctx, tx, mboxID.InternalID, internalIDs, state, state.imapLimits)
 }
 
 func (state *State) actionImportRecoveredMessage(
@@ -536,7 +536,7 @@ func (state *State) actionMoveMessages(
 		return nil, err
 	}
 
-	messageUIDs, updates, err := MoveMessagesFromMailbox(ctx, tx, mboxFromID.InternalID, mboxToID.InternalID, internalIDs, state)
+	messageUIDs, updates, err := MoveMessagesFromMailbox(ctx, tx, mboxFromID.InternalID, mboxToID.InternalID, internalIDs, state, state.imapLimits)
 	if err != nil {
 		return nil, err
 	}

internal/state/mailbox.go

@@ -146,6 +146,24 @@ func (m *Mailbox) GetMessagesWithoutFlagCount(flag string) int {
 }
 
 func (m *Mailbox) AppendRegular(ctx context.Context, literal []byte, flags imap.FlagSet, date time.Time) (imap.UID, error) {
+	if err := m.state.db().Read(ctx, func(ctx context.Context, client *ent.Client) error {
+		if messageCount, uid, err := db.GetMailboxMessageCountAndUID(ctx, client, m.snap.mboxID.InternalID); err != nil {
+			return err
+		} else {
+			if err := m.state.imapLimits.CheckMailBoxMessageCount(messageCount, 1); err != nil {
+				return err
+			}
+
+			if err := m.state.imapLimits.CheckUIDCount(uid, 1); err != nil {
+				return err
+			}
+		}
+
+		return nil
+	}); err != nil {
+		return 0, err
+	}
+
 	var appendIntoDrafts bool
 
 	attr, err := m.Attributes(ctx)

internal/state/state.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"github.com/ProtonMail/gluon/limits"
 	"strings"
 	"sync/atomic"
 
@@ -45,6 +46,8 @@ type State struct {
 
 	// invalid indicates whether this state became invalid and a clients needs to disconnect.
 	invalid bool
+
+	imapLimits limits.IMAP
 }
 
 var stateIDGenerator int64
@@ -53,14 +56,15 @@ func nextStateID() StateID {
 	return StateID(atomic.AddInt64(&stateIDGenerator, 1))
 }
 
-func NewState(user UserInterface, delimiter string) *State {
+func NewState(user UserInterface, delimiter string, imapLimits limits.IMAP) *State {
 	return &State{
 		user:         user,
 		StateID:      nextStateID(),
 		doneCh:       make(chan struct{}),
 		snap:         nil,
 		delimiter:    delimiter,
 		updatesQueue: queue.NewQueuedChannel[Update](32, 128),
+		imapLimits:   imapLimits,
 	}
 }
 
@@ -164,6 +168,10 @@ func (state *State) Examine(ctx context.Context, name string, fn func(*Mailbox)
 }
 
 func (state *State) Create(ctx context.Context, name string) error {
+	if err := state.imapLimits.CheckUIDValidity(state.user.GetGlobalUIDValidity()); err != nil {
+		return err
+	}
+
 	if strings.HasPrefix(strings.ToLower(name), ids.GluonRecoveryMailboxNameLowerCase) {
 		return fmt.Errorf("operation not allowed")
 	}
@@ -179,6 +187,12 @@ func (state *State) Create(ctx context.Context, name string) error {
 	}
 
 	mboxesToCreate, err := db.ReadResult(ctx, state.db(), func(ctx context.Context, client *ent.Client) ([]string, error) {
+		if mailboxCount, err := db.GetMailboxCount(ctx, client); err != nil {
+			return nil, err
+		} else if err := state.imapLimits.CheckMailBoxCount(mailboxCount); err != nil {
+			return nil, err
+		}
+
 		var mboxesToCreate []string
 		// If the mailbox name is suffixed with the server's hierarchy separator, remove the separator and still create
 		// the mailbox

internal/state/updates_mailbox.go

@@ -2,6 +2,7 @@ package state
 
 import (
 	"context"
+	"github.com/ProtonMail/gluon/limits"
 
 	"github.com/ProtonMail/gluon/imap"
 	"github.com/ProtonMail/gluon/internal/db"
@@ -19,7 +20,21 @@ func MoveMessagesFromMailbox(
 	mboxFromID, mboxToID imap.InternalMailboxID,
 	messageIDs []imap.InternalMessageID,
 	s *State,
+	imapLimits limits.IMAP,
 ) ([]db.UIDWithFlags, []Update, error) {
+	messageCount, uid, err := db.GetMailboxMessageCountAndUID(ctx, tx.Client(), mboxToID)
+	if err != nil {
+		return nil, nil, err
+	}
+
+	if err := imapLimits.CheckMailBoxMessageCount(messageCount, len(messageIDs)); err != nil {
+		return nil, nil, err
+	}
+
+	if err := imapLimits.CheckUIDCount(uid, len(messageIDs)); err != nil {
+		return nil, nil, err
+	}
+
 	if mboxFromID != mboxToID {
 		if err := db.RemoveMessagesFromMailbox(ctx, tx, messageIDs, mboxFromID); err != nil {
 			return nil, nil, err
@@ -50,7 +65,20 @@ func MoveMessagesFromMailbox(
 }
 
 // AddMessagesToMailbox adds the messages to the given mailbox.
-func AddMessagesToMailbox(ctx context.Context, tx *ent.Tx, mboxID imap.InternalMailboxID, messageIDs []imap.InternalMessageID, s *State) ([]db.UIDWithFlags, Update, error) {
+func AddMessagesToMailbox(ctx context.Context, tx *ent.Tx, mboxID imap.InternalMailboxID, messageIDs []imap.InternalMessageID, s *State, imapLimits limits.IMAP) ([]db.UIDWithFlags, Update, error) {
+	messageCount, uid, err := db.GetMailboxMessageCountAndUID(ctx, tx.Client(), mboxID)
+	if err != nil {
+		return nil, nil, err
+	}
+
+	if err := imapLimits.CheckMailBoxMessageCount(messageCount, len(messageIDs)); err != nil {
+		return nil, nil, err
+	}
+
+	if err := imapLimits.CheckUIDCount(uid, len(messageIDs)); err != nil {
+		return nil, nil, err
+	}
+
 	messageUIDs, err := db.AddMessagesToMailbox(ctx, tx, messageIDs, mboxID)
 	if err != nil {
 		return nil, nil, err