From 2981e73f048b72b2b75b776ebbadd0228661deae Mon Sep 17 00:00:00 2001
From: Asara <amarpreet@minhas.io>
Date: Sat, 17 Aug 2024 15:18:55 -0400
Subject: [PATCH] chore: cleanup code

---
 db/db.go          |  3 +--
 nostr/helpers.go  | 11 +++++++++++
 nostr/policies.go | 43 +++++++++++++++++++++++++++++++++++++++++
 nostr/relay.go    | 49 -----------------------------------------------
 4 files changed, 55 insertions(+), 51 deletions(-)
 create mode 100644 nostr/helpers.go
 create mode 100644 nostr/policies.go

diff --git a/db/db.go b/db/db.go
index d115120..d32ca23 100644
--- a/db/db.go
+++ b/db/db.go
@@ -14,8 +14,7 @@ func NewDB() (*sqlx.DB, error) {
 	if err != nil {
 		l.Panic().Msg(err.Error())
 	}
-	err = db.Ping()
-	if err != nil {
+	if err = db.Ping(); err != nil {
 		l.Panic().Msg(err.Error())
 	}
 	l.Debug().Msg("connected to database")
diff --git a/nostr/helpers.go b/nostr/helpers.go
new file mode 100644
index 0000000..1865d75
--- /dev/null
+++ b/nostr/helpers.go
@@ -0,0 +1,11 @@
+package nostr
+
+func checknPubsInDb(npubs []string) bool {
+	var uid int
+	for _, npub := range npubs {
+		if err := DB.QueryRow("SELECT id FROM users WHERE pubkey=$1", npub).Scan(&uid); err == nil {
+			return true
+		}
+	}
+	return false
+}
diff --git a/nostr/policies.go b/nostr/policies.go
new file mode 100644
index 0000000..8f5f651
--- /dev/null
+++ b/nostr/policies.go
@@ -0,0 +1,43 @@
+package nostr
+
+import (
+	"context"
+	"fmt"
+
+	"git.devvul.com/asara/gologger"
+	"git.devvul.com/asara/well-goknown/config"
+	"github.com/fiatjaf/khatru"
+	"github.com/nbd-wtf/go-nostr"
+)
+
+func RejectUnregisteredNpubs(ctx context.Context, event *nostr.Event) (reject bool, msg string) {
+	var err error
+	l := gologger.Get(config.GetConfig().LogLevel).With().Str("context", "nostr-reject-unregistered").Logger()
+
+	// always allow auth messages
+	if event.Kind == 22242 {
+		return false, ""
+	}
+
+	// ensure pubkey has authenticated
+	authenticatedUser := khatru.GetAuthed(ctx)
+	if authenticatedUser == "" {
+		l.Debug().Msgf("pubkey not authed: %s", event.PubKey)
+		return true, fmt.Sprintf("auth-required: interacting with this relay requires authentication")
+	}
+
+	npubs := []string{authenticatedUser}
+	// add recipients to npubs list
+	if event.Kind == 4 || event.Kind == 14 {
+		for _, npub := range event.Tags.GetAll([]string{"p"}) {
+			npubs = append(npubs, npub.Value())
+		}
+	}
+
+	// check if npubs are registered
+	if authz := checknPubsInDb(npubs); authz == false {
+		l.Debug().Msgf("kind: %v, pubkey: %s, error: %s", event.Kind, event.PubKey, err.Error())
+		return true, fmt.Sprintf("restricted: pubkey %s is not registered to any users", authenticatedUser)
+	}
+	return false, ""
+}
diff --git a/nostr/relay.go b/nostr/relay.go
index 1c6019f..3c67fad 100644
--- a/nostr/relay.go
+++ b/nostr/relay.go
@@ -2,15 +2,12 @@ package nostr
 
 import (
 	"context"
-	"fmt"
 
-	"git.devvul.com/asara/gologger"
 	"git.devvul.com/asara/well-goknown/config"
 	"github.com/fiatjaf/eventstore/postgresql"
 	"github.com/fiatjaf/khatru"
 	"github.com/fiatjaf/khatru/policies"
 	"github.com/jmoiron/sqlx"
-	"github.com/nbd-wtf/go-nostr"
 )
 
 var (
@@ -73,49 +70,3 @@ func NewRelay(version string) *khatru.Relay {
 	)
 	return relay
 }
-
-func RejectUnregisteredNpubs(ctx context.Context, event *nostr.Event) (reject bool, msg string) {
-	l := gologger.Get(config.GetConfig().LogLevel).With().Str("context", "nostr-reject-unregistered").Logger()
-
-	// always allow auth messages
-	if event.Kind == 22242 {
-		return false, ""
-	}
-
-	authenticatedUser := khatru.GetAuthed(ctx)
-	if authenticatedUser == "" {
-		l.Debug().Msgf("pubkey not authed: %s", event.PubKey)
-		return true, fmt.Sprintf("auth-required: interacting with this relay requires authentication")
-	}
-
-	// reject nip-04 messages to users who aren't registered
-	if event.Kind == 4 {
-		receiver := event.Tags.GetFirst([]string{"p"}).Value()
-		var rid int
-		err := DB.QueryRow("SELECT id FROM users WHERE pubkey=$1", receiver).Scan(&rid)
-		if err != nil {
-			rid = -1
-		}
-
-		var sid int
-		err = DB.QueryRow("SELECT id FROM users WHERE pubkey=$1", authenticatedUser).Scan(&sid)
-		if err != nil {
-			sid = -1
-		}
-
-		if rid != -1 && sid != -1 {
-			l.Debug().Msgf("pubkeys %s or %s not found to be registered", receiver, event.PubKey)
-			return true, fmt.Sprintf("restricted: nobody in this nip04 message is registered to the relay")
-		}
-		return false, ""
-	}
-
-	// check if user is registered
-	var uid int
-	err := DB.QueryRow("SELECT id FROM users WHERE pubkey=$1", authenticatedUser).Scan(&uid)
-	if err != nil {
-		l.Debug().Msgf("kind: %v, pubkey: %s, error: %s", event.Kind, event.PubKey, err.Error())
-		return true, fmt.Sprintf("restricted: pubkey %s is not registered to any users", authenticatedUser)
-	}
-	return false, ""
-}