mirror of
https://github.com/0ceanSlim/grain.git
synced 2024-11-22 00:27:14 +00:00
pubkey whitelist!
This commit is contained in:
parent
8bbf442493
commit
9d8db7df08
@ -4,7 +4,10 @@ mongodb:
|
|||||||
|
|
||||||
server:
|
server:
|
||||||
port: ":8080" # Port for the server to listen on
|
port: ":8080" # Port for the server to listen on
|
||||||
|
whitelist:
|
||||||
|
enabled: false
|
||||||
|
pubkeys: #["3fe0ab6cbdb7ee27148202249e3fb3b89423c6f6cda6ef43ea5057c3d93088e4",
|
||||||
|
#"cac0e43235806da094f0787a5b04e29ad04cb1a3c7ea5cf61edc1c338734082b"]
|
||||||
rate_limit:
|
rate_limit:
|
||||||
ws_limit: 100 # Global rate limit for WebSocket messages (50 messages per second)
|
ws_limit: 100 # Global rate limit for WebSocket messages (50 messages per second)
|
||||||
ws_burst: 200 # Global burst limit for WebSocket messages (allows a burst of 100 messages)
|
ws_burst: 200 # Global burst limit for WebSocket messages (allows a burst of 100 messages)
|
||||||
|
@ -2,24 +2,37 @@ package config
|
|||||||
|
|
||||||
import (
|
import (
|
||||||
"os"
|
"os"
|
||||||
|
"sync"
|
||||||
|
|
||||||
config "grain/config/types"
|
configTypes "grain/config/types"
|
||||||
|
|
||||||
"gopkg.in/yaml.v2"
|
"gopkg.in/yaml.v2"
|
||||||
)
|
)
|
||||||
|
|
||||||
func LoadConfig(filename string) (*config.ServerConfig, error) {
|
var (
|
||||||
|
cfg *configTypes.ServerConfig
|
||||||
|
once sync.Once
|
||||||
|
)
|
||||||
|
|
||||||
|
func LoadConfig(filename string) (*configTypes.ServerConfig, error) {
|
||||||
data, err := os.ReadFile(filename)
|
data, err := os.ReadFile(filename)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return nil, err
|
return nil, err
|
||||||
}
|
}
|
||||||
|
|
||||||
var config config.ServerConfig
|
var config configTypes.ServerConfig
|
||||||
|
|
||||||
err = yaml.Unmarshal(data, &config)
|
err = yaml.Unmarshal(data, &config)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return nil, err
|
return nil, err
|
||||||
}
|
}
|
||||||
|
|
||||||
return &config, nil
|
once.Do(func() {
|
||||||
}
|
cfg = &config
|
||||||
|
})
|
||||||
|
|
||||||
|
return cfg, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func GetConfig() *configTypes.ServerConfig {
|
||||||
|
return cfg
|
||||||
|
}
|
||||||
|
@ -30,7 +30,7 @@ type RateLimiter struct {
|
|||||||
}
|
}
|
||||||
|
|
||||||
var rateLimiterInstance *RateLimiter
|
var rateLimiterInstance *RateLimiter
|
||||||
var once sync.Once
|
var rateOnce sync.Once
|
||||||
|
|
||||||
func SetupRateLimiter(cfg *config.ServerConfig) {
|
func SetupRateLimiter(cfg *config.ServerConfig) {
|
||||||
rateLimiter := NewRateLimiter(
|
rateLimiter := NewRateLimiter(
|
||||||
@ -54,7 +54,7 @@ func SetupRateLimiter(cfg *config.ServerConfig) {
|
|||||||
}
|
}
|
||||||
|
|
||||||
func SetRateLimiter(rl *RateLimiter) {
|
func SetRateLimiter(rl *RateLimiter) {
|
||||||
once.Do(func() {
|
rateOnce.Do(func() {
|
||||||
rateLimiterInstance = rl
|
rateLimiterInstance = rl
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
|
@ -9,4 +9,5 @@ type ServerConfig struct {
|
|||||||
Port string `yaml:"port"`
|
Port string `yaml:"port"`
|
||||||
} `yaml:"server"`
|
} `yaml:"server"`
|
||||||
RateLimit RateLimitConfig `yaml:"rate_limit"`
|
RateLimit RateLimitConfig `yaml:"rate_limit"`
|
||||||
}
|
Whitelist WhitelistConfig `yaml:"whitelist"`
|
||||||
|
}
|
||||||
|
6
config/types/whitelistConfig.go
Normal file
6
config/types/whitelistConfig.go
Normal file
@ -0,0 +1,6 @@
|
|||||||
|
package config
|
||||||
|
|
||||||
|
type WhitelistConfig struct {
|
||||||
|
Enabled bool `yaml:"enabled"`
|
||||||
|
Pubkeys []string `yaml:"pubkeys"`
|
||||||
|
}
|
@ -59,6 +59,12 @@ func HandleKind(ctx context.Context, evt relay.Event, ws *websocket.Conn, eventS
|
|||||||
rateLimiter := config.GetRateLimiter()
|
rateLimiter := config.GetRateLimiter()
|
||||||
sizeLimiter := config.GetSizeLimiter()
|
sizeLimiter := config.GetSizeLimiter()
|
||||||
|
|
||||||
|
// Check whitelist
|
||||||
|
if !isWhitelisted(evt.PubKey) {
|
||||||
|
response.SendOK(ws, evt.ID, false, "not allowed: pubkey is not whitelisted")
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
category := determineCategory(evt.Kind)
|
category := determineCategory(evt.Kind)
|
||||||
|
|
||||||
if allowed, msg := rateLimiter.AllowEvent(evt.Kind, category); !allowed {
|
if allowed, msg := rateLimiter.AllowEvent(evt.Kind, category); !allowed {
|
||||||
@ -121,3 +127,17 @@ func determineCategory(kind int) string {
|
|||||||
return "unknown"
|
return "unknown"
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// Helper function to check if a pubkey is whitelisted
|
||||||
|
func isWhitelisted(pubKey string) bool {
|
||||||
|
cfg := config.GetConfig()
|
||||||
|
if !cfg.Whitelist.Enabled {
|
||||||
|
return true
|
||||||
|
}
|
||||||
|
for _, whitelistedKey := range cfg.Whitelist.Pubkeys {
|
||||||
|
if pubKey == whitelistedKey {
|
||||||
|
return true
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
Loading…
Reference in New Issue
Block a user