Open-IM-Server/pkg/common/storage/database/mgo/crypto.go

package mgo

import (
	"context"
	"time"

	"github.com/openimsdk/open-im-server/v3/pkg/common/storage/database"
	"github.com/openimsdk/open-im-server/v3/pkg/common/storage/model"
	"github.com/openimsdk/tools/errs"
	"go.mongodb.org/mongo-driver/bson"
	"go.mongodb.org/mongo-driver/mongo"
	"go.mongodb.org/mongo-driver/mongo/options"
)

// ---- CryptoDevice ----

type CryptoDeviceMgo struct {
	coll *mongo.Collection
}

func NewCryptoDeviceMongo(db *mongo.Database) (database.CryptoDevice, error) {
	coll := db.Collection("crypto_device")
	_, err := coll.Indexes().CreateMany(context.Background(), []mongo.IndexModel{
		{
			Keys:    bson.D{{Key: "user_id", Value: 1}, {Key: "device_id", Value: 1}},
			Options: options.Index().SetUnique(true),
		},
		{
			Keys: bson.D{{Key: "user_id", Value: 1}},
		},
	})
	if err != nil {
		return nil, err
	}
	return &CryptoDeviceMgo{coll: coll}, nil
}

func (m *CryptoDeviceMgo) Create(ctx context.Context, device *model.CryptoDevice) error {
	_, err := m.coll.InsertOne(ctx, device)
	return err
}

func (m *CryptoDeviceMgo) FindByUserID(ctx context.Context, userID string) ([]*model.CryptoDevice, error) {
	cursor, err := m.coll.Find(ctx, bson.M{"user_id": userID})
	if err != nil {
		return nil, err
	}
	var devices []*model.CryptoDevice
	if err := cursor.All(ctx, &devices); err != nil {
		return nil, err
	}
	return devices, nil
}

func (m *CryptoDeviceMgo) FindByUserIDAndDeviceID(ctx context.Context, userID, deviceID string) (*model.CryptoDevice, error) {
	var device model.CryptoDevice
	err := m.coll.FindOne(ctx, bson.M{"user_id": userID, "device_id": deviceID}).Decode(&device)
	if err != nil {
		if err == mongo.ErrNoDocuments {
			return nil, errs.ErrRecordNotFound.WrapMsg("crypto device not found", "userID", userID, "deviceID", deviceID)
		}
		return nil, err
	}
	return &device, nil
}

func (m *CryptoDeviceMgo) UpdateStatus(ctx context.Context, userID, deviceID, status string) error {
	result, err := m.coll.UpdateOne(ctx,
		bson.M{"user_id": userID, "device_id": deviceID},
		bson.M{"$set": bson.M{"status": status}},
	)
	if err != nil {
		return err
	}
	if result.MatchedCount == 0 {
		return errs.ErrRecordNotFound.WrapMsg("crypto device not found", "userID", userID, "deviceID", deviceID)
	}
	return nil
}

func (m *CryptoDeviceMgo) UpdateLastSeen(ctx context.Context, userID, deviceID string) error {
	result, err := m.coll.UpdateOne(ctx,
		bson.M{"user_id": userID, "device_id": deviceID},
		bson.M{"$set": bson.M{"last_seen_at": time.Now()}},
	)
	if err != nil {
		return err
	}
	if result.MatchedCount == 0 {
		return errs.ErrRecordNotFound.WrapMsg("crypto device not found", "userID", userID, "deviceID", deviceID)
	}
	return nil
}

// ---- GroupKeyVersion ----

type GroupKeyVersionMgo struct {
	coll *mongo.Collection
}

func NewGroupKeyVersionMongo(db *mongo.Database) (database.GroupKeyVersion, error) {
	coll := db.Collection("group_key_version")
	_, err := coll.Indexes().CreateOne(context.Background(), mongo.IndexModel{
		Keys:    bson.D{{Key: "group_id", Value: 1}},
		Options: options.Index().SetUnique(true),
	})
	if err != nil {
		return nil, err
	}
	return &GroupKeyVersionMgo{coll: coll}, nil
}

func (m *GroupKeyVersionMgo) Find(ctx context.Context, groupID string) (*model.GroupKeyVersion, error) {
	var v model.GroupKeyVersion
	err := m.coll.FindOne(ctx, bson.M{"group_id": groupID}).Decode(&v)
	if err != nil {
		if err == mongo.ErrNoDocuments {
			return nil, errs.ErrRecordNotFound.WrapMsg("group key version not found", "groupID", groupID)
		}
		return nil, err
	}
	return &v, nil
}

func (m *GroupKeyVersionMgo) IncrVersion(ctx context.Context, groupID string) (int64, error) {
	var result model.GroupKeyVersion
	err := m.coll.FindOneAndUpdate(ctx,
		bson.M{"group_id": groupID},
		bson.M{"$inc": bson.M{"group_key_version": int64(1)}},
		options.FindOneAndUpdate().SetUpsert(true).SetReturnDocument(options.After),
	).Decode(&result)
	if err != nil {
		return 0, err
	}
	return result.GroupKeyVersion, nil
}

// ---- GroupKeyEvent ----

type GroupKeyEventMgo struct {
	coll *mongo.Collection
}

const maxGroupKeyEventsPerQuery = 500

func NewGroupKeyEventMongo(db *mongo.Database) (database.GroupKeyEvent, error) {
	coll := db.Collection("group_key_event")
	_, err := coll.Indexes().CreateMany(context.Background(), []mongo.IndexModel{
		{
			Keys: bson.D{{Key: "group_id", Value: 1}, {Key: "group_key_version", Value: 1}},
		},
		{
			Keys:    bson.D{{Key: "event_id", Value: 1}},
			Options: options.Index().SetUnique(true),
		},
	})
	if err != nil {
		return nil, err
	}
	return &GroupKeyEventMgo{coll: coll}, nil
}

func (m *GroupKeyEventMgo) Create(ctx context.Context, event *model.GroupKeyEvent) error {
	_, err := m.coll.InsertOne(ctx, event)
	return err
}

func (m *GroupKeyEventMgo) FindSinceVersion(ctx context.Context, groupID string, sinceVersion int64) ([]*model.GroupKeyEvent, error) {
	cursor, err := m.coll.Find(ctx, bson.M{
		"group_id":          groupID,
		"group_key_version": bson.M{"$gt": sinceVersion},
	}, options.Find().
		SetSort(bson.D{{Key: "group_key_version", Value: 1}}).
		SetLimit(maxGroupKeyEventsPerQuery))
	if err != nil {
		return nil, err
	}
	var events []*model.GroupKeyEvent
	if err := cursor.All(ctx, &events); err != nil {
		return nil, err
	}
	return events, nil
}