feat(courierhttp/client): added support for custom host aliases and set default tls config

morlay · morlay · commit 8e79357f00d0 · 2025-07-09T16:45:59.000+08:00
diff --git a/pkg/courierhttp/client/host_alias.go b/pkg/courierhttp/client/host_alias.go
@@ -0,0 +1,155 @@
+package client
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"iter"
+	"maps"
+	"net"
+	"strings"
+
+	randv2 "math/rand/v2"
+)
+
+var ErrInvalidHostAlias = errors.New("invalid host alias")
+
+func ParseHostAlias(x string) (*HostAlias, error) {
+	if x == "" {
+		return nil, fmt.Errorf("empty host alias: %w", ErrInvalidHostAlias)
+	}
+
+	ha := &HostAlias{}
+
+	if strings.IndexByte(x, '[') == 0 {
+		// ipv6
+		end := strings.IndexByte(x, ']')
+		if end == -1 {
+			return nil, fmt.Errorf("ipv6 addr should warpped []: %w", ErrInvalidHostAlias)
+		}
+
+		ha.IP = net.ParseIP(x[1:end])
+
+		x = x[end:]
+
+		end = strings.IndexByte(x, ':')
+		if end == -1 {
+			return nil, fmt.Errorf(" ip should end with ':': %w", ErrInvalidHostAlias)
+		}
+		x = x[end+1:]
+	} else {
+		end := strings.IndexByte(x, ':')
+		if end == -1 {
+			return nil, fmt.Errorf(" ip should end with ':': %w", ErrInvalidHostAlias)
+		}
+		ha.IP = net.ParseIP(x[0:end])
+		x = x[end+1:]
+	}
+
+	if x == "" {
+		return nil, fmt.Errorf("invalid host alias")
+	}
+
+	ha.Hostnames = strings.Split(x, ",")
+
+	return ha, nil
+}
+
+type HostAlias struct {
+	IP        net.IP
+	Hostnames []string
+}
+
+func (x HostAlias) IsZero() bool {
+	return len(x.Hostnames) == 0 || len(x.IP) == 0
+}
+
+func (x *HostAlias) UnmarshalText(raw []byte) error {
+	ha, err := ParseHostAlias(string(raw))
+	if err != nil {
+		return err
+	}
+	*x = *ha
+	return nil
+}
+
+func (x HostAlias) MarshalText() ([]byte, error) {
+	return []byte(x.String()), nil
+}
+
+func (x HostAlias) String() string {
+	s := strings.Builder{}
+
+	if ip := x.IP.To4(); ip != nil {
+		s.WriteString(ip.String())
+	} else {
+		s.WriteString("[")
+		s.WriteString(x.IP.String())
+		s.WriteString("]")
+	}
+
+	s.WriteString(":")
+
+	for i, hostname := range x.Hostnames {
+		if i > 0 {
+			s.WriteString(",")
+		}
+		s.WriteString(hostname)
+	}
+
+	return s.String()
+}
+
+type Hosts map[string]map[string]struct{}
+
+func (hosts Hosts) WrapDialContext(dialContext func(ctx context.Context, network string, address string) (net.Conn, error)) func(ctx context.Context, network string, addr string) (net.Conn, error) {
+	return func(ctx context.Context, network string, addr string) (net.Conn, error) {
+		if len(hosts) == 0 {
+			return dialContext(ctx, network, addr)
+		}
+
+		host, port, err := net.SplitHostPort(addr)
+		if err != nil {
+			host = addr
+			port = "80"
+		}
+
+		if ips, ok := hosts[host]; ok && len(ips) > 0 {
+			resolved := net.JoinHostPort(hosts.selectIP(maps.Keys(ips), len(ips)), port)
+			return dialContext(ctx, network, resolved)
+		}
+
+		return dialContext(ctx, network, addr)
+	}
+}
+
+func (hosts Hosts) selectIP(ips iter.Seq[string], n int) (ip string) {
+	i := 0
+	idx := 0
+	if n > 1 {
+		idx = randv2.IntN(n) - 1
+	}
+
+	for x := range ips {
+		if i == idx {
+			ip = x
+			break
+		}
+		i++
+	}
+
+	return
+}
+
+func (hosts Hosts) AddHostAlias(alias HostAlias) {
+	if alias.IsZero() {
+		return
+	}
+
+	for _, hostname := range alias.Hostnames {
+		if hosts[hostname] == nil {
+			hosts[hostname] = make(map[string]struct{})
+		}
+		hosts[hostname][alias.IP.String()] = struct{}{}
+	}
+}
diff --git a/pkg/courierhttp/client/host_alias_test.go b/pkg/courierhttp/client/host_alias_test.go
@@ -0,0 +1,60 @@
+package client
+
+import (
+	"net"
+	"testing"
+
+	"github.com/octohelm/x/testing/bdd"
+)
+
+func TestHostAliases(t *testing.T) {
+	bdd.FromT(t).Given("host alias for ipv4", func(b bdd.T) {
+		ha := &HostAlias{
+			IP: net.ParseIP("127.0.0.1"),
+			Hostnames: []string{
+				"localhost",
+				"localhost1",
+			},
+		}
+
+		txt := bdd.Must(ha.MarshalText())
+
+		b.Then("match results",
+			bdd.Equal("127.0.0.1:localhost,localhost1", string(txt)),
+		)
+
+		b.When("unmarshal", func(b bdd.T) {
+			ha1 := &HostAlias{}
+
+			b.Then("success",
+				bdd.NoError(ha1.UnmarshalText(txt)),
+				bdd.Equal(ha, ha1),
+			)
+		})
+	})
+
+	bdd.FromT(t).Given("host alias for ipv6", func(b bdd.T) {
+		ha := &HostAlias{
+			IP: net.ParseIP("::1"),
+			Hostnames: []string{
+				"localhost",
+				"localhost1",
+			},
+		}
+
+		txt := bdd.Must(ha.MarshalText())
+
+		b.Then("match results",
+			bdd.Equal("[::1]:localhost,localhost1", string(txt)),
+		)
+
+		b.When("unmarshal", func(b bdd.T) {
+			ha1 := &HostAlias{}
+
+			b.Then("success",
+				bdd.NoError(ha1.UnmarshalText(txt)),
+				bdd.Equal(ha, ha1),
+			)
+		})
+	})
+}
diff --git a/pkg/courierhttp/client/http_default.go b/pkg/courierhttp/client/http_default.go
@@ -2,18 +2,19 @@ package client
 
 import (
 	"context"
+	"crypto/tls"
 	"net"
 	"net/http"
 	"time"
 )
 
