From 22da06b629ffc0fcaba48c0e32e2175ade0c5f30 Mon Sep 17 00:00:00 2001
From: fiatjaf <fiatjaf@gmail.com>
Date: Wed, 6 Dec 2023 12:14:58 -0300
Subject: [PATCH] new flow for auth based on "auth-required: " rejection
 messages.

---
 examples/readme-demo/main.go | 15 ++++++++++-----
 handlers.go                  |  6 ++++++
 helpers.go                   | 31 +++++++++++++++++++++++++++++++
 relay.go                     |  5 -----
 utils.go                     | 21 ---------------------
 5 files changed, 47 insertions(+), 31 deletions(-)
 create mode 100644 helpers.go

diff --git a/examples/readme-demo/main.go b/examples/readme-demo/main.go
index 938a8e3..65ba940 100644
--- a/examples/readme-demo/main.go
+++ b/examples/readme-demo/main.go
@@ -60,15 +60,20 @@ func main() {
 			return false, "" // anyone else can
 		},
 	)
-	relay.OnConnect = append(relay.OnConnect,
-		func(ctx context.Context) {
-			// request NIP-42 AUTH from everybody
-			relay.RequestAuth(ctx)
+
+	// you can request auth by rejecting an event or a request with the prefix "auth-required: "
+	relay.RejectFilter = append(relay.RejectFilter,
+		func(ctx context.Context, filter nostr.Filter) (reject bool, msg string) {
+			if pubkey := khatru.GetAuthed(ctx); pubkey != "" {
+				log.Printf("request from %s\n", pubkey)
+				return false, ""
+			}
+			return true, "auth-required: only authenticated users can read from this relay"
 		},
 	)
 	relay.OnAuth = append(relay.OnAuth,
 		func(ctx context.Context, pubkey string) {
-			// and when they auth we just log that for nothing
+			// and when they auth we can just log that for nothing
 			log.Println(pubkey + " is authed!")
 		},
 	)
diff --git a/handlers.go b/handlers.go
index a60a320..0f9f69b 100644
--- a/handlers.go
+++ b/handlers.go
@@ -134,6 +134,9 @@ func (rl *Relay) HandleWebsocket(w http.ResponseWriter, r *http.Request) {
 						ok = true
 					} else {
 						reason = nostr.NormalizeOKMessage(err.Error(), "blocked")
+						if isAuthRequired(reason) {
+							ws.WriteJSON(nostr.AuthEnvelope{Challenge: &ws.Challenge})
+						}
 					}
 					ws.WriteJSON(nostr.OKEnvelope{EventID: env.Event.ID, OK: ok, Reason: reason})
 				case *nostr.CountEnvelope:
@@ -154,6 +157,9 @@ func (rl *Relay) HandleWebsocket(w http.ResponseWriter, r *http.Request) {
 						err := rl.handleRequest(ctx, env.SubscriptionID, &eose, ws, filter)
 						if err == nil {
 							reason := nostr.NormalizeOKMessage(err.Error(), "blocked")
+							if isAuthRequired(reason) {
+								ws.WriteJSON(nostr.AuthEnvelope{Challenge: &ws.Challenge})
+							}
 							ws.WriteJSON(nostr.ClosedEnvelope{SubscriptionID: env.SubscriptionID, Reason: reason})
 							return
 						}
diff --git a/helpers.go b/helpers.go
new file mode 100644
index 0000000..dd2d858
--- /dev/null
+++ b/helpers.go
@@ -0,0 +1,31 @@
+package khatru
+
+import (
+	"hash/maphash"
+	"regexp"
+	"strings"
+	"unsafe"
+
+	"github.com/nbd-wtf/go-nostr"
+)
+
+const (
+	AUTH_CONTEXT_KEY = iota
+	WS_KEY
+)
+
+var nip20prefixmatcher = regexp.MustCompile(`^\w+: `)
+
+func pointerHasher[V any](_ maphash.Seed, k *V) uint64 {
+	return uint64(uintptr(unsafe.Pointer(k)))
+}
+
+func isOlder(previous, next *nostr.Event) bool {
+	return previous.CreatedAt < next.CreatedAt ||
+		(previous.CreatedAt == next.CreatedAt && previous.ID > next.ID)
+}
+
+func isAuthRequired(msg string) bool {
+	idx := strings.IndexByte(msg, ':')
+	return msg[0:idx] == "auth-required"
+}
diff --git a/relay.go b/relay.go
index 97f4ed1..bfe840c 100644
--- a/relay.go
+++ b/relay.go
@@ -82,8 +82,3 @@ type Relay struct {
 	PingPeriod     time.Duration // Send pings to peer with this period. Must be less than pongWait.
 	MaxMessageSize int64         // Maximum message size allowed from peer.
 }
-
-func (rl *Relay) RequestAuth(ctx context.Context) {
-	ws := GetConnection(ctx)
-	ws.WriteJSON(nostr.AuthEnvelope{Challenge: &ws.Challenge})
-}
diff --git a/utils.go b/utils.go
index 20e1e0c..f3960d3 100644
--- a/utils.go
+++ b/utils.go
@@ -2,20 +2,8 @@ package khatru
 
 import (
 	"context"
-	"hash/maphash"
-	"regexp"
-	"unsafe"
-
-	"github.com/nbd-wtf/go-nostr"
 )
 
-const (
-	AUTH_CONTEXT_KEY = iota
-	WS_KEY           = iota
-)
-
-var nip20prefixmatcher = regexp.MustCompile(`^\w+: `)
-
 func GetConnection(ctx context.Context) *WebSocket {
 	return ctx.Value(WS_KEY).(*WebSocket)
 }
@@ -27,12 +15,3 @@ func GetAuthed(ctx context.Context) string {
 	}
 	return authedPubkey.(string)
 }
-
-func pointerHasher[V any](_ maphash.Seed, k *V) uint64 {
-	return uint64(uintptr(unsafe.Pointer(k)))
-}
-
-func isOlder(previous, next *nostr.Event) bool {
-	return previous.CreatedAt < next.CreatedAt ||
-		(previous.CreatedAt == next.CreatedAt && previous.ID > next.ID)
-}