pubkey whitelist!

This commit is contained in:
Chris kerr 2024-08-03 14:27:58 -04:00
parent 8bbf442493
commit 9d8db7df08
6 changed files with 53 additions and 10 deletions

View File

@ -4,7 +4,10 @@ mongodb:
server:
port: ":8080" # Port for the server to listen on
whitelist:
enabled: false
pubkeys: #["3fe0ab6cbdb7ee27148202249e3fb3b89423c6f6cda6ef43ea5057c3d93088e4",
#"cac0e43235806da094f0787a5b04e29ad04cb1a3c7ea5cf61edc1c338734082b"]
rate_limit:
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)

View File

@ -2,24 +2,37 @@ package config
import (
"os"
"sync"
config "grain/config/types"
configTypes "grain/config/types"
"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)
if err != nil {
return nil, err
}
var config config.ServerConfig
var config configTypes.ServerConfig
err = yaml.Unmarshal(data, &config)
if err != nil {
return nil, err
}
return &config, nil
}
once.Do(func() {
cfg = &config
})
return cfg, nil
}
func GetConfig() *configTypes.ServerConfig {
return cfg
}

View File

@ -30,7 +30,7 @@ type RateLimiter struct {
}
var rateLimiterInstance *RateLimiter
var once sync.Once
var rateOnce sync.Once
func SetupRateLimiter(cfg *config.ServerConfig) {
rateLimiter := NewRateLimiter(
@ -54,7 +54,7 @@ func SetupRateLimiter(cfg *config.ServerConfig) {
}
func SetRateLimiter(rl *RateLimiter) {
once.Do(func() {
rateOnce.Do(func() {
rateLimiterInstance = rl
})
}

View File

@ -9,4 +9,5 @@ type ServerConfig struct {
Port string `yaml:"port"`
} `yaml:"server"`
RateLimit RateLimitConfig `yaml:"rate_limit"`
}
Whitelist WhitelistConfig `yaml:"whitelist"`
}

View File

@ -0,0 +1,6 @@
package config
type WhitelistConfig struct {
Enabled bool `yaml:"enabled"`
Pubkeys []string `yaml:"pubkeys"`
}

View File

@ -59,6 +59,12 @@ func HandleKind(ctx context.Context, evt relay.Event, ws *websocket.Conn, eventS
rateLimiter := config.GetRateLimiter()
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)
if allowed, msg := rateLimiter.AllowEvent(evt.Kind, category); !allowed {
@ -121,3 +127,17 @@ func determineCategory(kind int) string {
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
}