feat(GODT-1970): Soft limits (#242)

Add limit checks for mailbox count, messages per mailbox, UID and
UIDValidity so that we don't run into issues related exceeding the
maximum representable size for each of these data types (uint32).

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