mirror of
https://github.com/0ceanSlim/grain.git
synced 2024-11-24 17:33:22 +00:00
refactor more main funcs
This commit is contained in:
parent
ba95bc849b
commit
8c2e4ed926
@ -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 {
|
||||
|
@ -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
|
||||
|
@ -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,
|
||||
|
54
main.go
54
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)
|
||||
|
@ -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) {
|
||||
|
@ -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)
|
||||
|
Loading…
Reference in New Issue
Block a user