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"`
|
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 {
|
||||||
|
@ -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
|
||||||
|
@ -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
54
main.go
@ -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)
|
||||||
|
@ -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) {
|
||||||
|
@ -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)
|
||||||
|
Loading…
Reference in New Issue
Block a user