ente/server/pkg/controller/storagebonus/referral.go

package storagebonus

import (
	"database/sql"
	"errors"
	goaway "github.com/TwiN/go-away"
	"github.com/ente-io/museum/pkg/utils/random"
	"strings"

	"github.com/ente-io/museum/ente"
	entity "github.com/ente-io/museum/ente/storagebonus"
	"github.com/ente-io/museum/pkg/controller/email"
	"github.com/ente-io/museum/pkg/controller/lock"

	"github.com/ente-io/museum/pkg/repo"
	"github.com/ente-io/museum/pkg/repo/storagebonus"
	"github.com/ente-io/museum/pkg/utils/auth"
	"github.com/ente-io/stacktrace"
	"github.com/gin-gonic/gin"
)

const (
	codeLength                 = 6
	referralAmountInGb         = 10
	maxClaimableReferralAmount = 2000
	defaultPlanType            = entity.TenGbOnUpgrade
)

// Controller exposes functions to interact with family module
type Controller struct {
	UserRepo                    *repo.UserRepository
	StorageBonus                *storagebonus.Repository
	LockController              *lock.LockController
	CronRunning                 bool
	EmailNotificationController *email.EmailNotificationController
}

func (c *Controller) GetUserReferralView(ctx *gin.Context) (*entity.GetUserReferralView, error) {
	// Get the user id from the context
	userID := auth.GetUserID(ctx.Request.Header)

	// Use goroutines to fetch UserRepo.Get, HasAppliedReferral
	user, err := c.UserRepo.Get(userID)
	if err != nil {
		return nil, stacktrace.Propagate(err, "failed to get user")
	}
	appliedReferral, err := c.StorageBonus.HasAppliedReferral(ctx, userID)
	if err != nil {
		return nil, stacktrace.Propagate(err, "")
	}
	isFamilyMember := user.FamilyAdminID != nil && *user.FamilyAdminID != userID
	enableApplyCode := !appliedReferral && !isFamilyMember
	// Get the referral code for the user or family admin
	codeUser := userID
	if isFamilyMember {
		codeUser = *user.FamilyAdminID
	}
	referralCode, err2 := c.GetOrCreateReferralCode(ctx, codeUser)
	if err2 != nil {
		return nil, stacktrace.Propagate(err2, "failed to get or create referral code")
	}
	storageClaimed, err2 := c.GetActiveReferralBonusValue(ctx, codeUser)
	if err2 != nil {
		return nil, stacktrace.Propagate(err2, "failed to get storage claimed")
	}

	return &entity.GetUserReferralView{
		PlanInfo: entity.PlanInfo{
			IsEnabled:               true,
			PlanType:                defaultPlanType,
			StorageInGB:             referralAmountInGb,
			MaxClaimableStorageInGB: maxClaimableReferralAmount,
		},
		Code:            referralCode,
		EnableApplyCode: enableApplyCode,
		IsFamilyMember:  isFamilyMember,
		HasAppliedCode:  appliedReferral,
		ClaimedStorage:  *storageClaimed,
	}, nil
}

func (c *Controller) ApplyReferralCode(ctx *gin.Context, code string) error {
	// Get user id from the context
	userID := auth.GetUserID(ctx.Request.Header)
	user, err := c.UserRepo.Get(userID)
	if err != nil {
		return stacktrace.Propagate(err, "failed to get user")
	}

	codeOwnerID, err := c.StorageBonus.GetUserIDByCode(ctx, code)
	if err != nil {
		return stacktrace.Propagate(err, "failed to get user id by code")
	}
	// Verify that the codeOwnerID is not deleted yet
	_, err = c.UserRepo.Get(*codeOwnerID)
	if err != nil {
		if errors.Is(err, ente.ErrUserDeleted) {
			return stacktrace.Propagate(entity.InvalidCodeErr, "code belongs to deleted user")
		}
		return stacktrace.Propagate(err, "failed to get user")
	}

	if user.FamilyAdminID != nil && userID != *user.FamilyAdminID {
		return stacktrace.Propagate(entity.CanNotApplyCodeErr, "user is member of a family plan")
	}

	err = c.StorageBonus.TrackReferralAndInviteeBonus(ctx, userID, *codeOwnerID, defaultPlanType)
	if err != nil {
		return stacktrace.Propagate(err, "failed to apply code")
	}
	return nil
}

func (c *Controller) GetOrCreateReferralCode(ctx *gin.Context, userID int64) (*string, error) {
	referralCode, err := c.StorageBonus.GetCode(ctx, userID)
	if err != nil {
		if !errors.Is(err, sql.ErrNoRows) {
			return nil, stacktrace.Propagate(err, "failed to get storagebonus code")
		}
		code, err := random.GenerateAlphaNumString(codeLength)
		if err != nil {
			return nil, stacktrace.Propagate(err, "")
		}
		err = c.StorageBonus.InsertCode(ctx, userID, code)
		if err != nil {
			return nil, stacktrace.Propagate(err, "failed to insert storagebonus code")
		}
		referralCode = &code
	}
	return referralCode, nil
}

func (c *Controller) UpdateReferralCode(ctx *gin.Context, userID int64, code string, isAdminEdit bool) error {
	code = strings.ToUpper(code)
	if !random.IsAlphanumeric(code) {
		return stacktrace.Propagate(ente.NewBadRequestWithMessage("code is not alphanumeric"), "")
	}
	if len(code) < 4 || len(code) > 20 {
		return stacktrace.Propagate(ente.NewBadRequestWithMessage("code length should be between 4 and 8"), "")
	}

	// Check if the code contains any offensive language using the go-away library
	if goaway.IsProfane(code) {
		return stacktrace.Propagate(ente.NewBadRequestWithMessage("Referral code contains offensive language and cannot be used"), "")
	}

	err := c.StorageBonus.AddNewCode(ctx, userID, code, isAdminEdit)
	if err != nil {
		return stacktrace.Propagate(err, "failed to update referral code")
	}
	return nil
}