~evanj/cms

ref: 7092668b038fb5c831981d303f6553b94b9c0bf8 cms/internal/s/rl/rl.go -rw-r--r-- 2.4 KiB
7092668bEvan M Jones Fix(rate limiting): Don't rely on MySQL time zone for rate limit 1 year, 2 months ago
                                                                                
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
// Rate limiter.
package rl

import (
	"errors"
	"fmt"
	"log"
	"time"

	"git.sr.ht/~evanj/cms/internal/m/org"
	"git.sr.ht/~evanj/cms/internal/m/space"
	"git.sr.ht/~evanj/cms/internal/m/tier"
	"git.sr.ht/~evanj/cms/internal/m/user"
	"git.sr.ht/~evanj/cms/internal/s/cache"
)

var (
	requestLimits = map[string]int{
		// If not in map, unlimited.
		tier.Free.Name:     60,
		tier.Business.Name: 60,
	}

	spaceLimits = map[string]int{
		// If not in map, unlimited.
		tier.Free.Name:     1,
		tier.Business.Name: 5,
	}

	ErrHitLimit = errors.New("you have surpassed your usage limit: consider upgrading")
)

type RL struct {
	*cache.Cache

	log *log.Logger
	db  *cache.Cache // Or *db.DB depending on nest order.
}

func New(l *log.Logger, db *cache.Cache) RL {
	return RL{db, l, db}
}

// Limit requests made.

func (rl RL) requestLimit(o org.Org) error {
	now := time.Now().UTC()

	limit, ok := requestLimits[o.Tier().Name]
	if !ok {
		// If not in map, unlimited.
		return rl.db.ActionNew(o, now)
	}

	c, err := rl.db.ActionGetCount(o, now.Add(-1*time.Minute), now)
	if err != nil {
		return err
	}

	if c >= limit {
		return ErrHitLimit
	}

	return rl.db.ActionNew(o, now)
}

func (rl RL) UserGet(username, password string) (user.User, error) {
	u, e := rl.db.UserGet(username, password)
	if e != nil {
		return nil, e
	}
	return u, rl.requestLimit(u.Org())
}

func (rl RL) UserGetFromToken(token string) (user.User, error) {
	u, e := rl.db.UserGetFromToken(token)
	if e != nil {
		return nil, e
	}
	return u, rl.requestLimit(u.Org())
}

// Limit spaces created.

func (rl RL) spaceLimit(o org.Org, getter func() (space.Space, error)) (space.Space, error) {
	limit, ok := spaceLimits[o.Tier().Name]
	if !ok {
		// If not in map, unlimited.
		return getter()
	}

	c, err := rl.db.OrgGetSpaceCount(o)
	if err != nil {
		return nil, err
	}

	if c >= limit {
		return nil, fmt.Errorf("can't create new space: %w", ErrHitLimit)
	}

	return getter()
}

func (rl RL) SpaceNew(user user.User, name, desc string) (space.Space, error) {
	return rl.spaceLimit(user.Org(), func() (space.Space, error) {
		return rl.db.SpaceNew(user, name, desc)
	})
}

func (rl RL) SpaceCopy(user user.User, prevS space.Space, name, desc string) (space.Space, error) {
	return rl.spaceLimit(user.Org(), func() (space.Space, error) {
		return rl.db.SpaceCopy(user, prevS, name, desc)
	})
}

// TODO: Limit users created (and associated with org). Currently, no support
// for adding more users to an org.