OceanSlim/grain
1
0
Fork 0
mirror of https://github.com/0ceanSlim/grain.git synced 2024-11-23 00:57:14 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

configuration refactor

Browse Source
This commit is contained in:
0ceanSlim 2024-07-30 11:27:38 -04:00
parent 567d9010a4
commit ba95bc849b
9 changed files with 46 additions and 44 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
relay/utils/loadConfig.go → config/load.go
View File

@ -1,4 +1,4 @@
package utils
package config
import (
"os"

2
relay/utils/rateLimiter.go → config/rateLimiter.go
View File

@ -1,4 +1,4 @@
package utils
package config
import (
"fmt"

2
relay/utils/sizeLimiter.go → config/sizeLimiter.go
View File

@ -1,4 +1,4 @@
package utils
package config
import (
"sync"

50
main.go
View File

@ -5,9 +5,9 @@ import (
"log"
"net/http"
"grain/config"
"grain/relay"
"grain/relay/db"
"grain/relay/utils"
"grain/web"
"golang.org/x/net/websocket"
@ -15,19 +15,19 @@ import (
)
func main() {
config, err := loadConfiguration()
cfg, err := loadConfiguration()
if err != nil {
log.Fatal("Error loading config: ", err)
}
err = initializeDatabase(config)
err = initializeDatabase(cfg)
if err != nil {
log.Fatal("Error initializing database: ", err)
}
defer db.DisconnectDB()
setupRateLimiter(config)
setupSizeLimiter(config)
setupRateLimiter(cfg)
setupSizeLimiter(cfg)
err = loadRelayMetadata()
if err != nil {
@ -36,46 +36,46 @@ func main() {
mux := setupRoutes()
startServer(config, mux)
startServer(cfg, mux)
}
func loadConfiguration() (*utils.Config, error) {
return utils.LoadConfig("config.yml")
func loadConfiguration() (*config.Config, error) {
return config.LoadConfig("config.yml")
}
func initializeDatabase(config *utils.Config) error {
func initializeDatabase(config *config.Config) error {
_, err := db.InitDB(config.MongoDB.URI, config.MongoDB.Database)
return err
}
func setupRateLimiter(config *utils.Config) {
rateLimiter := utils.NewRateLimiter(
rate.Limit(config.RateLimit.WsLimit),
config.RateLimit.WsBurst,
rate.Limit(config.RateLimit.EventLimit),
config.RateLimit.EventBurst,
rate.Limit(config.RateLimit.ReqLimit),
config.RateLimit.ReqBurst,
func setupRateLimiter(cfg *config.Config) {
rateLimiter := config.NewRateLimiter(
rate.Limit(cfg.RateLimit.WsLimit),
cfg.RateLimit.WsBurst,
rate.Limit(cfg.RateLimit.EventLimit),
cfg.RateLimit.EventBurst,
rate.Limit(cfg.RateLimit.ReqLimit),
cfg.RateLimit.ReqBurst,
)
for _, kindLimit := range config.RateLimit.KindLimits {
for _, kindLimit := range cfg.RateLimit.KindLimits {
rateLimiter.AddKindLimit(kindLimit.Kind, rate.Limit(kindLimit.Limit), kindLimit.Burst)
}
for category, categoryLimit := range config.RateLimit.CategoryLimits {
for category, categoryLimit := range cfg.RateLimit.CategoryLimits {
rateLimiter.AddCategoryLimit(category, rate.Limit(categoryLimit.Limit), categoryLimit.Burst)
}
utils.SetRateLimiter(rateLimiter)
config.SetRateLimiter(rateLimiter)
}
func setupSizeLimiter(config *utils.Config) {
sizeLimiter := utils.NewSizeLimiter(config.RateLimit.MaxEventSize)
for _, kindSizeLimit := range config.RateLimit.KindSizeLimits {
func setupSizeLimiter(cfg *config.Config) {
sizeLimiter := config.NewSizeLimiter(cfg.RateLimit.MaxEventSize)
for _, kindSizeLimit := range cfg.RateLimit.KindSizeLimits {
sizeLimiter.AddKindSizeLimit(kindSizeLimit.Kind, kindSizeLimit.MaxSize)
}
utils.SetSizeLimiter(sizeLimiter)
config.SetSizeLimiter(sizeLimiter)
}
func loadRelayMetadata() error {
@ -92,7 +92,7 @@ func setupRoutes() *http.ServeMux {
return mux
}
func startServer(config *utils.Config, mux *http.ServeMux) {
func startServer(config *config.Config, mux *http.ServeMux) {
fmt.Printf("Server is running on http://localhost%s\n", config.Server.Port)
err := http.ListenAndServe(config.Server.Port, mux)
if err != nil {

5
relay/handlers/event.go
View File

@ -4,6 +4,7 @@ import (
"context"
"encoding/json"
"fmt"
"grain/config"
"grain/relay/db"
"grain/relay/handlers/kinds"
"grain/relay/handlers/response"
@ -55,8 +56,8 @@ func HandleKind(ctx context.Context, evt relay.Event, ws *websocket.Conn, eventS
}
collection := db.GetCollection(evt.Kind)
rateLimiter := utils.GetRateLimiter()
sizeLimiter := utils.GetSizeLimiter()
rateLimiter := config.GetRateLimiter()
sizeLimiter := config.GetSizeLimiter()
category := determineCategory(evt.Kind)

5
relay/relay.go
View File

@ -4,7 +4,8 @@ import (
"encoding/json"
"fmt"
"grain/relay/handlers"
"grain/relay/utils"
"grain/config"
"golang.org/x/net/websocket"
)
@ -13,7 +14,7 @@ func WebSocketHandler(ws *websocket.Conn) {
defer ws.Close()
var msg string
rateLimiter := utils.GetRateLimiter()
rateLimiter := config.GetRateLimiter()
for {
err := websocket.Message.Receive(ws, &msg)

4
tests/config_test.go
View File

@ -3,11 +3,11 @@ package tests
import (
"testing"
"grain/relay/utils"
"grain/config"
)
func TestConfigValidity(t *testing.T) {
config, err := utils.LoadConfig("../config.yml")
config, err := config.LoadConfig("../config.yml")
if err != nil {
t.Fatalf("Error loading config: %v", err)
}

12
tests/rateLimits_test.go
View File

@ -3,13 +3,13 @@ package tests
import (
"testing"
"grain/relay/utils"
"grain/config"
"golang.org/x/time/rate"
)
func TestWebSocketRateLimit(t *testing.T) {
rateLimiter := utils.NewRateLimiter(rate.Limit(1), 1, rate.Limit(100), 200, rate.Limit(100), 200)
rateLimiter := config.NewRateLimiter(rate.Limit(1), 1, rate.Limit(100), 200, rate.Limit(100), 200)
// First message should be allowed
if allowed, _ := rateLimiter.AllowWs(); !allowed {
@ -28,7 +28,7 @@ func TestWebSocketRateLimit(t *testing.T) {
}
func TestEventRateLimit(t *testing.T) {
rateLimiter := utils.NewRateLimiter(rate.Limit(100), 200, rate.Limit(1), 1, rate.Limit(100), 200)
rateLimiter := config.NewRateLimiter(rate.Limit(100), 200, rate.Limit(1), 1, rate.Limit(100), 200)
rateLimiter.AddKindLimit(1, rate.Limit(1), 1)
rateLimiter.AddCategoryLimit("regular", rate.Limit(1), 1)
@ -49,7 +49,7 @@ func TestEventRateLimit(t *testing.T) {
}
func TestReqRateLimit(t *testing.T) {
rateLimiter := utils.NewRateLimiter(rate.Limit(100), 200, rate.Limit(100), 200, rate.Limit(1), 1)
rateLimiter := config.NewRateLimiter(rate.Limit(100), 200, rate.Limit(100), 200, rate.Limit(1), 1)
// First REQ should be allowed
if allowed, _ := rateLimiter.AllowReq(); !allowed {
@ -68,7 +68,7 @@ func TestReqRateLimit(t *testing.T) {
}
func TestKindRateLimit(t *testing.T) {
rateLimiter := utils.NewRateLimiter(rate.Limit(100), 200, rate.Limit(100), 200, rate.Limit(100), 200)
rateLimiter := config.NewRateLimiter(rate.Limit(100), 200, rate.Limit(100), 200, rate.Limit(100), 200)
rateLimiter.AddKindLimit(1, rate.Limit(1), 1)
// First event of kind 1 should be allowed
@ -88,7 +88,7 @@ func TestKindRateLimit(t *testing.T) {
}
func TestCategoryRateLimit(t *testing.T) {
rateLimiter := utils.NewRateLimiter(rate.Limit(100), 200, rate.Limit(100), 200, rate.Limit(100), 200)
rateLimiter := config.NewRateLimiter(rate.Limit(100), 200, rate.Limit(100), 200, rate.Limit(100), 200)
rateLimiter.AddCategoryLimit("regular", rate.Limit(1), 1)
// First event in category "regular" should be allowed

8
tests/sizeLimits_test.go
View File

@ -1,12 +1,12 @@
package tests
import (
"grain/relay/utils"
"grain/config"
"testing"
)
func TestSizeLimiterGlobalMaxSize(t *testing.T) {
sizeLimiter := utils.NewSizeLimiter(1024) // Set global max size to 1024 bytes
sizeLimiter := config.NewSizeLimiter(1024) // Set global max size to 1024 bytes
// Test that an event within the global max size is allowed
if allowed, _ := sizeLimiter.AllowSize(0, 512); !allowed {
@ -25,7 +25,7 @@ func TestSizeLimiterGlobalMaxSize(t *testing.T) {
}
func TestSizeLimiterKindSpecificSize(t *testing.T) {
sizeLimiter := utils.NewSizeLimiter(1024) // Set global max size to 1024 bytes
sizeLimiter := config.NewSizeLimiter(1024) // Set global max size to 1024 bytes
sizeLimiter.AddKindSizeLimit(1, 512) // Set max size for kind 1 to 512 bytes
// Test that an event within the kind-specific max size is allowed
@ -55,7 +55,7 @@ func TestSizeLimiterKindSpecificSize(t *testing.T) {
}
func TestSizeLimiterNoKindSpecificLimit(t *testing.T) {
sizeLimiter := utils.NewSizeLimiter(1024) // Set global max size to 1024 bytes
sizeLimiter := config.NewSizeLimiter(1024) // Set global max size to 1024 bytes
// Test that an event for a kind without a specific limit is governed by the global limit
if allowed, _ := sizeLimiter.AllowSize(2, 512); !allowed {