Skip to content

Commit 8c08df3

Browse files
efectnefectn
committed
server: add support for rate limiting
1 parent b9f01d0 commit 8c08df3

File tree

5 files changed

+129
-69
lines changed

5 files changed

+129
-69
lines changed

‎cmd/main.go‎

Lines changed: 18 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,8 @@ import (
1313
"github.com/armbian/ansi-hastebin/internal/keygenerator"
1414
"github.com/armbian/ansi-hastebin/internal/server"
1515
"github.com/armbian/ansi-hastebin/internal/storage"
16-
"github.com/sirupsen/logrus"
16+
"github.com/rs/zerolog"
17+
"github.com/rs/zerolog/log"
1718
)
1819

1920
func handleConfig(location string) (*config.Config, storage.Storage, keygenerator.KeyGenerator) {
@@ -35,25 +36,25 @@ func handleConfig(location string) (*config.Config, storage.Storage, keygenerato
3536
case "s3":
3637
pasteStorage = storage.NewS3Storage(cfg.Storage.Host, cfg.Storage.Port, cfg.Storage.Username, cfg.Storage.Password, cfg.Storage.AWSRegion, cfg.Storage.Bucket)
3738
default:
38-
logrus.Fatalf("Unknown storage type: %s", cfg.Storage.Type)
39+
log.Fatal().Str("storage_type", cfg.Storage.Type).Msg("Unknown storage type")
3940
return nil, nil, nil
4041
}
4142

4243
// Set static documents from config
4344
for _, doc := range cfg.Documents {
4445
file, err := os.OpenFile(doc.Path, os.O_RDONLY, 0644)
4546
if err != nil {
46-
logrus.WithError(err).WithField("path", doc.Path).Fatal("Failed to open document")
47+
log.Fatal().Err(err).Str("path", doc.Path).Msg("Failed to open document")
4748
}
4849

4950
content, err := io.ReadAll(file)
5051
if err != nil {
51-
logrus.WithError(err).WithField("path", doc.Path).Fatal("Failed to read document")
52+
log.Fatal().Err(err).Str("path", doc.Path).Msg("Failed to read document")
5253
}
5354
file.Close()
5455

5556
if err := pasteStorage.Set(doc.Key, string(content), false); err != nil {
56-
logrus.WithError(err).WithField("key", doc.Key).Fatal("Failed to set document")
57+
log.Fatal().Err(err).Str("key", doc.Key).Msg("Failed to set document")
5758
}
5859
}
5960

@@ -65,10 +66,21 @@ func handleConfig(location string) (*config.Config, storage.Storage, keygenerato
6566
case "phonetic":
6667
keyGenerator = keygenerator.NewPhoneticKeyGenerator()
6768
default:
68-
logrus.Fatalf("Unknown key generator: %s", cfg.KeyGenerator)
69+
log.Fatal().Str("key_generator", cfg.KeyGenerator).Msg("Unknown key generator")
6970
return nil, nil, nil
7071
}
7172

73+
// Adjust logger
74+
logLevel, err := zerolog.ParseLevel(cfg.Logging.Level)
75+
if err != nil {
76+
log.Fatal().Err(err).Str("level", cfg.Logging.Level).Msg("Failed to parse log level")
77+
}
78+
log.Logger = log.Level(logLevel)
79+
80+
if cfg.Logging.Colorize {
81+
log.Logger = log.Output(zerolog.ConsoleWriter{Out: os.Stdout})
82+
}
83+
7284
return cfg, pasteStorage, keyGenerator
7385
}
7486

‎config.yaml‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,11 @@ documents:
1515
- key: "about"
1616
path: "./about.md"
1717

18+
rate_limiting:
19+
enable: true
20+
limit: 500
21+
window: 15
22+
1823
logging:
1924
level: "info"
2025
type: "text"

‎config/config.go‎

Lines changed: 56 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -5,14 +5,27 @@ import (
55
"strconv"
66
"strings"
77

8-
"github.com/sirupsen/logrus"
8+
"github.com/rs/zerolog/log"
99
"gopkg.in/yaml.v3"
1010
)
1111

1212
type LoggingConfig struct {
13-
Level string `yaml:"level"`
14-
Type string `yaml:"type"`
15-
Colorize bool `yaml:"colorize"`
13+
// Level is the logging level to use
14+
Level string `yaml:"level"`
15+
16+
// Colorize is a flag to enable colorized output
17+
Colorize bool `yaml:"colorize"`
18+
}
19+
20+
type RateLimitingConfig struct {
21+
// Enable is a flag to enable rate limiting
22+
Enable bool `yaml:"enable"`
23+
24+
// Limit is the maximum number of requests
25+
Limit int `yaml:"limit"`
26+
27+
// Window is the time window to limit requests
28+
Window int `yaml:"window"`
1629
}
1730

