package storagebonus import ( "database/sql" "errors" "fmt" "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" enteTime "github.com/ente-io/museum/pkg/utils/time" "github.com/ente-io/stacktrace" "github.com/gin-gonic/gin" ) const ( codeLength = 6 referralAmountInGb = 10 maxClaimableReferralAmount = 2000 numOfDaysToClaimReferral = 32 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 && user.CreationTime > enteTime.MicrosecondBeforeDays(numOfDaysToClaimReferral) && !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.CreationTime < enteTime.MicrosecondBeforeDays(numOfDaysToClaimReferral) { return stacktrace.Propagate(entity.CanNotApplyCodeErr, "account is too old to apply code") } else 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 := 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 } // generateAlphaNumString returns AlphaNumeric code of given length // which exclude number 0 and letter O. The code always starts with an // alphabet func generateAlphaNumString(length int) (string, error) { // Define the alphabet and numbers to be used in the string. alphabet := "ABCDEFGHIJKLMNPQRSTUVWXYZ" // Define the alphabet and numbers to be used in the string. alphaNum := fmt.Sprintf("%s123456789", alphabet) // Allocate a byte slice with the desired length. result := make([]byte, length) // Generate the first letter as an alphabet. r0, err := auth.GenerateRandomInt(int64(len(alphabet))) if err != nil { return "", stacktrace.Propagate(err, "") } result[0] = alphabet[r0] // Generate the remaining characters as alphanumeric. for i := 1; i < length; i++ { ri, err := auth.GenerateRandomInt(int64(len(alphaNum))) if err != nil { return "", stacktrace.Propagate(err, "") } result[i] = alphaNum[ri] } return string(result), nil }