Change network to accept DialFunc

mtglib/network/init.go

@@ -11,7 +11,6 @@ import (
 const (
 	DefaultTimeout     = 10 * time.Second
 	DefaultIdleTimeout = time.Minute
-	DefaultHTTPTimeout = 5 * time.Second
 	DefaultBufferSize  = 4096
 
 	ProxyDialerOpenThreshold        = 5
@@ -20,7 +19,8 @@ const (
 
 	DefaultDOHHostname = "9.9.9.9"
 
-	DNSTimeout = 5 * time.Second
+	DNSTimeout  = 5 * time.Second
+	HTTPTimeout = 10 * time.Second
 )
 
 var (
@@ -28,6 +28,8 @@ var (
 	ErrCannotDialWithAllProxies = errors.New("cannot dial with all proxies")
 )
 
+type DialFunc func(ctx context.Context, protocol, address string) (net.Conn, error)
+
 type Dialer interface {
 	Dial(network, address string) (net.Conn, error)
 	DialContext(ctx context.Context, network, address string) (net.Conn, error)
@@ -37,7 +39,6 @@ type Network interface {
 	Dialer
 
 	DNSResolve(network, hostname string) (ips []string, err error)
-	MakeHTTPClient(timeout time.Duration) *http.Client
+	MakeHTTPClient(DialFunc) *http.Client
 	IdleTimeout() time.Duration
-	PrepareHTTPClient(*http.Client)
 }

mtglib/network/network.go

@@ -12,10 +12,22 @@ import (
 	doh "github.com/babolivier/go-doh-client"
 )
 
+type networkHTTPTransport struct {
+	userAgent string
+	next      http.RoundTripper
+}
+
+func (n networkHTTPTransport) RoundTrip(req *http.Request) (*http.Response, error) {
+	req.Header.Set("User-Agent", n.userAgent)
+
+	return n.next.RoundTrip(req)
+}
+
 type network struct {
-	idleTimeout time.Duration
 	dialer      Dialer
 	dns         doh.Resolver
+	idleTimeout time.Duration
+	userAgent   string
 }
 
 func (n *network) Dial(protocol, address string) (net.Conn, error) {
@@ -107,23 +119,15 @@ func (n *network) IdleTimeout() time.Duration {
 	return n.idleTimeout
 }
 
-func (n *network) MakeHTTPClient(timeout time.Duration) *http.Client {
-	if timeout <= 0 {
-		timeout = DefaultHTTPTimeout
+func (n *network) MakeHTTPClient(dialFunc DialFunc) *http.Client {
+	if dialFunc == nil {
+		dialFunc = n.DialContext
 	}
 
-	return &http.Client{
-		Timeout: timeout,
-		Transport: &http.Transport{
-			DialContext: n.DialContext,
-		},
-	}
+	return makeHTTPClient(n.userAgent, dialFunc)
 }
 
-func (n *network) PatchHTTPClient(_ *http.Client) {
-}
-
-func NewNetwork(dialer Dialer, dohHostname string, idleTimeout time.Duration) (Network, error) {
+func NewNetwork(dialer Dialer, userAgent, dohHostname string, idleTimeout time.Duration) (Network, error) {
 	switch {
 	case idleTimeout < 0:
 		return nil, fmt.Errorf("timeout should be positive number %s", idleTimeout)
@@ -138,6 +142,7 @@ func NewNetwork(dialer Dialer, dohHostname string, idleTimeout time.Duration) (N
 	return &network{
 		dialer:      dialer,
 		idleTimeout: idleTimeout,
+		userAgent:   userAgent,
 		dns: doh.Resolver{
 			Host:  dohHostname,
 			Class: doh.IN,
@@ -150,3 +155,15 @@ func NewNetwork(dialer Dialer, dohHostname string, idleTimeout time.Duration) (N
 		},
 	}, nil
 }
+
+func makeHTTPClient(userAgent string, dialFunc DialFunc) *http.Client {
+	return &http.Client{
+		Timeout: HTTPTimeout,
+		Transport: networkHTTPTransport{
+			userAgent: userAgent,
+			next: &http.Transport{
+				DialContext: dialFunc,
+			},
+		},
+	}
+}