thread-master/backend/internal/model/threads_account/repository/secrets_mongo.go

package repository

import (
	"context"

	"haixun-backend/internal/library/clock"
	"haixun-backend/internal/library/crypto"
	app "haixun-backend/internal/library/errors"
	"haixun-backend/internal/library/errors/code"
	"haixun-backend/internal/model/threads_account/domain/entity"
	domrepo "haixun-backend/internal/model/threads_account/domain/repository"

	"go.mongodb.org/mongo-driver/bson"
	"go.mongodb.org/mongo-driver/mongo"
	"go.mongodb.org/mongo-driver/mongo/options"
)

type secretsMongoRepository struct {
	collection *mongo.Collection
	cipher     *crypto.Cipher
}

func NewSecretsMongoRepository(db *mongo.Database, cipher *crypto.Cipher) domrepo.SecretsRepository {
	if db == nil {
		return &secretsMongoRepository{cipher: cipher}
	}
	return &secretsMongoRepository{collection: db.Collection(entity.SecretsCollectionName), cipher: cipher}
}

func (r *secretsMongoRepository) EnsureIndexes(ctx context.Context) error {
	if r.collection == nil {
		return nil
	}
	_, err := r.collection.Indexes().CreateOne(ctx, mongo.IndexModel{
		Keys: bson.D{{Key: "_id", Value: 1}},
	})
	return err
}

func (r *secretsMongoRepository) FindByAccountID(ctx context.Context, accountID string) (*entity.Secrets, error) {
	if r.collection == nil {
		return nil, app.For(code.ThreadsAccount).DBUnavailable("Mongo is not configured")
	}
	var out entity.Secrets
	err := r.collection.FindOne(ctx, bson.M{"_id": accountID}).Decode(&out)
	if err == mongo.ErrNoDocuments {
		return nil, nil
	}
	if err != nil {
		return nil, err
	}
	if r.cipher != nil && out.BrowserStorageState != "" {
		plain, decErr := r.cipher.Decrypt(out.BrowserStorageState)
		if decErr != nil {
			return nil, decErr
		}
		out.BrowserStorageState = plain
	}
	if r.cipher != nil && out.APIAccessToken != "" {
		plain, decErr := r.cipher.Decrypt(out.APIAccessToken)
		if decErr != nil {
			return nil, decErr
		}
		out.APIAccessToken = plain
	}
	return &out, nil
}

func (r *secretsMongoRepository) SaveBrowserStorageState(ctx context.Context, accountID, storageState string) (*entity.Secrets, error) {
	if r.collection == nil {
		return nil, app.For(code.ThreadsAccount).DBUnavailable("Mongo is not configured")
	}
	now := clock.NowUnixNano()
	storedValue := storageState
	if r.cipher != nil {
		enc, encErr := r.cipher.Encrypt(storageState)
		if encErr != nil {
			return nil, encErr
		}
		storedValue = enc
	}
	var out entity.Secrets
	err := r.collection.FindOneAndUpdate(
		ctx,
		bson.M{"_id": accountID},
		bson.M{
			"$set": bson.M{
				"browser_storage_state": storedValue,
				"update_at":             now,
			},
			"$setOnInsert": bson.M{"_id": accountID},
		},
		options.FindOneAndUpdate().SetUpsert(true).SetReturnDocument(options.After),
	).Decode(&out)
	if err != nil {
		return nil, err
	}
	// Return plaintext to callers regardless of at-rest encryption.
	out.BrowserStorageState = storageState
	return &out, nil
}

func (r *secretsMongoRepository) SaveAPIAccessToken(ctx context.Context, accountID, accessToken string, expiresAt int64) (*entity.Secrets, error) {
	if r.collection == nil {
		return nil, app.For(code.ThreadsAccount).DBUnavailable("Mongo is not configured")
	}
	now := clock.NowUnixNano()
	storedToken := accessToken
	if r.cipher != nil {
		enc, encErr := r.cipher.Encrypt(accessToken)
		if encErr != nil {
			return nil, encErr
		}
		storedToken = enc
	}
	var out entity.Secrets
	err := r.collection.FindOneAndUpdate(
		ctx,
		bson.M{"_id": accountID},
		bson.M{
			"$set": bson.M{
				"api_access_token":     storedToken,
				"api_token_expires_at": expiresAt,
				"update_at":            now,
			},
			"$setOnInsert": bson.M{"_id": accountID},
		},
		options.FindOneAndUpdate().SetUpsert(true).SetReturnDocument(options.After),
	).Decode(&out)
	if err != nil {
		return nil, err
	}
	out.APIAccessToken = accessToken
	out.APITokenExpiresAt = expiresAt
	return &out, nil
}

func (r *secretsMongoRepository) ListAPIAccessTokensExpiringBefore(
	ctx context.Context,
	beforeNs int64,
) ([]*entity.Secrets, error) {
	if r.collection == nil {
		return nil, app.For(code.ThreadsAccount).DBUnavailable("Mongo is not configured")
	}
	if beforeNs <= 0 {
		return nil, nil
	}
	cursor, err := r.collection.Find(ctx, bson.M{
		"api_access_token": bson.M{"$ne": ""},
		"api_token_expires_at": bson.M{
			"$gt":  int64(0),
			"$lte": beforeNs,
		},
	})
	if err != nil {
		return nil, err
	}
	defer cursor.Close(ctx)
	var items []*entity.Secrets
	if err := cursor.All(ctx, &items); err != nil {
		return nil, err
	}
	for _, item := range items {
		if item == nil {
			continue
		}
		if r.cipher != nil && item.APIAccessToken != "" {
			plain, decErr := r.cipher.Decrypt(item.APIAccessToken)
			if decErr != nil {
				return nil, decErr
			}
			item.APIAccessToken = plain
		}
	}
	return items, nil
}

func (r *secretsMongoRepository) ClearAPIAccessToken(ctx context.Context, accountID string) error {
	if r.collection == nil {
		return app.For(code.ThreadsAccount).DBUnavailable("Mongo is not configured")
	}
	now := clock.NowUnixNano()
	res, err := r.collection.UpdateOne(
		ctx,
		bson.M{"_id": accountID},
		bson.M{
			"$set": bson.M{
				"api_access_token":     "",
				"api_token_expires_at": int64(0),
				"update_at":            now,
			},
		},
	)
	if err != nil {
		return err
	}
	if res.MatchedCount == 0 {
		return nil
	}
	return nil
}