From b6f77a48bf7ff655e096052caa3c8ca8f31da63f Mon Sep 17 00:00:00 2001
From: bndw <ben@bdw.to>
Date: Wed, 19 Jul 2023 16:37:28 -0700
Subject: [PATCH] feat: allowed pubkeys list

---
 config.go        | 15 ++++++++-------
 local/config.yml |  1 +
 relay.go         | 22 ++++++++++++++++++++++
 relay_test.go    | 26 ++++++++++++++++++++++++++
 4 files changed, 57 insertions(+), 7 deletions(-)
 create mode 100644 relay_test.go

diff --git a/config.go b/config.go
index 06a7cda..6b46d37 100644
--- a/config.go
+++ b/config.go
@@ -7,13 +7,14 @@ import (
 )
 
 type Config struct {
-	Port             int    `yaml:"port"`
-	DatabaseURL      string `yaml:"database_url"`
-	AllowedKinds     []int  `yaml:"allowed_kinds"`
-	Nip11Pubkey      string `yaml:"nip11_pubkey"`
-	Nip11Contact     string `yaml:"nip11_contact"`
-	Nip11Description string `yaml:"nip11_description"`
-	Nip11Version     string `yaml:"nip11_version"`
+	Port             int      `yaml:"port"`
+	DatabaseURL      string   `yaml:"database_url"`
+	AllowedKinds     []int    `yaml:"allowed_kinds"`
+	AllowedPubkeys   []string `yaml:"allowed_pubkeys"`
+	Nip11Pubkey      string   `yaml:"nip11_pubkey"`
+	Nip11Contact     string   `yaml:"nip11_contact"`
+	Nip11Description string   `yaml:"nip11_description"`
+	Nip11Version     string   `yaml:"nip11_version"`
 }
 
 // Load Config from a yaml file at path.
diff --git a/local/config.yml b/local/config.yml
index 16e8f84..0ce4693 100644
--- a/local/config.yml
+++ b/local/config.yml
@@ -5,3 +5,4 @@ nip11_description: custom relay for https://stemstr.app
 nip11_version: 0.1
 database_url: "postgresql://admin:password@localhost:5432/relay?sslmode=disable"
 allowed_kinds: []
+allowed_pubkeys: []
diff --git a/relay.go b/relay.go
index a396771..9d4da56 100644
--- a/relay.go
+++ b/relay.go
@@ -5,6 +5,7 @@ import (
 	"encoding/json"
 	"fmt"
 	"log"
+	"strings"
 
 	"github.com/fiatjaf/relayer/v2"
 	"github.com/fiatjaf/relayer/v2/storage/postgresql"
@@ -70,6 +71,10 @@ func (r Relay) Init() error {
 }
 
 func (r Relay) AcceptEvent(ctx context.Context, evt *nostr.Event) bool {
+	if !pubkeyIsAllowed(r.cfg.AllowedPubkeys, evt.PubKey) {
+		return false
+	}
+
 	// block events that are too large
 	jsonb, _ := json.Marshal(evt)
 	if len(jsonb) > 10000 {
@@ -106,6 +111,23 @@ func (r Relay) Start() error {
 	return server.Start("0.0.0.0", r.cfg.Port)
 }
 
+func pubkeyIsAllowed(pubkeys []string, pubkey string) bool {
+	// If no whitelist of pubkeys are provided, it's allowed
+	if len(pubkeys) == 0 {
+		return true
+	}
+
+	allowed := false
+	for _, allowedPubkey := range pubkeys {
+		if strings.EqualFold(allowedPubkey, pubkey) {
+			allowed = true
+			break
+		}
+	}
+
+	return allowed
+}
+
 var defaultAllowedKinds = []int{
 	nostr.KindSetMetadata,
 	nostr.KindTextNote,
diff --git a/relay_test.go b/relay_test.go
new file mode 100644
index 0000000..1f1c085
--- /dev/null
+++ b/relay_test.go
@@ -0,0 +1,26 @@
+package main
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestPubkeyIsAllowed(t *testing.T) {
+	var tests = []struct {
+		name     string
+		pubkeys  []string
+		pubkey   string
+		expected bool
+	}{
+		{"default open", []string{}, "123", true},
+		{"whitelisted", []string{"123"}, "123", true},
+		{"whitelisted", []string{"123", "456"}, "123", true},
+		{"not whitelisted", []string{"123", "456"}, "789", false},
+	}
+
+	for _, tt := range tests {
+		result := pubkeyIsAllowed(tt.pubkeys, tt.pubkey)
+		assert.Equal(t, tt.expected, result)
+	}
+}