refactor more main funcs

This commit is contained in:
0ceanSlim 2024-07-30 11:51:05 -04:00
parent ba95bc849b
commit 8c2e4ed926
6 changed files with 49 additions and 49 deletions

View File

@ -53,6 +53,10 @@ type KindLimitConfig struct {
Burst int `yaml:"burst"` Burst int `yaml:"burst"`
} }
func LoadConfiguration() (*Config, error) {
return LoadConfig("config.yml")
}
func LoadConfig(filename string) (*Config, error) { func LoadConfig(filename string) (*Config, error) {
data, err := os.ReadFile(filename) data, err := os.ReadFile(filename)
if err != nil { if err != nil {

View File

@ -31,6 +31,27 @@ type RateLimiter struct {
var rateLimiterInstance *RateLimiter var rateLimiterInstance *RateLimiter
var once sync.Once 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) { func SetRateLimiter(rl *RateLimiter) {
once.Do(func() { once.Do(func() {
rateLimiterInstance = rl rateLimiterInstance = rl

View File

@ -10,6 +10,15 @@ type SizeLimiter struct {
mu sync.RWMutex 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 { func NewSizeLimiter(globalMaxSize int) *SizeLimiter {
return &SizeLimiter{ return &SizeLimiter{
globalMaxSize: globalMaxSize, globalMaxSize: globalMaxSize,

54
main.go
View File

@ -11,25 +11,24 @@ import (
"grain/web" "grain/web"
"golang.org/x/net/websocket" "golang.org/x/net/websocket"
"golang.org/x/time/rate"
) )
func main() { func main() {
cfg, err := loadConfiguration() cfg, err := config.LoadConfiguration()
if err != nil { if err != nil {
log.Fatal("Error loading config: ", err) log.Fatal("Error loading config: ", err)
} }
err = initializeDatabase(cfg) err = db.InitializeDatabase(cfg)
if err != nil { if err != nil {
log.Fatal("Error initializing database: ", err) log.Fatal("Error initializing database: ", err)
} }
defer db.DisconnectDB() defer db.DisconnectDB()
setupRateLimiter(cfg) config.SetupRateLimiter(cfg)
setupSizeLimiter(cfg) config.SetupSizeLimiter(cfg)
err = loadRelayMetadata() err = web.LoadRelayMetadataJSON()
if err != nil { if err != nil {
log.Fatal("Failed to load relay metadata: ", err) log.Fatal("Failed to load relay metadata: ", err)
} }
@ -39,49 +38,6 @@ func main() {
startServer(cfg, mux) 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 { func setupRoutes() *http.ServeMux {
mux := http.NewServeMux() mux := http.NewServeMux()
mux.HandleFunc("/", ListenAndServe) mux.HandleFunc("/", ListenAndServe)

View File

@ -3,6 +3,7 @@ package db
import ( import (
"context" "context"
"fmt" "fmt"
"grain/config"
"go.mongodb.org/mongo-driver/bson" "go.mongodb.org/mongo-driver/bson"
"go.mongodb.org/mongo-driver/mongo" "go.mongodb.org/mongo-driver/mongo"
@ -16,6 +17,10 @@ var collections = make(map[int]*mongo.Collection)
func GetClient() *mongo.Client { func GetClient() *mongo.Client {
return client return client
} }
func InitializeDatabase(config *config.Config) error {
_, err := InitDB(config.MongoDB.URI, config.MongoDB.Database)
return err
}
// Initialize MongoDB client // Initialize MongoDB client
func InitDB(uri, database string) (*mongo.Client, error) { func InitDB(uri, database string) (*mongo.Client, error) {

View File

@ -18,6 +18,10 @@ type RelayMetadata struct {
var relayMetadata RelayMetadata var relayMetadata RelayMetadata
func LoadRelayMetadataJSON() error {
return LoadRelayMetadata("relay_metadata.json")
}
func LoadRelayMetadata(filename string) error { func LoadRelayMetadata(filename string) error {
data, err := os.ReadFile(filename) data, err := os.ReadFile(filename)
if err != nil { if err != nil {
@ -31,6 +35,7 @@ func LoadRelayMetadata(filename string) error {
return nil return nil
} }
func RelayInfoHandler(w http.ResponseWriter, r *http.Request) { func RelayInfoHandler(w http.ResponseWriter, r *http.Request) {
if r.Header.Get("Accept") != "application/nostr+json" { if r.Header.Get("Accept") != "application/nostr+json" {
http.Error(w, "Unsupported Media Type", http.StatusUnsupportedMediaType) http.Error(w, "Unsupported Media Type", http.StatusUnsupportedMediaType)