1831
type StorageConfig struct {
@@ -94,6 +107,9 @@ type Config struct {
94107
// Logging is the logging configuration
95108
Logging LoggingConfig `yaml:"logging"`
96109

110+
// RateLimiting is the rate limiting configuration
111+
RateLimiting RateLimitingConfig `yaml:"rate_limiting"`
112+
97113
// Documents is the list of documents to load statically
98114
Documents []DocumentConfig `yaml:"documents"`
99115
}
@@ -112,7 +128,6 @@ var DefaultConfig = &Config{
112128
},
113129
Logging: LoggingConfig{
114130
Level: "info",
115-
Type: "text",
116131
},
117132
Documents: []DocumentConfig{
118133
{
@@ -129,12 +144,12 @@ func NewConfig(configFile string) *Config {
129144
// Read the configuration file
130145
data, err := os.ReadFile(configFile)
131146
if err != nil && !os.IsNotExist(err) {
132-
logrus.WithError(err).Fatal("Failed to read configuration file")
147+
log.Fatal().Err(err).Msg("Failed to read configuration file")
133148
}
134149

135150
// Unmarshal the configuration file
136151
if err := yaml.Unmarshal(data, cfg); err != nil {
137-
logrus.WithError(err).Fatal("Failed to unmarshal configuration file")
152+
log.Fatal().Err(err).Msg("Failed to unmarshal configuration file")
138153
}
139154

140155
// Override with environment variables
@@ -145,39 +160,39 @@ func NewConfig(configFile string) *Config {
145160
if port := os.Getenv("PORT"); port != "" {
146161
portInt, err := strconv.Atoi(port)
147162
if err != nil {
148-
logrus.WithError(err).Fatal("Failed to parse PORT environment variable")
163+
log.Fatal().Err(err).Msg("Failed to parse PORT environment variable")
149164
}
150165
cfg.Port = portInt
151166
}
152167

153168
if keyLength := os.Getenv("KEY_LENGTH"); keyLength != "" {
154169
keyLengthInt, err := strconv.Atoi(keyLength)
155170
if err != nil {
156-
logrus.WithError(err).Fatal("Failed to parse KEY_LENGTH environment variable")
171+
log.Fatal().Err(err).Msg("Failed to parse KEY_LENGTH environment variable")
157172
}
158173
cfg.KeyLength = keyLengthInt
159174
}
160175

161176
if maxLength := os.Getenv("MAX_LENGTH"); maxLength != "" {
162177
maxLengthInt, err := strconv.Atoi(maxLength)
163178
if err != nil {
164-
logrus.WithError(err).Fatal("Failed to parse MAX_LENGTH environment variable")
179+
log.Fatal().Err(err).Msg("Failed to parse MAX_LENGTH environment variable")
165180
}
166181
cfg.MaxLength = maxLengthInt
167182
}
168183

169184
if staticMaxAge := os.Getenv("STATIC_MAX_AGE"); staticMaxAge != "" {
170185
staticMaxAgeInt, err := strconv.Atoi(staticMaxAge)
171186
if err != nil {
172-
logrus.WithError(err).Fatal("Failed to parse STATIC_MAX_AGE environment variable")
187+
log.Fatal().Err(err).Msg("Failed to parse STATIC_MAX_AGE environment variable")
173188
}
174189
cfg.StaticMaxAge = staticMaxAgeInt
175190
}
176191

177192
if recompressStaticAssets := os.Getenv("RECOMPRESS_STATIC_ASSETS"); recompressStaticAssets != "" {
178193
recompressStaticAssetsBool, err := strconv.ParseBool(recompressStaticAssets)
179194
if err != nil {
180-
logrus.WithError(err).Fatal("Failed to parse RECOMPRESS_STATIC_ASSETS environment variable")
195+
log.Fatal().Err(err).Msg("Failed to parse RECOMPRESS_STATIC_ASSETS environment variable")
181196
}
182197
cfg.RecompressStaticAssets = recompressStaticAssetsBool
183198
}
@@ -197,7 +212,7 @@ func NewConfig(configFile string) *Config {
197212
if storagePort := os.Getenv("STORAGE_PORT"); storagePort != "" {
198213
storagePortInt, err := strconv.Atoi(storagePort)
199214
if err != nil {
200-
logrus.WithError(err).Fatal("Failed to parse STORAGE_PORT environment variable")
215+
log.Fatal().Err(err).Msg("Failed to parse STORAGE_PORT environment variable")
201216
}
202217
cfg.Storage.Port = storagePortInt
203218
}
@@ -230,18 +245,41 @@ func NewConfig(configFile string) *Config {
230245
cfg.Logging.Level = loggingLevel
231246
}
232247

233-
if loggingType := os.Getenv("LOGGING_TYPE"); loggingType != "" {
234-
cfg.Logging.Type = loggingType
235-
}
236-
237248
if loggingColorize := os.Getenv("LOGGING_COLORIZE"); loggingColorize != "" {
238249
loggingColorizeBool, err := strconv.ParseBool(loggingColorize)
239250
if err != nil {
240-
logrus.WithError(err).Fatal("Failed to parse LOGGING_COLORIZE environment variable")
251+
log.Fatal().Err(err).Msg("Failed to parse LOGGING_COLORIZE environment variable")
241252
}
242253
cfg.Logging.Colorize = loggingColorizeBool
243254
}
244255

256+
if rateLimitingEnable := os.Getenv("RATE_LIMITING_ENABLE"); rateLimitingEnable != "" {
257+
rateLimitingEnableBool, err := strconv.ParseBool(rateLimitingEnable)
258+
if err != nil {
259+
log.Fatal().Err(err).Msg("Failed to parse RATE_LIMITING_ENABLE environment variable")
260+
}
261+
262+
cfg.RateLimiting.Enable = rateLimitingEnableBool
263+
}
264+
265+
if rateLimitingLimit := os.Getenv("RATE_LIMITING_LIMIT"); rateLimitingLimit != "" {
266+
rateLimitingLimitInt, err := strconv.Atoi(rateLimitingLimit)
267+
if err != nil {
268+
log.Fatal().Err(err).Msg("Failed to parse RATE_LIMITING_LIMIT environment variable")
269+
}
270+
271+
cfg.RateLimiting.Limit = rateLimitingLimitInt
272+
}
273+
274+
if rateLimitingWindow := os.Getenv("RATE_LIMITING_WINDOW"); rateLimitingWindow != "" {
275+
rateLimitingWindowInt, err := strconv.Atoi(rateLimitingWindow)
276+
if err != nil {
277+
log.Fatal().Err(err).Msg("Failed to parse RATE_LIMITING_WINDOW environment variable")
278+
}
279+
280+
cfg.RateLimiting.Window = rateLimitingWindowInt
281+
}
282+
245283
// Walk environment variables for documents
246284
for _, env := range os.Environ() {
247285
if len(env) > 10 && env[:10] == "DOCUMENTS_" {
@@ -292,9 +330,5 @@ func NewConfig(configFile string) *Config {
292330
cfg.Logging.Level = DefaultConfig.Logging.Level
293331
}
294332

295-
if cfg.Logging.Type == "" {
296-
cfg.Logging.Type = DefaultConfig.Logging.Type
297-
}
298-
299333
return cfg
300334
}

‎config/config_test.go‎

Lines changed: 37 additions & 35 deletions
Original file line numberDiff line numberDiff line change
@@ -4,44 +4,47 @@ import (
44
"os"
55
"testing"
66

7-
"github.com/stretchr/testify/assert"
7+
"github.com/stretchr/testify/require"
88
)
99

1010
func TestNewConfig_DefaultValues(t *testing.T) {
1111
cfg := NewConfig("nonexistent.yaml") // Should use defaults since file doesn't exist
1212

13-
assert.Equal(t, "0.0.0.0", cfg.Host)
14-
assert.Equal(t, 7777, cfg.Port)
15-
assert.Equal(t, 10, cfg.KeyLength)
16-
assert.Equal(t, "phonetic", cfg.KeyGenerator)
17-
assert.Equal(t, "file", cfg.Storage.Type)
18-
assert.Equal(t, "data", cfg.Storage.FilePath)
19-
assert.Equal(t, "info", cfg.Logging.Level)
20-
assert.Equal(t, "text", cfg.Logging.Type)
13+
require.Equal(t, "0.0.0.0", cfg.Host)
14+
require.Equal(t, 7777, cfg.Port)
15+
require.Equal(t, 10, cfg.KeyLength)
16+
require.Equal(t, "phonetic", cfg.KeyGenerator)
17+
require.Equal(t, "file", cfg.Storage.Type)
18+
require.Equal(t, "data", cfg.Storage.FilePath)
19+
require.Equal(t, "info", cfg.Logging.Level)
2120
}
2221

2322
func TestNewConfig_OverrideWithEnvVars(t *testing.T) {
24-
os.Setenv("HOST", "127.0.0.1")
25-
os.Setenv("PORT", "8080")
26-
os.Setenv("KEY_LENGTH", "15")
27-
os.Setenv("MAX_LENGTH", "5000000")
28-
os.Setenv("STORAGE_TYPE", "redis")
29-
os.Setenv("STORAGE_HOST", "localhost")
30-
os.Setenv("STORAGE_PORT", "6379")
31-
os.Setenv("LOGGING_LEVEL", "debug")
23+
t.Setenv("HOST", "127.0.0.1")
24+
t.Setenv("PORT", "8080")
25+
t.Setenv("KEY_LENGTH", "15")
26+
t.Setenv("MAX_LENGTH", "5000000")
27+
t.Setenv("STORAGE_TYPE", "redis")
28+
t.Setenv("STORAGE_HOST", "localhost")
29+
t.Setenv("STORAGE_PORT", "6379")
30+
t.Setenv("LOGGING_LEVEL", "debug")
31+
t.Setenv("RATE_LIMITING_ENABLE", "true")
32+
t.Setenv("RATE_LIMITING_LIMIT", "100")
3233

3334
defer os.Clearenv()
3435

3536
cfg := NewConfig("nonexistent.yaml") // Load with environment variables
3637

37-
assert.Equal(t, "127.0.0.1", cfg.Host)
38-
assert.Equal(t, 8080, cfg.Port)
39-
assert.Equal(t, 15, cfg.KeyLength)
40-
assert.Equal(t, 5000000, cfg.MaxLength)
41-
assert.Equal(t, "redis", cfg.Storage.Type)
42-
assert.Equal(t, "localhost", cfg.Storage.Host)
43-
assert.Equal(t, 6379, cfg.Storage.Port)
44-
assert.Equal(t, "debug", cfg.Logging.Level)
38+
require.Equal(t, "127.0.0.1", cfg.Host)
39+
require.Equal(t, 8080, cfg.Port)
40+
require.Equal(t, 15, cfg.KeyLength)
41+
require.Equal(t, 5000000, cfg.MaxLength)
42+
require.Equal(t, "redis", cfg.Storage.Type)
43+
require.Equal(t, "localhost", cfg.Storage.Host)
44+
require.Equal(t, 6379, cfg.Storage.Port)
45+
require.Equal(t, "debug", cfg.Logging.Level)
46+
require.Equal(t, true, cfg.RateLimiting.Enable)
47+
require.Equal(t, 100, cfg.RateLimiting.Limit)
4548
}
4649

4750
func TestNewConfig_LoadFromYAML(t *testing.T) {
@@ -60,21 +63,20 @@ logging:
6063

6164
// Write to a temporary file
6265
tmpFile, err := os.CreateTemp("", "config_test_*.yaml")
63-
assert.NoError(t, err)
66+
require.NoError(t, err)
6467
defer os.Remove(tmpFile.Name())
6568

6669
_, err = tmpFile.Write([]byte(yamlContent))
67-
assert.NoError(t, err)
70+
require.NoError(t, err)
6871
tmpFile.Close()
6972

7073
cfg := NewConfig(tmpFile.Name()) // Load config from YAML file
7174

72-
assert.Equal(t, "192.168.1.1", cfg.Host)
73-
assert.Equal(t, 9090, cfg.Port)
74-
assert.Equal(t, 20, cfg.KeyLength)
75-
assert.Equal(t, "mongodb", cfg.Storage.Type)
76-
assert.Equal(t, "mongo.example.com", cfg.Storage.Host)
77-
assert.Equal(t, 27017, cfg.Storage.Port)
78-
assert.Equal(t, "warn", cfg.Logging.Level)
79-
assert.Equal(t, "json", cfg.Logging.Type)
75+
require.Equal(t, "192.168.1.1", cfg.Host)
76+
require.Equal(t, 9090, cfg.Port)
77+
require.Equal(t, 20, cfg.KeyLength)
78+
require.Equal(t, "mongodb", cfg.Storage.Type)
79+
require.Equal(t, "mongo.example.com", cfg.Storage.Host)
80+
require.Equal(t, 27017, cfg.Storage.Port)
81+
require.Equal(t, "warn", cfg.Logging.Level)
8082
}

0 commit comments

Comments
 (0)