From c8ed954a9dac9675fc671ad9e209f8121c7ae05b Mon Sep 17 00:00:00 2001 From: Chris kerr Date: Sun, 1 Sep 2024 20:51:02 -0400 Subject: [PATCH] refactoring --- config/Blacklist.go | 204 +++++++++++++++++ .../checkWhitelist.go => config/Whitelist.go | 10 +- main.go | 8 +- server/db/{db.go => dbMongo.go} | 0 server/db/storeMongo.go | 48 ++++ server/handlers/event.go | 56 +---- server/utils/checkBlacklist.go | 210 ------------------ 7 files changed, 271 insertions(+), 265 deletions(-) create mode 100644 config/Blacklist.go rename server/utils/checkWhitelist.go => config/Whitelist.go (88%) rename server/db/{db.go => dbMongo.go} (100%) create mode 100644 server/db/storeMongo.go delete mode 100644 server/utils/checkBlacklist.go diff --git a/config/Blacklist.go b/config/Blacklist.go new file mode 100644 index 0000000..4adf638 --- /dev/null +++ b/config/Blacklist.go @@ -0,0 +1,204 @@ +package config + +import ( + "fmt" + types "grain/config/types" + "grain/server/utils" + "log" + "os" + "strings" + "sync" + "time" + + "gopkg.in/yaml.v2" +) + +// CheckBlacklist checks if a pubkey is in the blacklist based on event content +func CheckBlacklist(pubkey, eventContent string) (bool, string) { + blacklistConfig := GetConfig().Blacklist + + if !blacklistConfig.Enabled { + return false, "" + } + + log.Printf("Checking blacklist for pubkey: %s", pubkey) + + // Check for permanent blacklist by pubkey or npub + if isPubKeyPermanentlyBlacklisted(pubkey, blacklistConfig) { + log.Printf("Pubkey %s is permanently blacklisted", pubkey) + return true, fmt.Sprintf("pubkey %s is permanently blacklisted", pubkey) + } + + // Check for temporary ban + if isPubKeyTemporarilyBlacklisted(pubkey) { + log.Printf("Pubkey %s is temporarily blacklisted", pubkey) + return true, fmt.Sprintf("pubkey %s is temporarily blacklisted", pubkey) + } + + // Check for permanent ban based on wordlist + for _, word := range blacklistConfig.PermanentBanWords { + if strings.Contains(eventContent, word) { + err := AddToPermanentBlacklist(pubkey) + if err != nil { + return true, fmt.Sprintf("pubkey %s is permanently banned and failed to save: %v", pubkey, err) + } + return true, "blocked: pubkey is permanently banned" + } + } + + // Check for temporary ban based on wordlist + for _, word := range blacklistConfig.TempBanWords { + if strings.Contains(eventContent, word) { + err := AddToTemporaryBlacklist(pubkey, blacklistConfig) + if err != nil { + return true, fmt.Sprintf("pubkey %s is temporarily banned and failed to save: %v", pubkey, err) + } + return true, "blocked: pubkey is temporarily banned" + } + } + + return false, "" +} + +// Checks if a pubkey is temporarily blacklisted +func isPubKeyTemporarilyBlacklisted(pubkey string) bool { + mu.Lock() + defer mu.Unlock() + + entry, exists := tempBannedPubkeys[pubkey] + if !exists { + log.Printf("Pubkey %s not found in temporary blacklist", pubkey) + return false + } + + now := time.Now() + if now.After(entry.unbanTime) { + log.Printf("Temporary ban for pubkey %s has expired. Count: %d", pubkey, entry.count) + return false + } + + log.Printf("Pubkey %s is currently temporarily blacklisted. Count: %d, Unban time: %s", pubkey, entry.count, entry.unbanTime) + return true +} + +func ClearTemporaryBans() { + mu.Lock() + defer mu.Unlock() + tempBannedPubkeys = make(map[string]*tempBanEntry) +} + +var ( + tempBannedPubkeys = make(map[string]*tempBanEntry) + mu sync.Mutex +) + +type tempBanEntry struct { + count int + unbanTime time.Time +} + +// Adds a pubkey to the temporary blacklist +func AddToTemporaryBlacklist(pubkey string, blacklistConfig types.BlacklistConfig) error { + mu.Lock() + defer mu.Unlock() + + entry, exists := tempBannedPubkeys[pubkey] + if !exists { + log.Printf("Creating new temporary ban entry for pubkey %s", pubkey) + entry = &tempBanEntry{ + count: 0, + unbanTime: time.Now(), + } + tempBannedPubkeys[pubkey] = entry + } else { + log.Printf("Updating existing temporary ban entry for pubkey %s. Current count: %d", pubkey, entry.count) + if time.Now().After(entry.unbanTime) { + log.Printf("Previous ban for pubkey %s has expired. Keeping count at %d", pubkey, entry.count) + } + } + + // Increment the count + entry.count++ + entry.unbanTime = time.Now().Add(time.Duration(blacklistConfig.TempBanDuration) * time.Second) + + log.Printf("Pubkey %s temporary ban count updated to: %d, MaxTempBans: %d, New unban time: %s", pubkey, entry.count, blacklistConfig.MaxTempBans, entry.unbanTime) + + if entry.count > blacklistConfig.MaxTempBans { + log.Printf("Attempting to move pubkey %s to permanent blacklist", pubkey) + delete(tempBannedPubkeys, pubkey) + + // Release the lock before calling AddToPermanentBlacklist + mu.Unlock() + err := AddToPermanentBlacklist(pubkey) + mu.Lock() // Re-acquire the lock + + if err != nil { + log.Printf("Error adding pubkey %s to permanent blacklist: %v", pubkey, err) + return err + } + log.Printf("Successfully added pubkey %s to permanent blacklist", pubkey) + } + + return nil +} + +// Checks if a pubkey is permanently blacklisted (only using config.yml) +func isPubKeyPermanentlyBlacklisted(pubKey string, blacklistConfig types.BlacklistConfig) bool { + if !blacklistConfig.Enabled { + return false + } + + // Check pubkeys + for _, blacklistedKey := range blacklistConfig.PermanentBlacklistPubkeys { + if pubKey == blacklistedKey { + return true + } + } + + // Check npubs + for _, npub := range blacklistConfig.PermanentBlacklistNpubs { + decodedPubKey, err := utils.DecodeNpub(npub) + if err != nil { + fmt.Println("Error decoding npub:", err) + continue + } + if pubKey == decodedPubKey { + return true + } + } + + return false +} + +func AddToPermanentBlacklist(pubkey string) error { + // Remove the mutex lock from here + blacklistConfig := GetConfig().Blacklist + + // Check if already blacklisted + if isPubKeyPermanentlyBlacklisted(pubkey, blacklistConfig) { + return fmt.Errorf("pubkey %s is already in the permanent blacklist", pubkey) + } + + // Add pubkey to the blacklist + blacklistConfig.PermanentBlacklistPubkeys = append(blacklistConfig.PermanentBlacklistPubkeys, pubkey) + + // Persist changes to config.yml + return saveBlacklistConfig(blacklistConfig) +} + +func saveBlacklistConfig(blacklistConfig types.BlacklistConfig) error { + configData := GetConfig() + configData.Blacklist = blacklistConfig + + data, err := yaml.Marshal(configData) + if err != nil { + return fmt.Errorf("failed to marshal config: %v", err) + } + + err = os.WriteFile("config.yml", data, 0644) + if err != nil { + return fmt.Errorf("failed to write config to file: %v", err) + } + + return nil +} diff --git a/server/utils/checkWhitelist.go b/config/Whitelist.go similarity index 88% rename from server/utils/checkWhitelist.go rename to config/Whitelist.go index 8148e30..bb6be0b 100644 --- a/server/utils/checkWhitelist.go +++ b/config/Whitelist.go @@ -1,14 +1,14 @@ -package utils +package config import ( "fmt" - "grain/config" + "grain/server/utils" "strconv" ) // Helper function to check if a pubkey or npub is whitelisted func IsPubKeyWhitelisted(pubKey string) bool { - cfg := config.GetConfig() + cfg := GetConfig() if !cfg.PubkeyWhitelist.Enabled { return true } @@ -22,7 +22,7 @@ func IsPubKeyWhitelisted(pubKey string) bool { // Check npubs for _, npub := range cfg.PubkeyWhitelist.Npubs { - decodedPubKey, err := DecodeNpub(npub) + decodedPubKey, err := utils.DecodeNpub(npub) if err != nil { fmt.Println("Error decoding npub:", err) continue @@ -36,7 +36,7 @@ func IsPubKeyWhitelisted(pubKey string) bool { } func IsKindWhitelisted(kind int) bool { - cfg := config.GetConfig() + cfg := GetConfig() if !cfg.KindWhitelist.Enabled { return true } diff --git a/main.go b/main.go index 6143279..36780a8 100644 --- a/main.go +++ b/main.go @@ -49,7 +49,7 @@ func main() { config.SetupRateLimiter(cfg) config.SetupSizeLimiter(cfg) - utils.ClearTemporaryBans() + config.ClearTemporaryBans() err = utils.LoadRelayMetadataJSON() if err != nil { @@ -70,9 +70,9 @@ func main() { case <-signalChan: log.Println("Shutting down server...") - server.Close() // Stop the server - db.DisconnectDB(client) // Disconnect from MongoDB - wg.Wait() // Wait for all goroutines to finish + server.Close() // Stop the server + db.DisconnectDB(client) // Disconnect from MongoDB + wg.Wait() // Wait for all goroutines to finish return } } diff --git a/server/db/db.go b/server/db/dbMongo.go similarity index 100% rename from server/db/db.go rename to server/db/dbMongo.go diff --git a/server/db/storeMongo.go b/server/db/storeMongo.go new file mode 100644 index 0000000..21f643e --- /dev/null +++ b/server/db/storeMongo.go @@ -0,0 +1,48 @@ +package db + +import ( + "context" + "fmt" + "grain/server/handlers/kinds" + "grain/server/handlers/response" + nostr "grain/server/types" + + "golang.org/x/net/websocket" +) + +func StoreMongoEvent(ctx context.Context, evt nostr.Event, ws *websocket.Conn) { + collection := GetCollection(evt.Kind) + + var err error + switch { + case evt.Kind == 0: + err = kinds.HandleKind0(ctx, evt, collection, ws) + case evt.Kind == 1: + err = kinds.HandleKind1(ctx, evt, collection, ws) + case evt.Kind == 2: + err = kinds.HandleKind2(ctx, evt, ws) + case evt.Kind == 3: + err = kinds.HandleReplaceableKind(ctx, evt, collection, ws) + case evt.Kind == 5: + err = kinds.HandleKind5(ctx, evt, GetClient(), ws) + case evt.Kind >= 4 && evt.Kind < 45: + err = kinds.HandleRegularKind(ctx, evt, collection, ws) + case evt.Kind >= 1000 && evt.Kind < 10000: + err = kinds.HandleRegularKind(ctx, evt, collection, ws) + case evt.Kind >= 10000 && evt.Kind < 20000: + err = kinds.HandleReplaceableKind(ctx, evt, collection, ws) + case evt.Kind >= 20000 && evt.Kind < 30000: + fmt.Println("Ephemeral event received and ignored:", evt.ID) + case evt.Kind >= 30000 && evt.Kind < 40000: + err = kinds.HandleParameterizedReplaceableKind(ctx, evt, collection, ws) + default: + err = kinds.HandleUnknownKind(ctx, evt, collection, ws) + } + + if err != nil { + response.SendOK(ws, evt.ID, false, fmt.Sprintf("error: %v", err)) + return + } + + response.SendOK(ws, evt.ID, true, "") +} diff --git a/server/handlers/event.go b/server/handlers/event.go index 7518d43..22ebe84 100644 --- a/server/handlers/event.go +++ b/server/handlers/event.go @@ -6,11 +6,11 @@ import ( "fmt" "grain/config" "grain/server/db" - "grain/server/handlers/kinds" + "grain/server/handlers/response" "grain/server/utils" - relay "grain/server/types" + nostr "grain/server/types" "golang.org/x/net/websocket" ) @@ -36,7 +36,7 @@ func HandleEvent(ws *websocket.Conn, message []interface{}) { return } - var evt relay.Event + var evt nostr.Event err = json.Unmarshal(eventBytes, &evt) if err != nil { fmt.Println("Error unmarshaling event data:", err) @@ -60,13 +60,14 @@ func HandleEvent(ws *websocket.Conn, message []interface{}) { return } - storeEvent(context.TODO(), evt, ws) + // This is where I'll handle storage for multiple database types in the future + db.StoreMongoEvent(context.TODO(), evt, ws) fmt.Println("Event processed:", evt.ID) }) } -func handleBlacklistAndWhitelist(ws *websocket.Conn, evt relay.Event) bool { +func handleBlacklistAndWhitelist(ws *websocket.Conn, evt nostr.Event) bool { if config.GetConfig().DomainWhitelist.Enabled { domains := config.GetConfig().DomainWhitelist.Domains pubkeys, err := utils.FetchPubkeysFromDomains(domains) @@ -80,17 +81,17 @@ func handleBlacklistAndWhitelist(ws *websocket.Conn, evt relay.Event) bool { } } - if blacklisted, msg := utils.CheckBlacklist(evt.PubKey, evt.Content); blacklisted { + if blacklisted, msg := config.CheckBlacklist(evt.PubKey, evt.Content); blacklisted { response.SendOK(ws, evt.ID, false, msg) return false } - if config.GetConfig().KindWhitelist.Enabled && !utils.IsKindWhitelisted(evt.Kind) { + if config.GetConfig().KindWhitelist.Enabled && !config.IsKindWhitelisted(evt.Kind) { response.SendOK(ws, evt.ID, false, "not allowed: event kind is not whitelisted") return false } - if config.GetConfig().PubkeyWhitelist.Enabled && !utils.IsPubKeyWhitelisted(evt.PubKey) { + if config.GetConfig().PubkeyWhitelist.Enabled && !config.IsPubKeyWhitelisted(evt.PubKey) { response.SendOK(ws, evt.ID, false, "not allowed: pubkey or npub is not whitelisted") return false } @@ -98,7 +99,7 @@ func handleBlacklistAndWhitelist(ws *websocket.Conn, evt relay.Event) bool { return true } -func handleRateAndSizeLimits(ws *websocket.Conn, evt relay.Event, eventSize int) bool { +func handleRateAndSizeLimits(ws *websocket.Conn, evt nostr.Event, eventSize int) bool { rateLimiter := config.GetRateLimiter() sizeLimiter := config.GetSizeLimiter() category := determineCategory(evt.Kind) @@ -116,43 +117,6 @@ func handleRateAndSizeLimits(ws *websocket.Conn, evt relay.Event, eventSize int) return true } -func storeEvent(ctx context.Context, evt relay.Event, ws *websocket.Conn) { - collection := db.GetCollection(evt.Kind) - - var err error - switch { - case evt.Kind == 0: - err = kinds.HandleKind0(ctx, evt, collection, ws) - case evt.Kind == 1: - err = kinds.HandleKind1(ctx, evt, collection, ws) - case evt.Kind == 2: - err = kinds.HandleKind2(ctx, evt, ws) - case evt.Kind == 3: - err = kinds.HandleReplaceableKind(ctx, evt, collection, ws) - case evt.Kind == 5: - err = kinds.HandleKind5(ctx, evt, db.GetClient(), ws) - case evt.Kind >= 4 && evt.Kind < 45: - err = kinds.HandleRegularKind(ctx, evt, collection, ws) - case evt.Kind >= 1000 && evt.Kind < 10000: - err = kinds.HandleRegularKind(ctx, evt, collection, ws) - case evt.Kind >= 10000 && evt.Kind < 20000: - err = kinds.HandleReplaceableKind(ctx, evt, collection, ws) - case evt.Kind >= 20000 && evt.Kind < 30000: - fmt.Println("Ephemeral event received and ignored:", evt.ID) - case evt.Kind >= 30000 && evt.Kind < 40000: - err = kinds.HandleParameterizedReplaceableKind(ctx, evt, collection, ws) - default: - err = kinds.HandleUnknownKind(ctx, evt, collection, ws) - } - - if err != nil { - response.SendOK(ws, evt.ID, false, fmt.Sprintf("error: %v", err)) - return - } - - response.SendOK(ws, evt.ID, true, "") -} - func determineCategory(kind int) string { switch { case kind == 0, kind == 3, kind >= 10000 && kind < 20000: diff --git a/server/utils/checkBlacklist.go b/server/utils/checkBlacklist.go deleted file mode 100644 index 4112100..0000000 --- a/server/utils/checkBlacklist.go +++ /dev/null @@ -1,210 +0,0 @@ -package utils - -import ( - "fmt" - "grain/config" - cfg "grain/config/types" - "log" - "os" - "strings" - "sync" - "time" - - "gopkg.in/yaml.v2" -) - -// CheckBlacklist checks if a pubkey is in the blacklist based on event content -func CheckBlacklist(pubkey, eventContent string) (bool, string) { - cfg := config.GetConfig().Blacklist - - if !cfg.Enabled { - return false, "" - } - - log.Printf("Checking blacklist for pubkey: %s", pubkey) - - // Check for permanent blacklist by pubkey or npub - if isPubKeyPermanentlyBlacklisted(pubkey) { - log.Printf("Pubkey %s is permanently blacklisted", pubkey) - return true, fmt.Sprintf("pubkey %s is permanently blacklisted", pubkey) - } - - // Check for temporary ban - if isPubKeyTemporarilyBlacklisted(pubkey) { - log.Printf("Pubkey %s is temporarily blacklisted", pubkey) - return true, fmt.Sprintf("pubkey %s is temporarily blacklisted", pubkey) - } - - // Check for permanent ban based on wordlist - for _, word := range cfg.PermanentBanWords { - if strings.Contains(eventContent, word) { - err := AddToPermanentBlacklist(pubkey) - if err != nil { - return true, fmt.Sprintf("pubkey %s is permanently banned and failed to save: %v", pubkey, err) - } - return true, "blocked: pubkey is permanently banned" - } - } - - // Check for temporary ban based on wordlist - for _, word := range cfg.TempBanWords { - if strings.Contains(eventContent, word) { - err := AddToTemporaryBlacklist(pubkey) - if err != nil { - return true, fmt.Sprintf("pubkey %s is temporarily banned and failed to save: %v", pubkey, err) - } - return true, "blocked: pubkey is temporarily banned" - } - } - - return false, "" -} - - -// Checks if a pubkey is temporarily blacklisted -func isPubKeyTemporarilyBlacklisted(pubkey string) bool { - mu.Lock() - defer mu.Unlock() - - entry, exists := tempBannedPubkeys[pubkey] - if !exists { - log.Printf("Pubkey %s not found in temporary blacklist", pubkey) - return false - } - - now := time.Now() - if now.After(entry.unbanTime) { - log.Printf("Temporary ban for pubkey %s has expired. Count: %d", pubkey, entry.count) - return false - } - - log.Printf("Pubkey %s is currently temporarily blacklisted. Count: %d, Unban time: %s", pubkey, entry.count, entry.unbanTime) - return true -} - -func ClearTemporaryBans() { - mu.Lock() - defer mu.Unlock() - tempBannedPubkeys = make(map[string]*tempBanEntry) -} - -var ( - tempBannedPubkeys = make(map[string]*tempBanEntry) - mu sync.Mutex -) - -type tempBanEntry struct { - count int - unbanTime time.Time -} - -// Adds a pubkey to the temporary blacklist -func AddToTemporaryBlacklist(pubkey string) error { - mu.Lock() - defer mu.Unlock() - - cfg := config.GetConfig().Blacklist - - entry, exists := tempBannedPubkeys[pubkey] - if !exists { - log.Printf("Creating new temporary ban entry for pubkey %s", pubkey) - entry = &tempBanEntry{ - count: 0, - unbanTime: time.Now(), - } - tempBannedPubkeys[pubkey] = entry - } else { - log.Printf("Updating existing temporary ban entry for pubkey %s. Current count: %d", pubkey, entry.count) - if time.Now().After(entry.unbanTime) { - log.Printf("Previous ban for pubkey %s has expired. Keeping count at %d", pubkey, entry.count) - } - } - - // Increment the count - entry.count++ - entry.unbanTime = time.Now().Add(time.Duration(cfg.TempBanDuration) * time.Second) - - log.Printf("Pubkey %s temporary ban count updated to: %d, MaxTempBans: %d, New unban time: %s", pubkey, entry.count, cfg.MaxTempBans, entry.unbanTime) - - if entry.count > cfg.MaxTempBans { - log.Printf("Attempting to move pubkey %s to permanent blacklist", pubkey) - delete(tempBannedPubkeys, pubkey) - - // Release the lock before calling AddToPermanentBlacklist - mu.Unlock() - err := AddToPermanentBlacklist(pubkey) - mu.Lock() // Re-acquire the lock - - if err != nil { - log.Printf("Error adding pubkey %s to permanent blacklist: %v", pubkey, err) - return err - } - log.Printf("Successfully added pubkey %s to permanent blacklist", pubkey) - } - - return nil -} - -// Checks if a pubkey is permanently blacklisted (only using config.yml) -func isPubKeyPermanentlyBlacklisted(pubKey string) bool { - cfg := config.GetConfig().Blacklist // Get the latest configuration - - if !cfg.Enabled { - return false - } - - // Check pubkeys - for _, blacklistedKey := range cfg.PermanentBlacklistPubkeys { - if pubKey == blacklistedKey { - return true - } - } - - // Check npubs - for _, npub := range cfg.PermanentBlacklistNpubs { - decodedPubKey, err := DecodeNpub(npub) - if err != nil { - fmt.Println("Error decoding npub:", err) - continue - } - if pubKey == decodedPubKey { - return true - } - } - - return false -} - -func AddToPermanentBlacklist(pubkey string) error { - // Remove the mutex lock from here - cfg := config.GetConfig().Blacklist - - // Check if already blacklisted - if isPubKeyPermanentlyBlacklisted(pubkey) { - return fmt.Errorf("pubkey %s is already in the permanent blacklist", pubkey) - } - - // Add pubkey to the blacklist - cfg.PermanentBlacklistPubkeys = append(cfg.PermanentBlacklistPubkeys, pubkey) - - // Persist changes to config.yml - return saveBlacklistConfig(cfg) -} - -func saveBlacklistConfig(blacklistConfig cfg.BlacklistConfig) error { - configData := config.GetConfig() - configData.Blacklist = blacklistConfig - - data, err := yaml.Marshal(configData) - if err != nil { - return fmt.Errorf("failed to marshal config: %v", err) - } - - err = os.WriteFile("config.yml", data, 0644) - if err != nil { - return fmt.Errorf("failed to write config to file: %v", err) - } - - return nil -} -