Repository: caddyserver/caddy (language: go)

Issue:
Question about load balancing
So what I am wanting is to load balance with least conn algorithm. But while doing so, I do not want my client to lose his session details as the flow of my app continues. So for that, is it right to use header up to set the session cookie? **Look at the example below**

flow:
client server request -> backend pick least conn server -> pass cookie header when it is present to uplink server stream -> server receives request

```
reverse_proxy :5000 :3000 {
    lb_policy least_conn
    header_up  Cookie {cookie.*}
}
``` 


Discussion / hints:
```reverse_proxy :5000 :3000 {
    lb_policy least_conn {
             header cookie
    }
}``` From the docs i see we can provide a header map to upstream specific headers, so this is the best for my purpose right?
No, you need to specifically use the `cookie` policy for it to be sticky. Only that policy considers a cookie value for determining the upstream to choose.

That `header_up` line is useless, all headers including cookies are sent upstream transparently. That line will probably break the cookies header, actually.

Something I've been considering is making the secondary policy configurable. Currently, most policies just do `random` as the fallback, but it would be useful to use a different policy instead when falling back.
> No, you need to specifically use the `cookie` policy for it to be sticky. Only that policy considers a cookie value for determining the upstream to choose.
> 
> That `header_up` line is useless, all headers including cookies are sent upstream transparently. That line will probably break the cookies header, actually.
> 
> Something I've been considering is making the secondary policy configurable. Currently, most policies just do `random` as the fallback, but it would be useful to use a different policy instead when falling back.

But the only supported HMAC hashing is 256 and not 512, my backend uses sha512 not 256 i read from the docs that it will do set cookie when theree was no cookie present.

The only thing i want is upstream a session when it is present not to set a session cookie when it was not present
Is the below config valid for that?


```
# tbh idk if the syntax is valid but i guess it is<name> [<options...>]
reverse_proxy :5000 :3000 {
    lb_policy least_conn {
          cookie sessionId "secret 1234" # uses sha256 hmac hashing to set cookie when no cookie is present but i wish sha512..
    }
}
```
> The only thing i want is upstream a session when it is present not to set a session cookie when it was not present

That's not a supported usecase currently. We could probably make that an option, i.e. `read_only` to avoid writing a cookie value.

> But the only supported HMAC hashing is 256 and not 512, my backend uses sha512 not 256 i read from the docs that it will do set cookie when theree was no cookie present.

Keep in mind that the contents of the hash need to be a precise value, i.e. the upstream address as known by Caddy. You can't use just any session cookie for this. So for that reason, the hashing algorithm Caddy uses doesn't need to match the algorithm your backend uses for sessions. Instead, your backend should use the same hashing algorithm as Caddy for the sticky cookie. To clarify, you would need to use a second cookie for this, your session cookie can't be reused.

> Is the below config valid for that?

No, fallback functionality doesn't exist yet.
To explain further, the reason your session cookie can't be reused is that this would mean Caddy would need to store in memory the list of all the sessions it has ever seen as a mapping to the upstream we proxied it to. That would use _way too much_ memory. The `cookie` policy uses no storage, because it delegates the storage to the cookie itself, so the client is the one that stores which upstream it had connected to.

I'm not convinced a `read_only` option makes sense, because it would mean your upstream needs to be aware of which upstream was chosen (so you'd need to either pass the upstream address as Caddy knows it as a request header upstream, or you'd need to hardcode or guess the upstream address in your backend by the machine it's running on). You should just let Caddy create the cookie itself.
> To explain further, the reason your session cookie can't be reused is that this would mean Caddy would need to store in memory the list of all the sessions it has ever seen as a mapping to the upstream we proxied it to. That would use _way too much_ memory. The `cookie` policy uses no storage, because it delegates the storage to the cookie itself, so the client is the one that stores which upstream it had connected to.
> 
> I'm not convinced a `read_only` option makes sense, because it would mean your upstream needs to be aware of which upstream was chosen (so you'd need to either pass the upstream address as Caddy knows it as a request header upstream, or you'd need to hardcode or guess the upstream address in your backend by the machine it's running on). You should just let Caddy create the cookie itself.

