fix: Ensure CreateMessage event adds messages to mailboxes

When we receive a create message event from the connector, we need to
ensure that, if it exists, it is still added to all the mailboxes.

This patch also adds a couple of error to detect misuse of some
functions.

internal/backend/connector_updates.go

@@ -202,32 +202,31 @@ func (user *user) applyMessagesCreated(ctx context.Context, update *imap.Message
 
 			internalID, ok := messagesToCreateFilter[message.Message.ID]
 			if !ok {
-				exists, err := db.HasMessageWithRemoteID(ctx, client, message.Message.ID)
-				if err != nil {
-					return err
-				} else if exists {
-					// Message exists, can't be replaced
-					continue
-				}
+				messageID, err := db.GetMessageIDFromRemoteID(ctx, client, message.Message.ID)
+				if ent.IsNotFound(err) {
+					internalID = imap.NewInternalMessageID()
 
-				internalID = imap.NewInternalMessageID()
+					literal, err := rfc822.SetHeaderValue(message.Literal, ids.InternalIDKey, internalID.String())
+					if err != nil {
+						return fmt.Errorf("failed to set internal ID: %w", err)
+					}
 
-				literal, err := rfc822.SetHeaderValue(message.Literal, ids.InternalIDKey, internalID.String())
-				if err != nil {
-					return fmt.Errorf("failed to set internal ID: %w", err)
-				}
+					request := &db.CreateMessageReq{
+						Message:    message.Message,
+						Literal:    literal,
+						Body:       message.ParsedMessage.Body,
+						Structure:  message.ParsedMessage.Structure,
+						Envelope:   message.ParsedMessage.Envelope,
+						InternalID: internalID,
+					}
 
-				request := &db.CreateMessageReq{
-					Message:    message.Message,
-					Literal:    literal,
-					Body:       message.ParsedMessage.Body,
-					Structure:  message.ParsedMessage.Structure,
-					Envelope:   message.ParsedMessage.Envelope,
-					InternalID: internalID,
+					messagesToCreate = append(messagesToCreate, request)
+					messagesToCreateFilter[message.Message.ID] = internalID
+				} else if err == nil {
+					internalID = messageID
+				} else {
+					return err
 				}
-
-				messagesToCreate = append(messagesToCreate, request)
-				messagesToCreateFilter[message.Message.ID] = internalID
 			}
 
 			for _, mboxID := range message.MailboxIDs {
@@ -289,9 +288,20 @@ func (user *user) applyMessagesCreated(ctx context.Context, update *imap.Message
 	return user.db.Write(ctx, func(ctx context.Context, tx *ent.Tx) error {
 		// Assign all the messages to the mailbox
 		for mboxID, msgList := range messageForMBox {
-			if _, err := user.applyMessagesAddedToMailbox(ctx, tx, mboxID, msgList); err != nil {
+			inMailbox, err := db.FilterMailboxContainsInternalID(ctx, tx.Client(), mboxID, msgList)
+			if err != nil {
 				return err
 			}
+
+			toAdd := xslices.Filter(msgList, func(id imap.InternalMessageID) bool {
+				return !slices.Contains(inMailbox, id)
+			})
+
+			if len(toAdd) != 0 {
+				if _, err := user.applyMessagesAddedToMailbox(ctx, tx, mboxID, toAdd); err != nil {
+					return err
+				}
+			}
 		}
 
 		return nil

internal/db/mailbox.go

@@ -367,6 +367,27 @@ func FilterMailboxContains(ctx context.Context, client *ent.Client, mboxID imap.
 	}), nil
 }
 
+func FilterMailboxContainsInternalID(ctx context.Context, client *ent.Client, mboxID imap.InternalMailboxID, messageIDs []imap.InternalMessageID) ([]imap.InternalMessageID, error) {
+	type result struct {
+		InternalID imap.InternalMessageID `json:"uid_message"`
+	}
+
+	var r []result
+
+	if err := client.UID.Query().Where(func(s *sql.Selector) {
+		s.Where(sql.And(sql.EQ(uid.MailboxColumn, mboxID), sql.In(uid.MessageColumn, xslices.Map(messageIDs, func(id imap.InternalMessageID) interface{} {
+			return id
+		})...)))
+		s.Select(uid.MessageColumn)
+	}).Select().Scan(ctx, &r); err != nil {
+		return nil, err
+	}
+
+	return xslices.Map(r, func(r result) imap.InternalMessageID {
+		return r.InternalID
+	}), nil
+}
+
 func GetMailboxFlags(ctx context.Context, client *ent.Client, mboxID imap.InternalMailboxID) (imap.FlagSet, error) {
 	mbox, err := client.Mailbox.Query().Where(mailbox.ID(mboxID)).WithFlags().Only(ctx)
 	if err != nil {
@@ -399,3 +420,10 @@ func GetMailboxAttributes(ctx context.Context, client *ent.Client, mboxID imap.I
 		return flag.Value
 	})), nil
 }
+
+func IsMessageInMailbox(ctx context.Context, client *ent.Client, mboxID imap.InternalMailboxID, messageID imap.InternalMailboxID) (bool, error) {
+	return client.UID.Query().Where(func(s *sql.Selector) {
+		s.Where(sql.And(sql.EQ(uid.MailboxColumn, mboxID), sql.EQ(uid.MessageColumn, messageID)))
+		s.Select(uid.MessageColumn)
+	}).Exist(ctx)
+}

internal/db/message.go

@@ -77,6 +77,10 @@ func CreateMessages(ctx context.Context, tx *ent.Tx, reqs ...*CreateMessageReq)
 }
 
 func AddMessagesToMailbox(ctx context.Context, tx *ent.Tx, messageIDs []imap.InternalMessageID, mboxID imap.InternalMailboxID) ([]UIDWithFlags, error) {
+	if len(messageIDs) == 0 {
+		return nil, nil
+	}
+
 	messageUIDs := make(map[imap.InternalMessageID]imap.UID)
 
 	mbox, err := tx.Mailbox.Query().Where(mailbox.ID(mboxID)).Select(mailbox.FieldUIDNext).Only(ctx)
@@ -99,8 +103,10 @@ func AddMessagesToMailbox(ctx context.Context, tx *ent.Tx, messageIDs []imap.Int
 
 	// Avoid too many SQL variables error.
 	for _, chunk := range xslices.Chunk(builders, ChunkLimit) {
-		if _, err := tx.UID.CreateBulk(chunk...).Save(ctx); err != nil {
+		if m, err := tx.UID.CreateBulk(chunk...).Save(ctx); err != nil {
 			return nil, err
+		} else if len(m) == 0 {
+			return nil, fmt.Errorf("no messages were added to the mailbox")
 		}
 	}
 
@@ -318,6 +324,10 @@ func GetMessageUIDsWithFlagsAfterAddOrUIDBump(ctx context.Context, client *ent.C
 		return v1.UID < v2.UID
 	})
 
+	if len(result) == 0 {
+		return nil, fmt.Errorf("result can never be null after UID bump")
+	}
+
 	return result, nil
 }