From 941eae5f615aeaf038f62002e673a7bf4886f1c7 Mon Sep 17 00:00:00 2001
From: Emily Lange <git@indeednotjames.com>
Date: Mon, 27 Feb 2023 18:23:09 +0100
Subject: reverseproxy: allow specifying ip version for dynamic `a` upstream
 (#5401)

Co-authored-by: Francis Lavoie <lavofr@gmail.com>
---
 .../reverse_proxy_dynamic_upstreams.txt            |  8 +++--
 modules/caddyhttp/reverseproxy/caddyfile.go        | 25 ++++++++++++++-
 modules/caddyhttp/reverseproxy/upstreams.go        | 36 ++++++++++++++++++++--
 3 files changed, 64 insertions(+), 5 deletions(-)

diff --git a/caddytest/integration/caddyfile_adapt/reverse_proxy_dynamic_upstreams.txt b/caddytest/integration/caddyfile_adapt/reverse_proxy_dynamic_upstreams.txt
index 2f2cbcd..384cc05 100644
--- a/caddytest/integration/caddyfile_adapt/reverse_proxy_dynamic_upstreams.txt
+++ b/caddytest/integration/caddyfile_adapt/reverse_proxy_dynamic_upstreams.txt
@@ -11,6 +11,7 @@
 			resolvers 8.8.8.8 8.8.4.4
 			dial_timeout 2s
 			dial_fallback_delay 300ms
+			versions ipv6
 		}
 	}
 }
@@ -66,7 +67,10 @@
 												"8.8.4.4"
 											]
 										},
-										"source": "a"
+										"source": "a",
+										"versions": {
+											"ipv6": true
+										}
 									},
 									"handler": "reverse_proxy"
 								}
@@ -113,4 +117,4 @@
 			}
 		}
 	}
-}
\ No newline at end of file
+}
diff --git a/modules/caddyhttp/reverseproxy/caddyfile.go b/modules/caddyhttp/reverseproxy/caddyfile.go
index 1211188..cf84ba7 100644
--- a/modules/caddyhttp/reverseproxy/caddyfile.go
+++ b/modules/caddyhttp/reverseproxy/caddyfile.go
@@ -1324,6 +1324,7 @@ func (u *SRVUpstreams) UnmarshalCaddyfile(d *caddyfile.Dispenser) error {
 //	    resolvers           <resolvers...>
 //	    dial_timeout        <timeout>
 //	    dial_fallback_delay <timeout>
+//	    versions            ipv4|ipv6
 //	}
 func (u *AUpstreams) UnmarshalCaddyfile(d *caddyfile.Dispenser) error {
 	for d.Next() {
@@ -1397,8 +1398,30 @@ func (u *AUpstreams) UnmarshalCaddyfile(d *caddyfile.Dispenser) error {
 				}
 				u.FallbackDelay = caddy.Duration(dur)
 
+			case "versions":
+				args := d.RemainingArgs()
+				if len(args) == 0 {
+					return d.Errf("must specify at least one version")
+				}
+
+				if u.Versions == nil {
+					u.Versions = &ipVersions{}
+				}
+
+				trueBool := true
+				for _, arg := range args {
+					switch arg {
+					case "ipv4":
+						u.Versions.IPv4 = &trueBool
+					case "ipv6":
+						u.Versions.IPv6 = &trueBool
+					default:
+						return d.Errf("unsupported version: '%s'", arg)
+					}
+				}
+
 			default:
-				return d.Errf("unrecognized srv option '%s'", d.Val())
+				return d.Errf("unrecognized a option '%s'", d.Val())
 			}
 		}
 	}
diff --git a/modules/caddyhttp/reverseproxy/upstreams.go b/modules/caddyhttp/reverseproxy/upstreams.go
index 7a90016..30bd7b5 100644
--- a/modules/caddyhttp/reverseproxy/upstreams.go
+++ b/modules/caddyhttp/reverseproxy/upstreams.go
@@ -213,6 +213,11 @@ func (sl srvLookup) isFresh() bool {
 	return time.Since(sl.freshness) < time.Duration(sl.srvUpstreams.Refresh)
 }
 
+type ipVersions struct {
+	IPv4 *bool `json:"ipv4,omitempty"`
+	IPv6 *bool `json:"ipv6,omitempty"`
+}
+
 // AUpstreams provides upstreams from A/AAAA lookups.
 // Results are cached and refreshed at the configured
 // refresh interval.
@@ -240,6 +245,11 @@ type AUpstreams struct {
 	// A negative value disables this.
 	FallbackDelay caddy.Duration `json:"dial_fallback_delay,omitempty"`
 
+	// The IP versions to resolve for. By default, both
+	// "ipv4" and "ipv6" will be enabled, which
+	// correspond to A and AAAA records respectively.
+	Versions *ipVersions `json:"versions,omitempty"`
+
 	resolver *net.Resolver
 }
 
@@ -286,7 +296,29 @@ func (au *AUpstreams) Provision(_ caddy.Context) error {
 
 func (au AUpstreams) GetUpstreams(r *http.Request) ([]*Upstream, error) {
 	repl := r.Context().Value(caddy.ReplacerCtxKey).(*caddy.Replacer)
-	auStr := repl.ReplaceAll(au.String(), "")
+
+	resolveIpv4 := au.Versions.IPv4 == nil || *au.Versions.IPv4
+	resolveIpv6 := au.Versions.IPv6 == nil || *au.Versions.IPv6
+
+	// Map ipVersion early, so we can use it as part of the cache-key.
+	// This should be fairly inexpensive and comes and the upside of
+	// allowing the same dynamic upstream (name + port combination)
+	// to be used multiple times with different ip versions.
+	//
+	// It also forced a cache-miss if a previously cached dynamic
+	// upstream changes its ip version, e.g. after a config reload,
+	// while keeping the cache-invalidation as simple as it currently is.
+	var ipVersion string
+	switch {
+	case resolveIpv4 && !resolveIpv6:
+		ipVersion = "ip4"
+	case !resolveIpv4 && resolveIpv6:
+		ipVersion = "ip6"
+	default:
+		ipVersion = "ip"
+	}
+
+	auStr := repl.ReplaceAll(au.String()+ipVersion, "")
 
 	// first, use a cheap read-lock to return a cached result quickly
 	aAaaaMu.RLock()
@@ -311,7 +343,7 @@ func (au AUpstreams) GetUpstreams(r *http.Request) ([]*Upstream, error) {
 	name := repl.ReplaceAll(au.Name, "")
 	port := repl.ReplaceAll(au.Port, "")
 
-	ips, err := au.resolver.LookupIPAddr(r.Context(), name)
+	ips, err := au.resolver.LookupIP(r.Context(), ipVersion, name)
 	if err != nil {
 		return nil, err
 	}
-- 
cgit v1.2.3