From 0caf8deb5646af75547303092c547083ff231b22 Mon Sep 17 00:00:00 2001
From: fiatjaf <fiatjaf@gmail.com>
Date: Thu, 26 Sep 2024 19:18:34 -0300
Subject: [PATCH] WithUserAgent() pool option.

---
 pool.go  | 17 +++++++++++++++--
 relay.go |  4 +---
 2 files changed, 16 insertions(+), 5 deletions(-)

diff --git a/pool.go b/pool.go
index ff132eb..cf7fb08 100644
--- a/pool.go
+++ b/pool.go
@@ -30,6 +30,7 @@ type SimplePool struct {
 	signatureChecker func(Event) bool
 	penaltyBoxMu     sync.Mutex
 	penaltyBox       map[string][2]float64
+	userAgent        string
 }
 
 type DirectedFilters struct {
@@ -120,10 +121,20 @@ func (h WithEventMiddleware) ApplyPoolOption(pool *SimplePool) {
 	pool.eventMiddleware = append(pool.eventMiddleware, h)
 }
 
+// WithUserAgent sets the user-agent header for all relay connections in the pool.
+func WithUserAgent(userAgent string) withUserAgentOpt { return withUserAgentOpt(userAgent) }
+
+type withUserAgentOpt string
+
+func (h withUserAgentOpt) ApplyPoolOption(pool *SimplePool) {
+	pool.userAgent = string(h)
+}
+
 var (
 	_ PoolOption = (WithAuthHandler)(nil)
 	_ PoolOption = (WithEventMiddleware)(nil)
 	_ PoolOption = WithPenaltyBox()
+	_ PoolOption = WithUserAgent("")
 )
 
 func (pool *SimplePool) EnsureRelay(url string) (*Relay, error) {
@@ -146,7 +157,6 @@ func (pool *SimplePool) EnsureRelay(url string) (*Relay, error) {
 	}
 
 	// try to connect
-	var err error
 	// we use this ctx here so when the pool dies everything dies
 	ctx, cancel := context.WithTimeout(pool.Context, time.Second*15)
 	defer cancel()
@@ -156,7 +166,10 @@ func (pool *SimplePool) EnsureRelay(url string) (*Relay, error) {
 		opts = append(opts, WithSignatureChecker(pool.signatureChecker))
 	}
 
-	if relay, err = RelayConnect(ctx, nm, opts...); err != nil {
+	relay = NewRelay(context.Background(), url, opts...)
+	relay.RequestHeader.Set("User-Agent", pool.userAgent)
+
+	if err := relay.Connect(ctx); err != nil {
 		if pool.penaltyBox != nil {
 			// putting relay in penalty box
 			pool.penaltyBoxMu.Lock()
diff --git a/relay.go b/relay.go
index 1ffb1fe..6d5a9e4 100644
--- a/relay.go
+++ b/relay.go
@@ -66,6 +66,7 @@ func NewRelay(ctx context.Context, url string, opts ...RelayOption) *Relay {
 			ok, _ := e.CheckSignature()
 			return ok
 		},
+		RequestHeader: make(http.Header, 1),
 	}
 
 	for _, opt := range opts {
@@ -160,9 +161,6 @@ func (r *Relay) ConnectWithTLS(ctx context.Context, tlsConfig *tls.Config) error
 		defer cancel()
 	}
 
-	if r.RequestHeader == nil {
-		r.RequestHeader = make(http.Header, 1)
-	}
 	if r.RequestHeader.Get("User-Agent") == "" {
 		r.RequestHeader.Set("User-Agent", "github.com/nbd-wtf/go-nostr")
 	}