-var reasonableRoundTripper http.RoundTripper = &http.Transport{
+var reasonableRoundTripper = &http.Transport{
 	Proxy: http.ProxyFromEnvironment,
 
-	DialContext: (&net.Dialer{
+	DialContext: defaultHosts.WrapDialContext((&net.Dialer{
 		Timeout:   30 * time.Second,
 		KeepAlive: 30 * time.Second,
-	}).DialContext,
+	}).DialContext),
 
 	MaxIdleConns:        100,
 	MaxIdleConnsPerHost: 10,
@@ -27,8 +28,25 @@ var reasonableRoundTripper http.RoundTripper = &http.Transport{
 	ForceAttemptHTTP2: true,
 }
 
+var defaultTlsConfig = &tls.Config{}
+
+var defaultHosts = Hosts{}
+
+func AddHostAlias(hostAliases ...HostAlias) {
+	for _, hostAlias := range hostAliases {
+		defaultHosts.AddHostAlias(hostAlias)
+	}
+}
+
+func SetDefaultTLSClientConfig(tlsConfig *tls.Config) {
+	if tlsConfig != nil {
+		defaultTlsConfig = tlsConfig.Clone()
+		reasonableRoundTripper.TLSClientConfig = tlsConfig.Clone()
+	}
+}
+
 func GetReasonableClientContext(ctx context.Context, httpTransports ...HttpTransport) *http.Client {
-	t := reasonableRoundTripper
+	t := http.RoundTripper(reasonableRoundTripper)
 
 	tc, ok := RoundTripperCreatorFromContext(ctx)
 	if ok {
@@ -39,20 +57,26 @@ func GetReasonableClientContext(ctx context.Context, httpTransports ...HttpTrans
 }
 
 func newRoundTripperWithoutKeepAlive() http.RoundTripper {
-	return &http.Transport{
+	t := &http.Transport{
 		Proxy: http.ProxyFromEnvironment,
 
-		DialContext: (&net.Dialer{
+		DialContext: defaultHosts.WrapDialContext((&net.Dialer{
 			Timeout:   30 * time.Second,
 			KeepAlive: 0,
-		}).DialContext,
+		}).DialContext),
 
 		DisableKeepAlives: true,
 
 		TLSHandshakeTimeout:   10 * time.Second,
 		ExpectContinueTimeout: 1 * time.Second,
 		ResponseHeaderTimeout: 60 * time.Second,
 	}
+
+	if defaultTlsConfig != nil {
+		t.TLSClientConfig = defaultTlsConfig.Clone()
+	}
+
+	return t
 }
 
 func GetShortConnClientContext(ctx context.Context, httpTransports ...HttpTransport) *http.Client {
