package repository import ( "context" "errors" "time" libmongo "gateway/internal/library/mongo" member "gateway/internal/model/member/domain" "gateway/internal/model/member/domain/entity" domrepo "gateway/internal/model/member/domain/repository" "go.mongodb.org/mongo-driver/v2/bson" mongodriver "go.mongodb.org/mongo-driver/v2/mongo" ) // MongoTOTPProfileRepository stores TOTP fields on the members document. type MongoTOTPProfileRepository struct { db libmongo.DocumentDBUseCase } // NewMongoTOTPProfileRepository creates a TOTP profile repo backed by members. func NewMongoTOTPProfileRepository(conf *libmongo.Conf) domrepo.TOTPProfileRepository { documentDB, err := libmongo.NewDocumentDB(conf, entity.Member{}.CollectionName()) if err != nil { panic(err) } return &MongoTOTPProfileRepository{db: documentDB} } func (r *MongoTOTPProfileRepository) Get(ctx context.Context, tenantID, uid string) (*domrepo.TOTPProfileRecord, error) { var doc entity.Member filter := bson.M{member.BSONFieldTenantID: tenantID, member.BSONFieldUID: uid} if err := r.db.GetClient().FindOne(ctx, &doc, filter); err != nil { if errors.Is(err, mongodriver.ErrNoDocuments) { return &domrepo.TOTPProfileRecord{}, nil } return nil, err } return memberToTOTPRecord(&doc), nil } func (r *MongoTOTPProfileRepository) Save(ctx context.Context, tenantID, uid string, rec *domrepo.TOTPProfileRecord) error { set := bson.M{ member.BSONFieldTOTPEnrolled: rec.Enrolled, member.BSONFieldTOTPSecretCipher: rec.SecretCipher, member.BSONFieldTOTPBackupCodesHash: rec.BackupCodesHash, member.BSONFieldTOTPEnrolledAt: rec.EnrolledAt, member.BSONFieldUpdateAt: time.Now().UTC().UnixMilli(), } filter := bson.M{member.BSONFieldTenantID: tenantID, member.BSONFieldUID: uid} res, err := r.db.GetClient().UpdateOne(ctx, filter, bson.M{bsonOpSet: set}) if err != nil { return err } if res.MatchedCount == 0 { return member.ErrNotFound } return nil } func (r *MongoTOTPProfileRepository) Clear(ctx context.Context, tenantID, uid string) error { set := bson.M{ member.BSONFieldTOTPEnrolled: false, member.BSONFieldTOTPSecretCipher: []byte{}, member.BSONFieldTOTPBackupCodesHash: []string{}, member.BSONFieldTOTPEnrolledAt: int64(0), member.BSONFieldUpdateAt: time.Now().UTC().UnixMilli(), } filter := bson.M{member.BSONFieldTenantID: tenantID, member.BSONFieldUID: uid} res, err := r.db.GetClient().UpdateOne(ctx, filter, bson.M{bsonOpSet: set}) if err != nil { return err } if res.MatchedCount == 0 { return member.ErrNotFound } return nil } func (r *MongoTOTPProfileRepository) ConsumeBackupCode(ctx context.Context, tenantID, uid, hash string) (bool, error) { filter := bson.M{ member.BSONFieldTenantID: tenantID, member.BSONFieldUID: uid, member.BSONFieldTOTPBackupCodesHash: hash, } update := bson.M{ "$pull": bson.M{member.BSONFieldTOTPBackupCodesHash: hash}, bsonOpSet: bson.M{member.BSONFieldUpdateAt: time.Now().UTC().UnixMilli()}, } res, err := r.db.GetClient().UpdateOne(ctx, filter, update) if err != nil { return false, err } return res.ModifiedCount > 0, nil } func (r *MongoTOTPProfileRepository) ReplaceBackupCodes(ctx context.Context, tenantID, uid string, hashes []string) error { set := bson.M{ member.BSONFieldTOTPBackupCodesHash: hashes, member.BSONFieldUpdateAt: time.Now().UTC().UnixMilli(), } filter := bson.M{member.BSONFieldTenantID: tenantID, member.BSONFieldUID: uid} res, err := r.db.GetClient().UpdateOne(ctx, filter, bson.M{bsonOpSet: set}) if err != nil { return err } if res.MatchedCount == 0 { return member.ErrNotFound } return nil } func memberToTOTPRecord(doc *entity.Member) *domrepo.TOTPProfileRecord { if doc == nil { return &domrepo.TOTPProfileRecord{} } return &domrepo.TOTPProfileRecord{ Enrolled: doc.TOTPEnrolled, SecretCipher: append([]byte(nil), doc.TOTPSecretCipher...), BackupCodesHash: append([]string(nil), doc.TOTPBackupCodesHash...), EnrolledAt: doc.TOTPEnrolledAt, } } var _ domrepo.TOTPProfileRepository = (*MongoTOTPProfileRepository)(nil)