template-monorepo/internal/model/member/repository/totp_profile_mongo.go

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)