From 85ec58cf8e7ac5ec966ca2093e4c6790e61f3b40 Mon Sep 17 00:00:00 2001
From: fiatjaf <fiatjaf@gmail.com>
Date: Thu, 2 Nov 2023 20:31:32 -0300
Subject: [PATCH] RestrictToSpecifiedKinds() filter.

---
 plugins/events.go | 36 ++++++++++++++++++++++++++++++++++++
 1 file changed, 36 insertions(+)

diff --git a/plugins/events.go b/plugins/events.go
index 61db6f4..1c6e748 100644
--- a/plugins/events.go
+++ b/plugins/events.go
@@ -6,6 +6,8 @@ import (
 	"github.com/nbd-wtf/go-nostr"
 )
 
+// PreventTooManyIndexableTags returns a function that can be used as a RejectFilter that will reject
+// events with more indexable (single-character) tags than the specified number.
 func PreventTooManyIndexableTags(max int) func(context.Context, *nostr.Event) (bool, string) {
 	return func(ctx context.Context, event *nostr.Event) (reject bool, msg string) {
 		ntags := 0
@@ -20,3 +22,37 @@ func PreventTooManyIndexableTags(max int) func(context.Context, *nostr.Event) (b
 		return false, ""
 	}
 }
+
+// RestrictToSpecifiedKinds returns a function that can be used as a RejectFilter that will reject
+// any events with kinds different than the specified ones.
+func RestrictToSpecifiedKinds(kinds ...uint16) func(context.Context, *nostr.Event) (bool, string) {
+	max := 0
+	min := 0
+	allowed := make(map[uint16]struct{}, len(kinds))
+	for _, kind := range kinds {
+		allowed[kind] = struct{}{}
+		if int(kind) > max {
+			max = int(kind)
+		}
+		if int(kind) < min {
+			min = int(kind)
+		}
+	}
+
+	return func(ctx context.Context, event *nostr.Event) (reject bool, msg string) {
+		// these are cheap and very questionable optimizations, but they exist for a reason:
+		// we would have to ensure that the kind number is within the bounds of a uint16 anyway
+		if event.Kind > max {
+			return true, "event kind not allowed"
+		}
+		if event.Kind < min {
+			return true, "event kind not allowed"
+		}
+
+		// hopefully this map of uint16s is very fast
+		if _, allowed := allowed[uint16(event.Kind)]; allowed {
+			return false, ""
+		}
+		return true, "event kind not allowed"
+	}
+}