From 63a7059c98b513ed67d5056e569e2252f0670a40 Mon Sep 17 00:00:00 2001
From: phuslu <phus.lu@gmail.com>
Date: Sun, 27 Oct 2024 21:43:18 +0800
Subject: [PATCH] fix header for http dialer

---
 README.md              |  6 ++++--
 client_dialer.go       | 15 +++++++++------
 client_test.go         | 10 ++++++----
 cmd/fastdig/fastdig.go |  8 ++++++--
 4 files changed, 25 insertions(+), 14 deletions(-)

diff --git a/README.md b/README.md
index 27891be..6f19964 100644
--- a/README.md
+++ b/README.md
@@ -113,7 +113,10 @@ func main() {
 		Addr: endpoint.String(),
 		Dialer: &fastdns.HTTPDialer{
 			Endpoint:  endpoint,
-			UserAgent: "fastdns/0.9",
+			Header: http.Header{
+				"content-type": {"application/dns-message"},
+				"user-agent":   {"fastdns/1.0"},
+			},
 			Transport: &http.Transport{
 				ForceAttemptHTTP2:   true,
 				MaxIdleConns:        100,
@@ -125,7 +128,6 @@ func main() {
 					ServerName:         endpoint.Hostname(),
 					ClientSessionCache: tls.NewLRUClientSessionCache(1024),
 				},
-				ExpectContinueTimeout: 1 * time.Second,
 			},
 		},
 	}
diff --git a/client_dialer.go b/client_dialer.go
index 6f6b461..f214f29 100644
--- a/client_dialer.go
+++ b/client_dialer.go
@@ -160,13 +160,13 @@ type HTTPDialer struct {
 	// This is the base address used for sending HTTP requests.
 	Endpoint *url.URL
 
-	// UserAgent defines the User-Agent header that will be included in the HTTP requests
-	// sent by this dialer. It can be customized for specific needs.
-	UserAgent string
-
 	// Transport allows for customizing the underlying transport mechanism used
 	// for making HTTP requests. If set, it overrides the default RoundTripper behavior.
 	Transport http.RoundTripper
+
+	// Header defines the request header that will be sent in the HTTP requests.
+	// It can be customized for specific needs, E.g. User-Agent.
+	Header http.Header
 }
 
 func (d *HTTPDialer) DialContext(ctx context.Context, network, addr string) (net.Conn, error) {
@@ -176,6 +176,9 @@ func (d *HTTPDialer) DialContext(ctx context.Context, network, addr string) (net
 	c.req.Body = nil
 	c.req.URL = d.Endpoint
 	c.req.Host = d.Endpoint.Host
+	if d.Header != nil {
+		c.req.Header = d.Header
+	}
 	c.reader.B = nil
 	c.writer.B = c.writer.B[:0]
 	c.resp = nil
@@ -286,8 +289,8 @@ var httpconnpool = sync.Pool{
 			req: &http.Request{
 				Method: http.MethodPost,
 				Header: http.Header{
-					"content-type": []string{"application/dns-message"},
-					"user-agent":   []string{"fastdns/1.0"},
+					"content-type": {"application/dns-message"},
+					"user-agent":   {"fastdns/1.0"},
 				},
 			},
 			reader: new(bufferreader),
diff --git a/client_test.go b/client_test.go
index 0f2f8e7..6a45638 100644
--- a/client_test.go
+++ b/client_test.go
@@ -100,8 +100,7 @@ func TestClientLookup(t *testing.T) {
 		{
 			Addr: "https://1.1.1.1/dns-query",
 			Dialer: &HTTPDialer{
-				Endpoint:  func() (u *url.URL) { u, _ = url.Parse("https://1.1.1.1/dns-query"); return }(),
-				UserAgent: "fastdns/0.9",
+				Endpoint: func() (u *url.URL) { u, _ = url.Parse("https://1.1.1.1/dns-query"); return }(),
 			},
 		},
 	}
@@ -289,8 +288,11 @@ func BenchmarkResolverFastdnsHTTP(b *testing.B) {
 	resolver := &Client{
 		Addr: server,
 		Dialer: &HTTPDialer{
-			Endpoint:  func() (u *url.URL) { u, _ = url.Parse("https://" + server + "/dns-query"); return }(),
-			UserAgent: "fastdns/0.9",
+			Endpoint: func() (u *url.URL) { u, _ = url.Parse("https://" + server + "/dns-query"); return }(),
+			Header: http.Header{
+				"content-type": {"application/dns-message"},
+				"user-agent":   {"fastdns/1.0"},
+			},
 			Transport: &http.Transport{
 				ForceAttemptHTTP2:   true,
 				MaxIdleConns:        100,
diff --git a/cmd/fastdig/fastdig.go b/cmd/fastdig/fastdig.go
index b596301..7ce53be 100644
--- a/cmd/fastdig/fastdig.go
+++ b/cmd/fastdig/fastdig.go
@@ -6,6 +6,7 @@ import (
 	"encoding/binary"
 	"fmt"
 	"net"
+	"net/http"
 	"net/netip"
 	"net/url"
 	"os"
@@ -29,8 +30,11 @@ func main() {
 			os.Exit(1)
 		}
 		client.Dialer = &fastdns.HTTPDialer{
-			Endpoint:  endpoint,
-			UserAgent: "fastdig/0.9",
+			Endpoint: endpoint,
+			Header: http.Header{
+				"content-type": {"application/dns-message"},
+				"user-agent":   {"fastdns/1.0"},
+			},
 		}
 	}