fix(GODT-2156): Use UUID instead of Integers for internal ID

Avoids collision issues if users try to copy messages between different Gluon based servers.
ProtonMail · Nov 25, 2022 · c337691 · c337691
1 parent 33e01d9
commit c337691
Show file tree

Hide file tree

Showing 32 changed files with 205 additions and 143 deletions.
diff --git a/benchmarks/gluon_bench/store_benchmarks/create.go b/benchmarks/gluon_bench/store_benchmarks/create.go
@@ -31,7 +31,7 @@ func (*Create) Run(ctx context.Context, st store.Store) (*reporter.BenchmarkRun,
 
 		for i := uint(0); i < *flags.StoreItemCount; i++ {
 			dc.Start()
-			err := s.Set(imap.InternalMessageID(uint64(i)), data)
+			err := s.Set(imap.NewInternalMessageID(), data)
 			dc.Stop()
 
 			if err != nil {

diff --git a/benchmarks/gluon_bench/store_benchmarks/utils.go b/benchmarks/gluon_bench/store_benchmarks/utils.go
@@ -17,7 +17,7 @@ func CreateRandomState(st store.Store, count uint) ([]imap.InternalMessageID, er
 	data := make([]byte, *flags.StoreItemSize)
 
 	for i := uint(0); i < count; i++ {
-		uuid := imap.InternalMessageID(uint64(i))
+		uuid := imap.NewInternalMessageID()
 
 		if err := st.Set(uuid, data); err != nil {
 			return nil, err

diff --git a/imap/strong_types.go b/imap/strong_types.go
@@ -1,11 +1,11 @@
 package imap
 
 import (
-	"encoding/binary"
 	"fmt"
 	"strconv"
 
 	"github.com/ProtonMail/gluon/internal/utils"
+	"github.com/google/uuid"
 )
 
 type MailboxID string
@@ -20,7 +20,9 @@ func (m MessageID) ShortID() string {
 	return utils.ShortID(string(m))
 }
 
-type InternalMessageID uint64
+type InternalMessageID struct {
+	uuid.UUID
+}
 
 type InternalMailboxID uint64
 
@@ -29,31 +31,28 @@ func (i InternalMailboxID) ShortID() string {
 }
 
 func (i InternalMessageID) ShortID() string {
-	return strconv.FormatUint(uint64(i), 10)
+	return utils.ShortID(i.String())
 }
 
 func (i InternalMailboxID) String() string {
 	return strconv.FormatUint(uint64(i), 10)
 }
 
 func (i InternalMessageID) String() string {
-	return strconv.FormatUint(uint64(i), 10)
+	return i.UUID.String()
+}
+
+func NewInternalMessageID() InternalMessageID {
+	return InternalMessageID{UUID: uuid.New()}
 }
 
 func InternalMessageIDFromString(id string) (InternalMessageID, error) {
-	num, err := strconv.ParseUint(id, 10, 64)
+	num, err := uuid.Parse(id)
 	if err != nil {
-		return 0, fmt.Errorf("invalid message id:%w", err)
+		return InternalMessageID{}, fmt.Errorf("invalid message id:%w", err)
 	}
 
-	return InternalMessageID(num), nil
-}
-
-func (i InternalMessageID) ToBytes() []byte {
-	bytes := make([]byte, 8)
-	binary.LittleEndian.PutUint64(bytes, uint64(i))
-
-	return bytes
+	return InternalMessageID{UUID: num}, nil
 }
 
 type UID uint32

diff --git a/internal/backend/connector_updates.go b/internal/backend/connector_updates.go
@@ -210,7 +210,7 @@ func (user *user) applyMessagesCreated(ctx context.Context, update *imap.Message
 					continue
 				}
 
-				internalID = user.nextMessageID()
+				internalID = imap.NewInternalMessageID()
 
 				literal, err := rfc822.SetHeaderValue(message.Literal, ids.InternalIDKey, internalID.String())
 				if err != nil {
@@ -597,7 +597,7 @@ func (user *user) applyMessageUpdated(ctx context.Context, update *imap.MessageU
 			}
 			// create new entry
 			{
-				newInternalID := user.nextMessageID()
+				newInternalID := imap.NewInternalMessageID()
 				literal, err := rfc822.SetHeaderValue(update.Literal, ids.InternalIDKey, newInternalID.String())
 				if err != nil {
 					return fmt.Errorf("failed to set internal ID: %w", err)

diff --git a/internal/backend/state_connector_impl.go b/internal/backend/state_connector_impl.go
@@ -69,10 +69,10 @@ func (sc *stateConnectorImpl) CreateMessage(
 
 	msg, newLiteral, err := sc.connector.CreateMessage(ctx, mboxID, literal, flags, date)
 	if err != nil {
-		return 0, imap.Message{}, nil, err
+		return imap.InternalMessageID{}, imap.Message{}, nil, err
 	}
 
-	return sc.user.nextMessageID(), msg, newLiteral, nil
+	return imap.NewInternalMessageID(), msg, newLiteral, nil
 }
 
 func (sc *stateConnectorImpl) AddMessagesToMailbox(

diff --git a/internal/backend/state_user_interface_impl.go b/internal/backend/state_user_interface_impl.go
@@ -93,7 +93,3 @@ func (s *StateUserInterfaceImpl) GetRecoveryMailboxID() ids.MailboxIDPair {
 		RemoteID:   ids.GluonInternalRecoveryMailboxRemoteID,
 	}
 }
-
-func (s *StateUserInterfaceImpl) NextRecoveryMessageID() imap.InternalMessageID {
-	return s.u.nextMessageID()
-}
diff --git a/internal/backend/user.go b/internal/backend/user.go
@@ -4,7 +4,6 @@ import (
 	"context"
 	"fmt"
 	"sync"
-	"sync/atomic"
 
 	"github.com/ProtonMail/gluon/connector"
 	"github.com/ProtonMail/gluon/imap"
@@ -37,7 +36,6 @@ type user struct {
 	updateWG     sync.WaitGroup
 	updateQuitCh chan struct{}
 
-	messageIDCounter  uint64
 	globalUIDValidity imap.UID
 
 	recoveryMailboxID imap.InternalMailboxID
@@ -68,15 +66,6 @@ func newUser(
 
 		return db.GetOrCreateMailbox(ctx, tx, mbox, delimiter, conn.GetUIDValidity())
 	})
-
-	if err != nil {
-		return nil, err
-	}
-
-	// Get the last message ID from the database so we can resume our counter properly.
-	highestMessageID, err := db.ReadResult(ctx, database, func(ctx context.Context, client *ent.Client) (imap.InternalMessageID, error) {
-		return db.GetHighestMessageID(ctx, client)
-	})
 	if err != nil {
 		return nil, err
 	}
@@ -93,7 +82,6 @@ func newUser(
 
 		states:            make(map[state.StateID]*state.State),
 		updateQuitCh:      make(chan struct{}),
-		messageIDCounter:  uint64(highestMessageID),
 		globalUIDValidity: conn.GetUIDValidity(),
 
 		recoveryMailboxID: recoveryMBox.ID,
@@ -311,10 +299,6 @@ func (user *user) closeStates() {
 	}
 }
 
-func (user *user) nextMessageID() imap.InternalMessageID {
-	return imap.InternalMessageID(atomic.AddUint64(&user.messageIDCounter, 1))
-}
-
 func (user *user) cleanupStaleStoreData(ctx context.Context) error {
 	storeIds, err := user.store.List()
 	if err != nil {

diff --git a/internal/db/ent/message.go b/internal/db/ent/message.go
diff --git a/internal/db/ent/message_create.go b/internal/db/ent/message_create.go
diff --git a/internal/db/ent/message_delete.go b/internal/db/ent/message_delete.go
diff --git a/internal/db/ent/message_query.go b/internal/db/ent/message_query.go
diff --git a/internal/db/ent/message_update.go b/internal/db/ent/message_update.go
diff --git a/internal/db/ent/messageflag.go b/internal/db/ent/messageflag.go
diff --git a/internal/db/ent/messageflag_create.go b/internal/db/ent/messageflag_create.go
diff --git a/internal/db/ent/messageflag_update.go b/internal/db/ent/messageflag_update.go
diff --git a/internal/db/ent/migrate/schema.go b/internal/db/ent/migrate/schema.go
diff --git a/internal/db/ent/schema/message.go b/internal/db/ent/schema/message.go
@@ -17,7 +17,7 @@ type Message struct {
 // Fields of the Message.
 func (Message) Fields() []ent.Field {
 	return []ent.Field{
-		field.Uint64("id").GoType(imap.InternalMessageID(0)).Unique().Immutable(),
+		field.UUID("id", imap.NewInternalMessageID()).Unique().Immutable(),
 		field.String("RemoteID").Optional().Unique().GoType(imap.MessageID("")),
 		field.Time("Date"),
 		field.Int("Size"),