Can wee see in future that we can choose whether caddy uses sha256 or sha512 to compute the cookie hash, bcs atm only sha256 from what i see, and thnx for the help tbh 
I don't see any reason to make the hashing algorithm customizable. As I said, the cookie value only makes sense to be controlled by Caddy. But if you wanted to force a specific upstream... if you can compute a sha512 hash, you can also compute a sha256 hash.
Oke thnx

Fix patch:
diff --git a/modules/caddyhttp/reverseproxy/selectionpolicies.go b/modules/caddyhttp/reverseproxy/selectionpolicies.go
index 4184df59607..a2985f196fa 100644
--- a/modules/caddyhttp/reverseproxy/selectionpolicies.go
+++ b/modules/caddyhttp/reverseproxy/selectionpolicies.go
@@ -18,6 +18,7 @@ import (
 	"crypto/hmac"
 	"crypto/sha256"
 	"encoding/hex"
+	"encoding/json"
 	"fmt"
 	"hash/fnv"
 	weakrand "math/rand"
@@ -29,6 +30,7 @@ import (
 	"time"
 
 	"github.com/caddyserver/caddy/v2"
+	"github.com/caddyserver/caddy/v2/caddyconfig"
 	"github.com/caddyserver/caddy/v2/caddyconfig/caddyfile"
 	"github.com/caddyserver/caddy/v2/modules/caddyhttp"
 )
@@ -372,6 +374,10 @@ func (r *URIHashSelection) UnmarshalCaddyfile(d *caddyfile.Dispenser) error {
 type QueryHashSelection struct {
 	// The query key whose value is to be hashed and used for upstream selection.
 	Key string `json:"key,omitempty"`
+
+	// The fallback policy to use if the query key is not present. Defaults to `random`.
+	FallbackRaw json.RawMessage `json:"fallback,omitempty" caddy:"namespace=http.reverse_proxy.selection_policies inline_key=policy"`
+	fallback    Selector
 }
 
 // CaddyModule returns the Caddy module information.
@@ -382,12 +388,24 @@ func (QueryHashSelection) CaddyModule() caddy.ModuleInfo {
 	}
 }
 
-// Select returns an available host, if any.
-func (s QueryHashSelection) Select(pool UpstreamPool, req *http.Request, _ http.ResponseWriter) *Upstream {
+// Provision sets up the module.
+func (s *QueryHashSelection) Provision(ctx caddy.Context) error {
 	if s.Key == "" {
-		return nil
+		return fmt.Errorf("query key is required")
+	}
+	if s.FallbackRaw == nil {
+		s.FallbackRaw = caddyconfig.JSONModuleObject(RandomSelection{}, "policy", "random", nil)
 	}
+	mod, err := ctx.LoadModule(s, "FallbackRaw")
+	if err != nil {
+		return fmt.Errorf("loading fallback selection policy: %s", err)
+	}
+	s.fallback = mod.(Selector)
+	return nil
+}
 
+// Select returns an available host, if any.
+func (s QueryHashSelection) Select(pool UpstreamPool, req *http.Request, _ http.ResponseWriter) *Upstream {
 	// Since the query may have multiple values for the same key,
 	// we'll join them to avoid a problem where the user can control
 	// the upstream that the request goes to by sending multiple values
@@ -397,7 +415,7 @@ func (s QueryHashSelection) Select(pool UpstreamPool, req *http.Request, _ http.
 	// different request, because the order of the values is significant.
 	vals := strings.Join(req.URL.Query()[s.Key], ",")
 	if vals == "" {
-		return RandomSelection{}.Select(pool, req, nil)
+		return s.fallback.Select(pool, req, nil)
 	}
 	return hostByHashing(pool, vals)
 }
@@ -410,6 +428,24 @@ func (s *QueryHashSelection) UnmarshalCaddyfile(d *caddyfile.Dispenser) error {
 		}
 		s.Key = d.Val()
 	}
+	for nesting := d.Nesting(); d.NextBlock(nesting); {
+		switch d.Val() {
+		case "fallback":
+			if !d.NextArg() {
+				return d.ArgErr()
+			}
+			if s.FallbackRaw != nil {
+				return d.Err("fallback selection policy already specified")
+			}
+			mod, err := loadFallbackPolicy(d)
+			if err != nil {
+				return err
+			}
+			s.FallbackRaw = mod
+		default:
+			return d.Errf("unrecognized option '%s'", d.Val())
+		}
+	}
 	return nil
 }
 
@@ -418,6 +454,10 @@ func (s *QueryHashSelection) UnmarshalCaddyfile(d *caddyfile.Dispenser) error {
 type HeaderHashSelection struct {
 	// The HTTP header field whose value is to be hashed and used for upstream selection.
 	Field string `json:"field,omitempty"`
+
+	// The fallback policy to use if the header is not present. Defaults to `random`.
+	FallbackRaw json.RawMessage `json:"fallback,omitempty" caddy:"namespace=http.reverse_proxy.selection_policies inline_key=policy"`
+	fallback    Selector
 }
 
 // CaddyModule returns the Caddy module information.
@@ -428,12 +468,24 @@ func (HeaderHashSelection) CaddyModule() caddy.ModuleInfo {
 	}
 }
 
-// Select returns an available host, if any.
-func (s HeaderHashSelection) Select(pool UpstreamPool, req *http.Request, _ http.ResponseWriter) *Upstream {
+// Provision sets up the module.
+func (s *HeaderHashSelection) Provision(ctx caddy.Context) error {
 	if s.Field == "" {
-		return nil
+		return fmt.Errorf("header field is required")
+	}
+	if s.FallbackRaw == nil {
+		s.FallbackRaw = caddyconfig.JSONModuleObject(RandomSelection{}, "policy", "random", nil)
 	}
+	mod, err := ctx.LoadModule(s, "FallbackRaw")
+	if err != nil {
+		return fmt.Errorf("loading fallback selection policy: %s", err)
+	}
+	s.fallback = mod.(Selector)
+	return nil
+}
 
+// Select returns an available host, if any.
+func (s HeaderHashSelection) Select(pool UpstreamPool, req *http.Request, _ http.ResponseWriter) *Upstream {
 	// The Host header should be obtained from the req.Host field
 	// since net/http removes it from the header map.
 	if s.Field == "Host" && req.Host != "" {
@@ -442,7 +494,7 @@ func (s HeaderHashSelection) Select(pool UpstreamPool, req *http.Request, _ http
 
 	val := req.Header.Get(s.Field)
 	if val == "" {
-		return RandomSelection{}.Select(pool, req, nil)
+		return s.fallback.Select(pool, req, nil)
 	}
 	return hostByHashing(pool, val)
 }
@@ -455,6 +507,24 @@ func (s *HeaderHashSelection) UnmarshalCaddyfile(d *caddyfile.Dispenser) error {
 		}
 		s.Field = d.Val()
 	}
+	for nesting := d.Nesting(); d.NextBlock(nesting); {
+		switch d.Val() {
+		case "fallback":
+			if !d.NextArg() {
+				return d.ArgErr()
+			}
+			if s.FallbackRaw != nil {
+				return d.Err("fallback selection policy already specified")
+			}
+			mod, err := loadFallbackPolicy(d)
+			if err != nil {
+				return err
+			}
+			s.FallbackRaw = mod
+		default:
+			return d.Errf("unrecognized option '%s'", d.Val())
+		}
+	}
 	return nil
 }
 
@@ -465,6 +535,10 @@ type CookieHashSelection struct {
 	Name string `json:"name,omitempty"`
 	// Secret to hash (Hmac256) chosen upstream in cookie
 	Secret string `json:"secret,omitempty"`
+
+	// The fallback policy to use if the cookie is not present. Defaults to `random`.
+	FallbackRaw json.RawMessage `json:"fallback,omitempty" caddy:"namespace=http.reverse_proxy.selection_policies inline_key=policy"`
+	fallback    Selector
 }
 
 // CaddyModule returns the Caddy module information.
@@ -475,15 +549,48 @@ func (CookieHashSelection) CaddyModule() caddy.ModuleInfo {
 	}
 }
 
-// Select returns an available host, if any.
-func (s CookieHashSelection) Select(pool UpstreamPool, req *http.Request, w http.ResponseWriter) *Upstream {
+// Provision sets up the module.
+func (s *CookieHashSelection) Provision(ctx caddy.Context) error {
 	if s.Name == "" {
 		s.Name = "lb"
 	}
+	if s.FallbackRaw == nil {
+		s.FallbackRaw = caddyconfig.JSONModuleObject(RandomSelection{}, "policy", "random", nil)
+	}
+	mod, err := ctx.LoadModule(s, "FallbackRaw")
+	if err != nil {
+		return fmt.Errorf("loading fallback selection policy: %s", err)
+	}
+	s.fallback = mod.(Selector)
+	return nil
+}
+
+// Select returns an available host, if any.
+func (s CookieHashSelection) Select(pool UpstreamPool, req *http.Request, w http.ResponseWriter) *Upstream {
+	// selects a new Host using the fallback policy (typically random)
+	// and write a sticky session cookie to the response.
+	selectNewHost := func() *Upstream {
+		upstream := s.fallback.Select(pool, req, w)
+		if upstream == nil {
+			return nil
+		}
+		sha, err := hashCookie(s.Secret, upstream.Dial)
+		if err != nil {
+			return upstream
+		}
+		http.SetCookie(w, &http.Cookie{
+			Name:   s.Name,
+			Value:  sha,
+			Path:   "/",
+			Secure: false,
+		})
+		return upstream
+	}
+
 	cookie, err := req.Cookie(s.Name)
-	// If there's no cookie, select new random host
+	// If there's no cookie, select a host using the fallback policy
 	if err != nil || cookie == nil {
-		return selectNewHostWithCookieHashSelection(pool, w, s.Secret, s.Name)
+		return selectNewHost()
 	}
 	// If the cookie is present, loop over the available upstreams until we find a match
 	cookieValue := cookie.Value
@@ -496,13 +603,15 @@ func (s CookieHashSelection) Select(pool UpstreamPool, req *http.Request, w http
 			return upstream
 		}
 	}
-	// If there is no matching host, select new random host
-	return selectNewHostWithCookieHashSelection(pool, w, s.Secret, s.Name)
+	// If there is no matching host, select a host using the fallback policy
+	return selectNewHost()
 }
 
 // UnmarshalCaddyfile sets up the module from Caddyfile tokens. Syntax:
 //
-//	lb_policy cookie [<name> [<secret>]]
+//	lb_policy cookie [<name> [<secret>]] {
+//		fallback <policy>
+//	}
 //
 // By default name is `lb`
 func (s *CookieHashSelection) UnmarshalCaddyfile(d *caddyfile.Dispenser) error {
@@ -517,22 +626,25 @@ func (s *CookieHashSelection) UnmarshalCaddyfile(d *caddyfile.Dispenser) error {
 	default:
 		return d.ArgErr()
 	}
-	return nil
-}
-
-// Select a new Host randomly and add a sticky session cookie
-func selectNewHostWithCookieHashSelection(pool []*Upstream, w http.ResponseWriter, cookieSecret string, cookieName string) *Upstream {
-	randomHost := selectRandomHost(pool)
-
-	if randomHost != nil {
-		// Hash (HMAC with some key for privacy) the upstream.Dial string as the cookie value
-		sha, err := hashCookie(cookieSecret, randomHost.Dial)
-		if err == nil {
-			// write the cookie.
-			http.SetCookie(w, &http.Cookie{Name: cookieName, Value: sha, Path: "/", Secure: false})
+	for nesting := d.Nesting(); d.NextBlock(nesting); {
+		switch d.Val() {
+		case "fallback":
+			if !d.NextArg() {
+				return d.ArgErr()
+			}
+			if s.FallbackRaw != nil {
+				return d.Err("fallback selection policy already specified")
+			}
+			mod, err := loadFallbackPolicy(d)
+			if err != nil {
+				return err
+			}
+			s.FallbackRaw = mod
+		default:
+			return d.Errf("unrecognized option '%s'", d.Val())
 		}
 	}
-	return randomHost
+	return nil
 }
 
 // hashCookie hashes (HMAC 256) some data with the secret
@@ -627,6 +739,20 @@ func hash(s string) uint32 {
 	return h.Sum32()
 }
 
+func loadFallbackPolicy(d *caddyfile.Dispenser) (json.RawMessage, error) {
+	name := d.Val()
+	modID := "http.reverse_proxy.selection_policies." + name
+	unm, err := caddyfile.UnmarshalModule(d, modID)
+	if err != nil {
+		return nil, err
+	}
+	sel, ok := unm.(Selector)
+	if !ok {
+		return nil, d.Errf("module %s (%T) is not a reverseproxy.Selector", modID, unm)
+	}
+	return caddyconfig.JSONModuleObject(sel, "policy", name, nil), nil
+}
+
 // Interface guards
 var (
 	_ Selector = (*RandomSelection)(nil)


Test patch:
diff --git a/modules/caddyhttp/reverseproxy/selectionpolicies_test.go b/modules/caddyhttp/reverseproxy/selectionpolicies_test.go
index d2b7b3d0be7..93dcb77acd4 100644
--- a/modules/caddyhttp/reverseproxy/selectionpolicies_test.go
+++ b/modules/caddyhttp/reverseproxy/selectionpolicies_test.go
@@ -20,6 +20,8 @@ import (
 	"net/http/httptest"
 	"testing"
 
+	"github.com/caddyserver/caddy/v2"
+	"github.com/caddyserver/caddy/v2/caddyconfig"
 	"github.com/caddyserver/caddy/v2/modules/caddyhttp"
 )
 
@@ -33,7 +35,7 @@ func testPool() UpstreamPool {
 
 func TestRoundRobinPolicy(t *testing.T) {
 	pool := testPool()
-	rrPolicy := new(RoundRobinSelection)
+	rrPolicy := RoundRobinSelection{}
 	req, _ := http.NewRequest("GET", "/", nil)
 
 	h := rrPolicy.Select(pool, req, nil)
@@ -74,7 +76,7 @@ func TestRoundRobinPolicy(t *testing.T) {
 
 func TestLeastConnPolicy(t *testing.T) {
 	pool := testPool()
-	lcPolicy := new(LeastConnSelection)
+	lcPolicy := LeastConnSelection{}
 	req, _ := http.NewRequest("GET", "/", nil)
 
 	pool[0].countRequest(10)
@@ -92,7 +94,7 @@ func TestLeastConnPolicy(t *testing.T) {
 
 func TestIPHashPolicy(t *testing.T) {
 	pool := testPool()
-	ipHash := new(IPHashSelection)
+	ipHash := IPHashSelection{}
 	req, _ := http.NewRequest("GET", "/", nil)
 
 	// We should be able to predict where every request is routed.
@@ -234,7 +236,7 @@ func TestIPHashPolicy(t *testing.T) {
 
 func TestClientIPHashPolicy(t *testing.T) {
 	pool := testPool()
-	ipHash := new(ClientIPHashSelection)
+	ipHash := ClientIPHashSelection{}
 	req, _ := http.NewRequest("GET", "/", nil)
 	req = req.WithContext(context.WithValue(req.Context(), caddyhttp.VarsCtxKey, make(map[string]any)))
 
@@ -377,7 +379,7 @@ func TestClientIPHashPolicy(t *testing.T) {
 
 func TestFirstPolicy(t *testing.T) {
 	pool := testPool()
-	firstPolicy := new(FirstSelection)
+	firstPolicy := FirstSelection{}
 	req := httptest.NewRequest(http.MethodGet, "/", nil)
 
 	h := firstPolicy.Select(pool, req, nil)
@@ -393,8 +395,15 @@ func TestFirstPolicy(t *testing.T) {
 }
 
 func TestQueryHashPolicy(t *testing.T) {
-	pool := testPool()
+	ctx, cancel := caddy.NewContext(caddy.Context{Context: context.Background()})
+	defer cancel()
 	queryPolicy := QueryHashSelection{Key: "foo"}
+	if err := queryPolicy.Provision(ctx); err != nil {
+		t.Errorf("Provision error: %v", err)
+		t.FailNow()
+	}
+
+	pool := testPool()
 
 	request := httptest.NewRequest(http.MethodGet, "/?foo=1", nil)
 	h := queryPolicy.Select(pool, request, nil)
@@ -463,7 +472,7 @@ func TestQueryHashPolicy(t *testing.T) {
 
 func TestURIHashPolicy(t *testing.T) {
 	pool := testPool()
-	uriPolicy := new(URIHashSelection)
+	uriPolicy := URIHashSelection{}
 
 	request := httptest.NewRequest(http.MethodGet, "/test", nil)
 	h := uriPolicy.Select(pool, request, nil)
@@ -552,8 +561,7 @@ func TestRandomChoicePolicy(t *testing.T) {
 	pool[2].countRequest(30)
 
 	request := httptest.NewRequest(http.MethodGet, "/test", nil)
-	randomChoicePolicy := new(RandomChoiceSelection)
-	randomChoicePolicy.Choose = 2
+	randomChoicePolicy := RandomChoiceSelection{Choose: 2}
 
 	h := randomChoicePolicy.Select(pool, request, nil)
 
@@ -568,6 +576,14 @@ func TestRandomChoicePolicy(t *testing.T) {
 }
 
 func TestCookieHashPolicy(t *testing.T) {
+	ctx, cancel := caddy.NewContext(caddy.Context{Context: context.Background()})
+	defer cancel()
+	cookieHashPolicy := CookieHashSelection{}
+	if err := cookieHashPolicy.Provision(ctx); err != nil {
+		t.Errorf("Provision error: %v", err)
+		t.FailNow()
+	}
+
 	pool := testPool()
 	pool[0].Dial = "localhost:8080"
 	pool[1].Dial = "localhost:8081"
@@ -577,7 +593,7 @@ func TestCookieHashPolicy(t *testing.T) {
 	pool[2].setHealthy(false)
 	request := httptest.NewRequest(http.MethodGet, "/test", nil)
 	w := httptest.NewRecorder()
-	cookieHashPolicy := new(CookieHashSelection)
+
 	h := cookieHashPolicy.Select(pool, request, w)
 	cookieServer1 := w.Result().Cookies()[0]
 	if cookieServer1 == nil {
@@ -614,3 +630,59 @@ func TestCookieHashPolicy(t *testing.T) {
 		t.Error("Expected cookieHashPolicy to set a new cookie.")
 	}
 }
+
+func TestCookieHashPolicyWithFirstFallback(t *testing.T) {
+	ctx, cancel := caddy.NewContext(caddy.Context{Context: context.Background()})
+	defer cancel()
+	cookieHashPolicy := CookieHashSelection{
+		FallbackRaw: caddyconfig.JSONModuleObject(FirstSelection{}, "policy", "first", nil),
+	}
+	if err := cookieHashPolicy.Provision(ctx); err != nil {
+		t.Errorf("Provision error: %v", err)
+		t.FailNow()
+	}
+
+	pool := testPool()
+	pool[0].Dial = "localhost:8080"
+	pool[1].Dial = "localhost:8081"
+	pool[2].Dial = "localhost:8082"
+	pool[0].setHealthy(true)
+	pool[1].setHealthy(true)
+	pool[2].setHealthy(true)
+	request := httptest.NewRequest(http.MethodGet, "/test", nil)
+	w := httptest.NewRecorder()
+
+	h := cookieHashPolicy.Select(pool, request, w)
+	cookieServer1 := w.Result().Cookies()[0]
+	if cookieServer1 == nil {
+		t.Fatal("cookieHashPolicy should set a cookie")
+	}
+	if cookieServer1.Name != "lb" {
+		t.Error("cookieHashPolicy should set a cookie with name lb")
+	}
+	if h != pool[0] {
+		t.Errorf("Expected cookieHashPolicy host to be the first only available host, got %s", h)
+	}
+	request = httptest.NewRequest(http.MethodGet, "/test", nil)
+	w = httptest.NewRecorder()
+	request.AddCookie(cookieServer1)
+	h = cookieHashPolicy.Select(pool, request, w)
+	if h != pool[0] {
+		t.Errorf("Expected cookieHashPolicy host to stick to the first host (matching cookie), got %s", h)
+	}
+	s := w.Result().Cookies()
+	if len(s) != 0 {
+		t.Error("Expected cookieHashPolicy to not set a new cookie.")
+	}
+	pool[0].setHealthy(false)
+	request = httptest.NewRequest(http.MethodGet, "/test", nil)
+	w = httptest.NewRecorder()
+	request.AddCookie(cookieServer1)
+	h = cookieHashPolicy.Select(pool, request, w)
+	if h != pool[1] {
+		t.Errorf("Expected cookieHashPolicy to select the next first available host, got %s", h)
+	}
+	if w.Result().Cookies() == nil {
+		t.Error("Expected cookieHashPolicy to set a new cookie.")
+	}
+}