From ec1721cfe33a1572d30bbcd9e87381980a6ce4e7 Mon Sep 17 00:00:00 2001
From: fiatjaf <fiatjaf@gmail.com>
Date: Wed, 25 Mar 2026 15:59:04 -0300
Subject: [PATCH] count: support CLOSED and AUTH.

fixes https://github.com/fiatjaf/nak/issues/117
---
 count.go   | 46 +++++++++++++++++++++++++++++++++++++++++++---
 helpers.go |  4 ++--
 2 files changed, 45 insertions(+), 5 deletions(-)

diff --git a/count.go b/count.go
index 493bcce..ff14209 100644
--- a/count.go
+++ b/count.go
@@ -7,8 +7,10 @@ import (
 	"strings"
 
 	"fiatjaf.com/nostr"
+	"fiatjaf.com/nostr/nip42"
 	"fiatjaf.com/nostr/nip45"
 	"fiatjaf.com/nostr/nip45/hyperloglog"
+	"github.com/fatih/color"
 	"github.com/mailru/easyjson"
 	"github.com/urfave/cli/v3"
 )
@@ -18,13 +20,51 @@ var count = &cli.Command{
 	Usage:                     "generates encoded COUNT messages and optionally use them to talk to relays",
 	Description:               `like 'nak req', but does a "COUNT" call instead. Will attempt to perform HyperLogLog aggregation if more than one relay is specified.`,
 	DisableSliceFlagSeparator: true,
-	Flags:                     reqFilterFlags,
-	ArgsUsage:                 "[relay...]",
+	Flags: append(defaultKeyFlags,
+		append(reqFilterFlags,
+			&cli.BoolFlag{
+				Name:  "auth",
+				Usage: "always perform nip42 \"AUTH\" when facing an \"auth-required: \" rejection and try again",
+			},
+			&cli.BoolFlag{
+				Name:     "force-pre-auth",
+				Aliases:  []string{"fpa"},
+				Usage:    "after connecting, for a nip42 \"AUTH\" message to be received, act on it and only then send the \"COUNT\"",
+				Category: CATEGORY_SIGNER,
+			},
+		)...,
+	),
+	ArgsUsage: "[relay...]",
 	Action: func(ctx context.Context, c *cli.Command) error {
 		biggerUrlSize := 0
 		relayUrls := c.Args().Slice()
 		if len(relayUrls) > 0 {
-			relays := connectToAllRelays(ctx, c, relayUrls, nil, nostr.PoolOptions{})
+			forcePreAuthSigner := authSigner
+			if !c.Bool("force-pre-auth") {
+				forcePreAuthSigner = nil
+			}
+			relays := connectToAllRelays(
+				ctx,
+				c,
+				relayUrls,
+				forcePreAuthSigner,
+				nostr.PoolOptions{
+					RelayOptions: nostr.RelayOptions{
+						AuthHandler: func(ctx context.Context, _ *nostr.Relay, authEvent *nostr.Event) error {
+							return authSigner(ctx, c, func(s string, args ...any) {
+								if strings.HasPrefix(s, "authenticating as") {
+									cleanUrl, _ := strings.CutPrefix(
+										nip42.GetRelayURLFromAuthEvent(*authEvent),
+										"wss://",
+									)
+									s = "authenticating to " + color.CyanString(cleanUrl) + " as" + s[len("authenticating as"):]
+								}
+								log(s+"\n", args...)
+							}, authEvent)
+						},
+					},
+				},
+			)
 			if len(relays) == 0 {
 				log("failed to connect to any of the given relays.\n")
 				os.Exit(3)
diff --git a/helpers.go b/helpers.go
index 7051abe..866d015 100644
--- a/helpers.go
+++ b/helpers.go
@@ -200,8 +200,8 @@ func connectToAllRelays(
 
 	opts.EventMiddleware = sys.TrackEventHints
 	opts.PenaltyBox = true
-	opts.RelayOptions = nostr.RelayOptions{
-		RequestHeader: http.Header{textproto.CanonicalMIMEHeaderKey("user-agent"): {"nak/s"}},
+	opts.RelayOptions.RequestHeader = http.Header{
+		textproto.CanonicalMIMEHeaderKey("user-agent"): {"nak/s"},
 	}
 	sys.Pool = nostr.NewPool(opts)