keeper.gno

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