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: 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)

View File

@ -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
}

View File

@ -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
}) })
} }

View File

@ -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"`
}

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() 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
}