Add IMAP quota command

connector/connector.go

@@ -10,6 +10,7 @@ import (
 
 var ErrOperationNotAllowed = errors.New("operation not allowed")
 var ErrMessageSizeExceedsLimits = errors.New("message size exceeds limits")
+var ErrOverQuota = errors.New("quota exceeded")
 
 // Connector connects the gluon server to a remote mail store.
 type Connector interface {
@@ -57,6 +58,13 @@ type Connector interface {
 	// MarkMessagesForwarded sets the forwarded value of the give messages.
 	MarkMessagesForwarded(ctx context.Context, cache IMAPStateWrite, messageIDs []imap.MessageID, forwarded bool) error
 
+	// GetQuota returns the quota for the given quota root name.
+	GetQuota(ctx context.Context, rootName string) (*imap.QuotaRoot, error)
+
+	// GetQuotaRoot returns the list of quota root names applicable to the given mailbox,
+	// along with the resource usage for each of those roots.
+	GetQuotaRoot(ctx context.Context, mailboxName string) ([]string, []*imap.QuotaRoot, error)
+
 	// GetUpdates returns a stream of updates that the gluon server should apply.
 	GetUpdates() <-chan imap.Update
 

connector/dummy.go

@@ -59,6 +59,12 @@ type Dummy struct {
 	allowMessageCreateWithUnknownMailboxID bool
 
 	updatesAllowedToFail int32
+
+	// quotaRoots holds configured quota roots for testing.
+	quotaRoots map[string]*imap.QuotaRoot
+
+	// mailboxQuotaRoots maps mailbox names to their applicable quota root names.
+	mailboxQuotaRoots map[string][]string
 }
 
 func NewDummy(usernames []string, password []byte, period time.Duration, flags, permFlags, attrs imap.FlagSet) *Dummy {
@@ -73,6 +79,8 @@ func NewDummy(usernames []string, password []byte, period time.Duration, flags,
 		updateQuitCh:        make(chan struct{}),
 		ticker:              ticker.New(period),
 		mailboxVisibilities: make(map[imap.MailboxID]imap.MailboxVisibility),
+		quotaRoots:          make(map[string]*imap.QuotaRoot),
+		mailboxQuotaRoots:   make(map[string][]string),
 	}
 
 	go func() {
@@ -341,6 +349,48 @@ func (conn *Dummy) SetMailboxVisibility(id imap.MailboxID, visibility imap.Mailb
 	conn.mailboxVisibilities[id] = visibility
 }
 
+func (conn *Dummy) GetQuota(_ context.Context, rootName string) (*imap.QuotaRoot, error) {
+	root, ok := conn.quotaRoots[rootName]
+	if !ok {
+		return nil, fmt.Errorf("no such quota root")
+	}
+
+	return root, nil
+}
+
+func (conn *Dummy) GetQuotaRoot(_ context.Context, mailboxName string) ([]string, []*imap.QuotaRoot, error) {
+	rootNames, ok := conn.mailboxQuotaRoots[mailboxName]
+	if !ok {
+		// Default: return the empty-string root if it exists.
+		if root, exists := conn.quotaRoots[""]; exists {
+			return []string{""}, []*imap.QuotaRoot{root}, nil
+		}
+
+		return nil, nil, nil
+	}
+
+	var roots []*imap.QuotaRoot
+
+	for _, name := range rootNames {
+		if root, exists := conn.quotaRoots[name]; exists {
+			roots = append(roots, root)
+		}
+	}
+
+	return rootNames, roots, nil
+}
+
+func (conn *Dummy) SetQuota(rootName string, resources ...imap.QuotaResource) {
+	conn.quotaRoots[rootName] = &imap.QuotaRoot{
+		RootName:  rootName,
+		Resources: resources,
+	}
+}
+
+func (conn *Dummy) SetMailboxQuotaRoot(mailboxName string, rootNames ...string) {
+	conn.mailboxQuotaRoots[mailboxName] = rootNames
+}
+
 func (conn *Dummy) pushUpdate(update imap.Update) {
 	conn.queueLock.Lock()
 	defer conn.queueLock.Unlock()

imap/capabilities.go

@@ -11,13 +11,14 @@ const (
 	MOVE      Capability = `MOVE`
 	ID        Capability = `ID`
 	AUTHPLAIN Capability = `AUTH=PLAIN`
+	QUOTA     Capability = `QUOTA`
 )
 
 func IsCapabilityAvailableBeforeAuth(c Capability) bool {
 	switch c {
 	case IMAP4rev1, StartTLS, IDLE, ID, AUTHPLAIN:
 		return true
-	case UNSELECT, UIDPLUS, MOVE:
+	case UNSELECT, UIDPLUS, MOVE, QUOTA:
 		return false
 	}
 

imap/command/getquota.go

@@ -0,0 +1,37 @@
+package command
+
+import (
+	"fmt"
+
+	"github.com/ProtonMail/gluon/rfcparser"
+)
+
+type GetQuota struct {
+	Root string
+}
+
+func (g GetQuota) String() string {
+	return fmt.Sprintf("GETQUOTA '%v'", g.Root)
+}
+
+func (g GetQuota) SanitizedString() string {
+	return g.String()
+}
+
+type GetQuotaCommandParser struct{}
+
+func (GetQuotaCommandParser) FromParser(p *rfcparser.Parser) (Payload, error) {
+	// getquota = "GETQUOTA" SP astring
+	if err := p.Consume(rfcparser.TokenTypeSP, "expected space after command"); err != nil {
+		return nil, err
+	}
+
+	root, err := p.ParseAString()
+	if err != nil {
+		return nil, err
+	}
+
+	return &GetQuota{
+		Root: root.Value,
+	}, nil
+}

imap/command/getquotaroot.go

@@ -0,0 +1,37 @@
+package command
+
+import (
+	"fmt"
+
+	"github.com/ProtonMail/gluon/rfcparser"
+)
+
+type GetQuotaRoot struct {
+	Mailbox string
+}
+
+func (g GetQuotaRoot) String() string {
+	return fmt.Sprintf("GETQUOTAROOT '%v'", g.Mailbox)
+}
+
+func (g GetQuotaRoot) SanitizedString() string {
+	return g.String()
+}
+
+type GetQuotaRootCommandParser struct{}
+
+func (GetQuotaRootCommandParser) FromParser(p *rfcparser.Parser) (Payload, error) {
+	// getquotaroot = "GETQUOTAROOT" SP mailbox
+	if err := p.Consume(rfcparser.TokenTypeSP, "expected space after command"); err != nil {
+		return nil, err
+	}
+
+	mailbox, err := ParseMailbox(p)
+	if err != nil {
+		return nil, err
+	}
+
+	return &GetQuotaRoot{
+		Mailbox: mailbox.Value,
+	}, nil
+}

imap/command/parser.go

@@ -60,34 +60,36 @@ func NewParser(s *rfcparser.Scanner, options ...Option) *Parser {
 	}
 
 	commands := map[string]Builder{
-		"list":        &ListCommandParser{},
-		"append":      &AppendCommandParser{},
-		"search":      &SearchCommandParser{},
-		"fetch":       &FetchCommandParser{},
-		"capability":  &CapabilityCommandParser{},
-		"idle":        &IdleCommandParser{},
-		"noop":        &NoopCommandParser{},
-		"logout":      &LogoutCommandParser{},
-		"check":       &CheckCommandParser{},
-		"close":       &CloseCommandParser{},
-		"expunge":     &ExpungeCommandParser{},
-		"unselect":    &UnselectCommandParser{},
-		"starttls":    &StartTLSCommandParser{},
-		"status":      &StatusCommandParser{},
-		"select":      &SelectCommandParser{},
-		"examine":     &ExamineCommandParser{},
-		"create":      &CreateCommandParser{},
-		"delete":      &DeleteCommandParser{},
-		"subscribe":   &SubscribeCommandParser{},
-		"unsubscribe": &UnsubscribeCommandParser{},
-		"rename":      &RenameCommandParser{},
-		"lsub":        &LSubCommandParser{},
-		"login":       &LoginCommandParser{},
-		"store":       &StoreCommandParser{},
-		"copy":        &CopyCommandParser{},
-		"move":        &MoveCommandParser{},
-		"uid":         NewUIDCommandParser(),
-		"id":          &IDCommandParser{},
+		"list":         &ListCommandParser{},
+		"append":       &AppendCommandParser{},
+		"search":       &SearchCommandParser{},
+		"fetch":        &FetchCommandParser{},
+		"capability":   &CapabilityCommandParser{},
+		"idle":         &IdleCommandParser{},
+		"noop":         &NoopCommandParser{},
+		"logout":       &LogoutCommandParser{},
+		"check":        &CheckCommandParser{},
+		"close":        &CloseCommandParser{},
+		"expunge":      &ExpungeCommandParser{},
+		"unselect":     &UnselectCommandParser{},
+		"starttls":     &StartTLSCommandParser{},
+		"status":       &StatusCommandParser{},
+		"select":       &SelectCommandParser{},
+		"examine":      &ExamineCommandParser{},
+		"create":       &CreateCommandParser{},
+		"delete":       &DeleteCommandParser{},
+		"subscribe":    &SubscribeCommandParser{},
+		"unsubscribe":  &UnsubscribeCommandParser{},
+		"rename":       &RenameCommandParser{},
+		"lsub":         &LSubCommandParser{},
+		"login":        &LoginCommandParser{},
+		"store":        &StoreCommandParser{},
+		"copy":         &CopyCommandParser{},
+		"move":         &MoveCommandParser{},
+		"uid":          NewUIDCommandParser(),
+		"id":           &IDCommandParser{},
+		"getquota":     &GetQuotaCommandParser{},
+		"getquotaroot": &GetQuotaRootCommandParser{},
 	}
 
 	if !builder.disableIMAPAuthenticate {

imap/quota.go

@@ -0,0 +1,14 @@
+package imap
+
+// QuotaResource represents a single resource within a quota root (e.g. STORAGE or MESSAGE).
+type QuotaResource struct {
+	ResourceName string
+	Usage        uint32
+	Limit        uint32
+}
+
+// QuotaRoot represents a named quota root and its associated resources.
+type QuotaRoot struct {
+	RootName  string
+	Resources []QuotaResource
+}

internal/backend/state_connector_impl.go

@@ -227,3 +227,15 @@ func (sc *stateConnectorImpl) newContextWithMetadata(ctx context.Context) contex
 
 	return ctx
 }
+
+func (sc *stateConnectorImpl) GetQuota(ctx context.Context, rootName string) (*imap.QuotaRoot, error) {
+	ctx = sc.newContextWithMetadata(ctx)
+
+	return sc.connector.GetQuota(ctx, rootName)
+}
+
+func (sc *stateConnectorImpl) GetQuotaRoot(ctx context.Context, mailboxName string) ([]string, []*imap.QuotaRoot, error) {
+	ctx = sc.newContextWithMetadata(ctx)
+
+	return sc.connector.GetQuotaRoot(ctx, mailboxName)
+}

internal/response/item_overquota.go

@@ -0,0 +1,11 @@
+package response
+
+type itemOverQuota struct{}
+
+func ItemOverQuota() *itemOverQuota {
+	return &itemOverQuota{}
+}
+
+func (c *itemOverQuota) String() string {
+	return "OVERQUOTA"
+}

internal/response/quota.go

@@ -0,0 +1,38 @@
+package response
+
+import (
+	"fmt"
+	"strings"
+
+	"github.com/ProtonMail/gluon/imap"
+)
+
+type quota struct {
+	root *imap.QuotaRoot
+}
+
+func Quota(root *imap.QuotaRoot) *quota {
+	return &quota{root: root}
+}
+
+func (r *quota) Send(s Session) error {
+	return s.WriteResponse(r.String())
+}
+
+func (r *quota) String() string {
+	var resources []string
+
+	for _, res := range r.root.Resources {
+		resources = append(resources, fmt.Sprintf("%v %v %v", res.ResourceName, res.Usage, res.Limit))
+	}
+
+	return fmt.Sprintf(`* QUOTA %v (%v)`, quoteQuotaRoot(r.root.RootName), strings.Join(resources, " "))
+}
+
+func quoteQuotaRoot(name string) string {
+	if name == "" {
+		return `""`
+	}
+
+	return fmt.Sprintf(`"%v"`, name)
+}