-
Notifications
You must be signed in to change notification settings - Fork 44
/
quota.go
87 lines (80 loc) · 2.45 KB
/
quota.go
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
package common
import (
"context"
"encoding/json"
"fmt"
"io"
"os"
"path/filepath"
"time"
"github.com/osbuild/image-builder/internal/db"
)
const (
day time.Duration = 24 * time.Hour
week time.Duration = 7 * day
DefaultSlidingWindow time.Duration = 2 * week
DefaultQuota int = 100
)
// the QUOTA_FILE needs to contain data arranged as such:
//
// {
// "000000":{
// "quota":2,
// "slidingWindow":1209600000000000
// },
// "000001":{
// "quota":0,
// "slidingWindow":1209600000000000
// },
// "default":{
// "quota":100,
// "slidingWindow":1209600000000000
// }
// }
//
// The unit for the sliding window is the nanosecond.
type Quota struct {
Quota int `json:"quota"`
SlidingWindow time.Duration `json:"slidingWindow"`
}
// Returns true if the number of requests made by OrgID during a sliding window is below a threshold.
// The duration of the sliding window and the value of the threshold must be set in a file pointed by the QUOTA_FILE
// environment variable.
// If the variable is unset (or an empty string), the check is disabled and always returns true.
func CheckQuota(ctx context.Context, orgID string, dB db.DB, quotaFile string) (bool, error) {
if quotaFile == "" {
return true, nil
}
var authorizedRequests int
var slidingWindow time.Duration
// read proper values from quotas' file
var quotas map[string]Quota
jsonFile, err := os.Open(filepath.Clean(quotaFile))
if _, ok := err.(*os.PathError); ok {
return false, fmt.Errorf("No config file for quotas found at %s\n", quotaFile)
} else {
rawJsonFile, err := io.ReadAll(jsonFile)
if err != nil {
return false, fmt.Errorf("Failed to read quota file %q: %s", quotaFile, err.Error())
}
err = json.Unmarshal(rawJsonFile, "as)
if err != nil {
return false, fmt.Errorf("Failed to unmarshal quota file %q: %s", quotaFile, err.Error())
}
if quota, ok := quotas[orgID]; ok {
authorizedRequests = quota.Quota
slidingWindow = quota.SlidingWindow
} else if quota, ok := quotas["default"]; ok {
authorizedRequests = quota.Quota
slidingWindow = quota.SlidingWindow
} else {
return false, fmt.Errorf("No default values in the quotas' file %s\n", quotaFile)
}
}
// read user created requests
count, err := dB.CountComposesSince(ctx, orgID, slidingWindow)
if err != nil {
return false, err
}
return count < authorizedRequests, nil
}