From 8c2e4ed926432f92ec208076e9294d7cb9f60332 Mon Sep 17 00:00:00 2001 From: 0ceanSlim Date: Tue, 30 Jul 2024 11:51:05 -0400 Subject: [PATCH] refactor more main funcs --- config/load.go | 4 ++++ config/rateLimiter.go | 21 +++++++++++++++++ config/sizeLimiter.go | 9 ++++++++ main.go | 54 ++++--------------------------------------- relay/db/db.go | 5 ++++ web/relay_metadata.go | 5 ++++ 6 files changed, 49 insertions(+), 49 deletions(-) diff --git a/config/load.go b/config/load.go index 33238fa..f94c098 100644 --- a/config/load.go +++ b/config/load.go @@ -53,6 +53,10 @@ type KindLimitConfig struct { Burst int `yaml:"burst"` } +func LoadConfiguration() (*Config, error) { + return LoadConfig("config.yml") +} + func LoadConfig(filename string) (*Config, error) { data, err := os.ReadFile(filename) if err != nil { diff --git a/config/rateLimiter.go b/config/rateLimiter.go index 84fdd2e..a8a6a14 100644 --- a/config/rateLimiter.go +++ b/config/rateLimiter.go @@ -31,6 +31,27 @@ type RateLimiter struct { var rateLimiterInstance *RateLimiter var once sync.Once +func SetupRateLimiter(cfg *Config) { + rateLimiter := 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 cfg.RateLimit.KindLimits { + rateLimiter.AddKindLimit(kindLimit.Kind, rate.Limit(kindLimit.Limit), kindLimit.Burst) + } + + for category, categoryLimit := range cfg.RateLimit.CategoryLimits { + rateLimiter.AddCategoryLimit(category, rate.Limit(categoryLimit.Limit), categoryLimit.Burst) + } + + SetRateLimiter(rateLimiter) +} + func SetRateLimiter(rl *RateLimiter) { once.Do(func() { rateLimiterInstance = rl diff --git a/config/sizeLimiter.go b/config/sizeLimiter.go index e9a703e..b50987b 100644 --- a/config/sizeLimiter.go +++ b/config/sizeLimiter.go @@ -10,6 +10,15 @@ type SizeLimiter struct { mu sync.RWMutex } +func SetupSizeLimiter(cfg *Config) { + sizeLimiter := NewSizeLimiter(cfg.RateLimit.MaxEventSize) + for _, kindSizeLimit := range cfg.RateLimit.KindSizeLimits { + sizeLimiter.AddKindSizeLimit(kindSizeLimit.Kind, kindSizeLimit.MaxSize) + } + + SetSizeLimiter(sizeLimiter) +} + func NewSizeLimiter(globalMaxSize int) *SizeLimiter { return &SizeLimiter{ globalMaxSize: globalMaxSize, diff --git a/main.go b/main.go index cd8098f..e58d257 100644 --- a/main.go +++ b/main.go @@ -11,25 +11,24 @@ import ( "grain/web" "golang.org/x/net/websocket" - "golang.org/x/time/rate" ) func main() { - cfg, err := loadConfiguration() + cfg, err := config.LoadConfiguration() if err != nil { log.Fatal("Error loading config: ", err) } - err = initializeDatabase(cfg) + err = db.InitializeDatabase(cfg) if err != nil { log.Fatal("Error initializing database: ", err) } defer db.DisconnectDB() - setupRateLimiter(cfg) - setupSizeLimiter(cfg) + config.SetupRateLimiter(cfg) + config.SetupSizeLimiter(cfg) - err = loadRelayMetadata() + err = web.LoadRelayMetadataJSON() if err != nil { log.Fatal("Failed to load relay metadata: ", err) } @@ -39,49 +38,6 @@ func main() { startServer(cfg, mux) } -func loadConfiguration() (*config.Config, error) { - return config.LoadConfig("config.yml") -} - -func initializeDatabase(config *config.Config) error { - _, err := db.InitDB(config.MongoDB.URI, config.MongoDB.Database) - return err -} - -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 cfg.RateLimit.KindLimits { - rateLimiter.AddKindLimit(kindLimit.Kind, rate.Limit(kindLimit.Limit), kindLimit.Burst) - } - - for category, categoryLimit := range cfg.RateLimit.CategoryLimits { - rateLimiter.AddCategoryLimit(category, rate.Limit(categoryLimit.Limit), categoryLimit.Burst) - } - - config.SetRateLimiter(rateLimiter) -} - -func setupSizeLimiter(cfg *config.Config) { - sizeLimiter := config.NewSizeLimiter(cfg.RateLimit.MaxEventSize) - for _, kindSizeLimit := range cfg.RateLimit.KindSizeLimits { - sizeLimiter.AddKindSizeLimit(kindSizeLimit.Kind, kindSizeLimit.MaxSize) - } - - config.SetSizeLimiter(sizeLimiter) -} - -func loadRelayMetadata() error { - return web.LoadRelayMetadata("relay_metadata.json") -} - func setupRoutes() *http.ServeMux { mux := http.NewServeMux() mux.HandleFunc("/", ListenAndServe) diff --git a/relay/db/db.go b/relay/db/db.go index 9ab4250..eb4927b 100644 --- a/relay/db/db.go +++ b/relay/db/db.go @@ -3,6 +3,7 @@ package db import ( "context" "fmt" + "grain/config" "go.mongodb.org/mongo-driver/bson" "go.mongodb.org/mongo-driver/mongo" @@ -16,6 +17,10 @@ var collections = make(map[int]*mongo.Collection) func GetClient() *mongo.Client { return client } +func InitializeDatabase(config *config.Config) error { + _, err := InitDB(config.MongoDB.URI, config.MongoDB.Database) + return err +} // Initialize MongoDB client func InitDB(uri, database string) (*mongo.Client, error) { diff --git a/web/relay_metadata.go b/web/relay_metadata.go index 9ea200e..1e28ab8 100644 --- a/web/relay_metadata.go +++ b/web/relay_metadata.go @@ -18,6 +18,10 @@ type RelayMetadata struct { var relayMetadata RelayMetadata +func LoadRelayMetadataJSON() error { + return LoadRelayMetadata("relay_metadata.json") +} + func LoadRelayMetadata(filename string) error { data, err := os.ReadFile(filename) if err != nil { @@ -31,6 +35,7 @@ func LoadRelayMetadata(filename string) error { return nil } + func RelayInfoHandler(w http.ResponseWriter, r *http.Request) { if r.Header.Get("Accept") != "application/nostr+json" { http.Error(w, "Unsupported Media Type", http.StatusUnsupportedMediaType)