sxchnge/backend/edge/shared/storage/refreshtokens/crud.go

package refreshtokens

import (
	"context"
	"errors"
	"time"

	"git.sendico.io/tech/sxchnge/bff/pkg/db/repository"
	"git.sendico.io/tech/sxchnge/bff/pkg/db/repository/builder"
	"git.sendico.io/tech/sxchnge/bff/pkg/merrors"
	"git.sendico.io/tech/sxchnge/bff/pkg/model"
	"git.sendico.io/tech/sxchnge/bff/pkg/mservice"
	"git.sendico.io/tech/sxchnge/bff/pkg/mutil/mzap"
	"go.mongodb.org/mongo-driver/v2/bson"
	"go.uber.org/zap"
)

func (db *RefreshTokenDB) Create(ctx context.Context, rt *model.RefreshToken) error {
	if rt.LastUsedAt.IsZero() {
		rt.LastUsedAt = time.Now()
	}

	if rt.AccountRef == nil {
		return merrors.InvalidArgument("Account reference must have a valid value", "refreshToken.accountRef")
	}

	matched, err := db.replaceActive(ctx, rt)
	if err != nil {
		return err
	}
	if matched {
		return nil
	}

	db.Logger.Debug("Registering refresh token", zap.String("client_id", rt.ClientID), zap.String("device_id", rt.DeviceID))
	if err := db.DBImp.Create(ctx, rt); err != nil {
		if !errors.Is(err, merrors.ErrDataConflict) {
			return err
		}

		// Another login can establish the same active slot after the patch and
		// before this insert. Converge on that row instead of returning the
		// unique-index conflict to the caller.
		db.Logger.Debug("Concurrent refresh session registration detected",
			zap.String("client_id", rt.ClientID), zap.String("device_id", rt.DeviceID))
		matched, reconcileErr := db.replaceActive(ctx, rt)
		if reconcileErr != nil {
			return reconcileErr
		}
		if !matched {
			return err
		}
	}
	return nil
}

func (db *RefreshTokenDB) replaceActive(ctx context.Context, rt *model.RefreshToken) (bool, error) {
	patch := repository.Patch().
		Set(repository.Field(TokenField), rt.RefreshToken).
		Set(repository.Field(ExpiresAtField), rt.ExpiresAt).
		Set(repository.Field(UserAgentField), rt.UserAgent).
		Set(repository.Field(IPAddressField), rt.IPAddress).
		Set(repository.Field(LastUsedAtField), rt.LastUsedAt).
		Set(repository.Field(IsRevokedField), rt.IsRevoked).
		Set(repository.Field(SessionVersionField), rt.SessionVersion)

	var stored model.RefreshToken
	matched, err := db.Repository.PatchOneByFilter(
		ctx,
		filterSlotByAccount(*rt.AccountRef, &rt.SessionIdentifier),
		patch,
		&stored,
	)
	if err != nil {
		db.Logger.Warn("Failed to replace active refresh session", zap.Error(err),
			zap.String("client_id", rt.ClientID), zap.String("device_id", rt.DeviceID))
		return false, err
	}
	if !matched {
		return false, nil
	}

	db.Logger.Debug("Replaced active refresh session",
		zap.String("client_id", rt.ClientID), zap.String("device_id", rt.DeviceID))
	rt.SetID(*stored.GetID())
	return true, nil
}

func (db *RefreshTokenDB) Update(ctx context.Context, rt *model.RefreshToken) error {
	rt.LastUsedAt = time.Now()

	// Use Patch instead of Update to avoid race conditions
	patch := repository.Patch().
		Set(repository.Field(TokenField), rt.RefreshToken).
		Set(repository.Field(ExpiresAtField), rt.ExpiresAt).
		Set(repository.Field(UserAgentField), rt.UserAgent).
		Set(repository.Field(IPAddressField), rt.IPAddress).
		Set(repository.Field(LastUsedAtField), rt.LastUsedAt).
		Set(repository.Field(IsRevokedField), rt.IsRevoked).
		Set(repository.Field(SessionVersionField), rt.SessionVersion)

	return db.Patch(ctx, *rt.GetID(), patch)
}

