keeper.gno

  1package referral
  2
  3import (
  4	"errors"
  5	"time"
  6
  7	bptree "gno.land/p/nt/bptree/v0"
  8)
  9
 10const (
 11	// MinTimeBetweenUpdates is minimum duration between operations (24 hours).
 12	MinTimeBetweenUpdates int64 = 24 * 60 * 60
 13)
 14
 15// keeper implements ReferralKeeper using BPTree storage.
 16// It includes rate limiting to prevent abuse.
 17type keeper struct {
 18	store   *bptree.BPTree // address(string) -> referral address(string)
 19	lastOps *bptree.BPTree // address(string) -> last operation timestamp(int64)
 20}
 21
 22var _ ReferralKeeper = &keeper{}
 23
 24// NewKeeper creates a new ReferralKeeper instance.
 25func NewKeeper() ReferralKeeper {
 26	return &keeper{
 27		store:   bptree.NewBPTreeN(16),
 28		lastOps: bptree.NewBPTreeN(16),
 29	}
 30}
 31
 32// register creates or updates a referral relationship between addresses.
 33// Setting refAddr to the contract's own address removes the referral.
 34func (k *keeper) register(addr, refAddr address) (address, error) {
 35	if err := k.validateAddresses(addr, refAddr); err != nil {
 36		return zeroAddress, err
 37	}
 38
 39	addrStr := addr.String()
 40	refAddrStr := refAddr.String()
 41
 42	if isRemovalRequest(refAddr) {
 43		if k.has(addr) {
 44			_, ok := k.store.Remove(addrStr)
 45			if !ok {
 46				return zeroAddress, errors.New(ErrNotFound)
 47			}
 48		}
 49
 50		return zeroAddress, nil
 51	}
 52
 53	if err := k.checkRateLimit(addrStr); err != nil {
 54		return zeroAddress, err
 55	}
 56
 57	k.store.Set(addrStr, refAddrStr)
 58	k.lastOps.Set(addrStr, time.Now().Unix())
 59
 60	return refAddr, nil
 61}
 62
 63// validateAddresses validates that addresses are properly formatted and not self-referencing.
 64func (k *keeper) validateAddresses(addr, refAddr address) error {
 65	if !addr.IsValid() || (!isRemovalRequest(refAddr) && !refAddr.IsValid()) {
 66		return errors.New(ErrInvalidAddress)
 67	}
 68	if addr == refAddr {
 69		return errors.New(ErrSelfReferral)
 70	}
 71	return nil
 72}
 73
 74// has returns true if a referral exists for the given address.
 75func (k *keeper) has(addr address) bool {
 76	exists := k.store.Get(addr.String()) != nil
 77	return exists
 78}
 79
 80// get retrieves the referral address for a given address.
 81// Returns ErrNotFound if no referral exists.
 82func (k *keeper) get(addr address) (address, error) {
 83	if !addr.IsValid() {
 84		return zeroAddress, errors.New(ErrInvalidAddress)
 85	}
 86
 87	val := k.store.Get(addr.String())
 88	if val == nil {
 89		return zeroAddress, errors.New(ErrNotFound)
 90	}
 91
 92	refAddr, ok := val.(string)
 93	if !ok {
 94		return zeroAddress, errors.New(ErrInvalidAddress)
 95	}
 96
 97	return address(refAddr), nil
 98}
 99
100// isEmpty returns true if no referrals exist in the store.
101func (k *keeper) isEmpty() bool {
102	return k.store.Size() == 0
103}
104
105// getLastOpTimestamp retrieves the last operation timestamp for a given address.
106// Returns ErrNotFound if no operation exists.
107func (k *keeper) getLastOpTimestamp(addr address) (int64, error) {
108	if !addr.IsValid() {
109		return 0, errors.New(ErrInvalidAddress)
110	}
111
112	val := k.lastOps.Get(addr.String())
113	if val == nil {
114		return 0, errors.New(ErrNotFound)
115	}
116
117	ts, ok := val.(int64)
118	if !ok {
119		return 0, errors.New(ErrInvalidTime)
120	}
121
122	return ts, nil
123}
124
125// checkRateLimit verifies if enough time has passed since the last operation.
126// Returns ErrTooManyRequests if rate limit is exceeded.
127func (k *keeper) checkRateLimit(addr string) error {
128	now := time.Now().Unix()
129
130	lastOpTimeRaw := k.lastOps.Get(addr)
131	if lastOpTimeRaw == nil {
132		return nil
133	}
134
135	lastOpTime, ok := lastOpTimeRaw.(int64)
136	if !ok {
137		return errors.New(ErrInvalidTime)
138	}
139
140	timeSinceLastOp := now - lastOpTime
141	if timeSinceLastOp < MinTimeBetweenUpdates {
142		return errors.New(ErrTooManyRequests)
143	}
144
145	return nil
146}