func (db *RefreshTokenDB) Delete(ctx context.Context, tokenRef bson.ObjectID) error {
	db.Logger.Debug("Deleting refresh token", mzap.ObjRef("refresh_token_ref", tokenRef))
	return db.DBImp.Delete(ctx, tokenRef)
}

func (db *RefreshTokenDB) Revoke(ctx context.Context, accountRef bson.ObjectID, session *model.SessionIdentifier) error {
	var rt model.RefreshToken
	f := filterByAccount(accountRef, session)
	if err := db.Repository.FindOneByFilter(ctx, f, &rt); err != nil {
		if errors.Is(err, merrors.ErrNoData) {
			db.Logger.Debug("Refresh token not found during revoke",
				mzap.AccRef(accountRef), zap.String("client_id", session.ClientID), zap.String("device_id", session.DeviceID))
			return nil
		}
		return err
	}

	// Use Patch to update the revocation status
	patch := repository.Patch().
		Set(repository.Field(IsRevokedField), true).
		Set(repository.Field(LastUsedAtField), time.Now())

	return db.Patch(ctx, *rt.GetID(), patch)
}

func (db *RefreshTokenDB) RevokeToken(ctx context.Context, token *model.ClientRefreshToken) (bool, error) {
	query := filter(&token.SessionIdentifier).
		And(repository.Query().Comparison(repository.Field(TokenField), builder.Eq, token.RefreshToken))
	patch := repository.Patch().
		Set(repository.Field(IsRevokedField), true).
		Set(repository.Field(LastUsedAtField), time.Now())
	modified, err := db.Repository.PatchMany(ctx, query, patch)
	return modified > 0, err
}

func (db *RefreshTokenDB) GetSession(ctx context.Context, accountRef bson.ObjectID, session *model.SessionIdentifier) (*model.RefreshToken, error) {
	var rt model.RefreshToken
	if err := db.Repository.FindOneByFilter(ctx, filterByAccount(accountRef, session), &rt); err != nil {
		return nil, err
	}
	if !rt.ExpiresAt.After(time.Now()) {
		return nil, merrors.ErrNoData
	}
	return &rt, nil
}

func (db *RefreshTokenDB) GetByCRT(ctx context.Context, t *model.ClientRefreshToken) (*model.RefreshToken, error) {
	var rt model.RefreshToken
	f := filter(&t.SessionIdentifier).And(repository.Query().Filter(repository.Field("token"), t.RefreshToken))
	if err := db.Repository.FindOneByFilter(ctx, f, &rt); err != nil {
		if !errors.Is(err, merrors.ErrNoData) {
			db.Logger.Warn("Failed to fetch refresh token", zap.Error(err),
				zap.String("client_id", t.ClientID), zap.String("device_id", t.DeviceID))
		}
		return nil, err
	}

	// Check if token is expired
	if rt.ExpiresAt.Before(time.Now()) {
		db.Logger.Debug("Refresh token expired", mzap.StorableRef(&rt),
			zap.String("client_id", t.ClientID), zap.String("device_id", t.DeviceID),
			zap.Time("expires_at", rt.ExpiresAt))
		return nil, merrors.AccessDenied(mservice.RefreshTokens, "read", *rt.GetID())
	}

	// Check if token is revoked
	if rt.IsRevoked {
		db.Logger.Debug("Refresh token is revoked", mzap.StorableRef(&rt),
			zap.String("client_id", t.ClientID), zap.String("device_id", t.DeviceID))
		return nil, merrors.ErrNoData
	}

	if rt.GetID() != nil && !rt.GetID().IsZero() {
		lastUsedAt := time.Now()
		patch := repository.Patch().Set(repository.Field(LastUsedAtField), lastUsedAt)
		if err := db.Patch(ctx, *rt.GetID(), patch); err != nil {
			db.Logger.Warn("Failed to update refresh token last usage timestamp", zap.Error(err),
				zap.String("client_id", t.ClientID), zap.String("device_id", t.DeviceID))
		} else {
			rt.LastUsedAt = lastUsedAt
		}
	}

	return &rt, nil
}