Fast and extensible multi-platform HTTP/1-2-3 web server with automatic HTTPS
Mirrored from https://github.com/Slazee/caddy.git
Synced Nov 13, 03:00 UTC
reverseproxy: Add `fallback` for some policies, instead of always random (#5488)
Francis Lavoie committed
May 5, 2023 at 17:08 UTC
48598e1f2a370c2440b38f0b77e4d74748111b9a
2 files changed
+237
-39
modules/caddyhttp/reverseproxy/selectionpolicies.go
+155
-29
@@ -18,6 +18,7 @@ import (
18
"crypto/hmac"
19
"crypto/sha256"
20
"encoding/hex"
21
+ "encoding/json"
22
"fmt"
23
"hash/fnv"
24
weakrand "math/rand"
@@ -29,6 +30,7 @@ import (
30
"time"
31
32
"github.com/caddyserver/caddy/v2"
33
+ "github.com/caddyserver/caddy/v2/caddyconfig"
34
"github.com/caddyserver/caddy/v2/caddyconfig/caddyfile"
35
"github.com/caddyserver/caddy/v2/modules/caddyhttp"
36
)
@@ -372,6 +374,10 @@ func (r *URIHashSelection) UnmarshalCaddyfile(d *caddyfile.Dispenser) error {
374
type QueryHashSelection struct {
375
// The query key whose value is to be hashed and used for upstream selection.
376
Key string `json:"key,omitempty"`
377
+
378
+ // The fallback policy to use if the query key is not present. Defaults to `random`.
379
+ FallbackRaw json.RawMessage `json:"fallback,omitempty" caddy:"namespace=http.reverse_proxy.selection_policies inline_key=policy"`
380
+ fallback Selector
381
}
382
383
// CaddyModule returns the Caddy module information.
@@ -382,12 +388,24 @@ func (QueryHashSelection) CaddyModule() caddy.ModuleInfo {
388
}
389
}
390
385
-// Select returns an available host, if any.
386
-func (s QueryHashSelection) Select(pool UpstreamPool, req *http.Request, _ http.ResponseWriter) *Upstream {
391
+// Provision sets up the module.
392
+func (s *QueryHashSelection) Provision(ctx caddy.Context) error {
393
if s.Key == "" {
388
- return nil
394
+ return fmt.Errorf("query key is required")
395
+ }
396
+ if s.FallbackRaw == nil {
397
+ s.FallbackRaw = caddyconfig.JSONModuleObject(RandomSelection{}, "policy", "random", nil)
398
}
399
+ mod, err := ctx.LoadModule(s, "FallbackRaw")
400
+ if err != nil {
401
+ return fmt.Errorf("loading fallback selection policy: %s", err)
402
+ }
403
+ s.fallback = mod.(Selector)
404
+ return nil
405
+}
406
407
+// Select returns an available host, if any.
408
+func (s QueryHashSelection) Select(pool UpstreamPool, req *http.Request, _ http.ResponseWriter) *Upstream {
409
// Since the query may have multiple values for the same key,
410
// we'll join them to avoid a problem where the user can control
411
// 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.
415
// different request, because the order of the values is significant.
416
vals := strings.Join(req.URL.Query()[s.Key], ",")
417
if vals == "" {
400
- return RandomSelection{}.Select(pool, req, nil)
418
+ return s.fallback.Select(pool, req, nil)
419
}
420
return hostByHashing(pool, vals)
421
}
@@ -410,6 +428,24 @@ func (s *QueryHashSelection) UnmarshalCaddyfile(d *caddyfile.Dispenser) error {
428
}
429
s.Key = d.Val()
430
}
431
+ for nesting := d.Nesting(); d.NextBlock(nesting); {
432
+ switch d.Val() {
433
+ case "fallback":
434
+ if !d.NextArg() {
435
+ return d.ArgErr()
436
+ }
437
+ if s.FallbackRaw != nil {
438
+ return d.Err("fallback selection policy already specified")
439
+ }
440
+ mod, err := loadFallbackPolicy(d)
441
+ if err != nil {
442
+ return err
443
+ }
444
+ s.FallbackRaw = mod
445
+ default:
446
+ return d.Errf("unrecognized option '%s'", d.Val())
447
+ }
448
+ }
449
return nil
450
}
451
@@ -418,6 +454,10 @@ func (s *QueryHashSelection) UnmarshalCaddyfile(d *caddyfile.Dispenser) error {
454
type HeaderHashSelection struct {
455
// The HTTP header field whose value is to be hashed and used for upstream selection.
456
Field string `json:"field,omitempty"`
457
+
458
+ // The fallback policy to use if the header is not present. Defaults to `random`.
459
+ FallbackRaw json.RawMessage `json:"fallback,omitempty" caddy:"namespace=http.reverse_proxy.selection_policies inline_key=policy"`
460
+ fallback Selector
461
}
462
463
// CaddyModule returns the Caddy module information.
@@ -428,12 +468,24 @@ func (HeaderHashSelection) CaddyModule() caddy.ModuleInfo {
468
}
469
}
470
431
-// Select returns an available host, if any.
432
-func (s HeaderHashSelection) Select(pool UpstreamPool, req *http.Request, _ http.ResponseWriter) *Upstream {
471
+// Provision sets up the module.
472
+func (s *HeaderHashSelection) Provision(ctx caddy.Context) error {
473
if s.Field == "" {
434
- return nil
474
+ return fmt.Errorf("header field is required")
475
+ }
476
+ if s.FallbackRaw == nil {
477
+ s.FallbackRaw = caddyconfig.JSONModuleObject(RandomSelection{}, "policy", "random", nil)
478
}
479
+ mod, err := ctx.LoadModule(s, "FallbackRaw")
480
+ if err != nil {
481
+ return fmt.Errorf("loading fallback selection policy: %s", err)
482
+ }
483
+ s.fallback = mod.(Selector)
484
+ return nil
485
+}
486
487
+// Select returns an available host, if any.
488
+func (s HeaderHashSelection) Select(pool UpstreamPool, req *http.Request, _ http.ResponseWriter) *Upstream {
489
// The Host header should be obtained from the req.Host field
490
// since net/http removes it from the header map.
491
if s.Field == "Host" && req.Host != "" {
@@ -442,7 +494,7 @@ func (s HeaderHashSelection) Select(pool UpstreamPool, req *http.Request, _ http
494
495
val := req.Header.Get(s.Field)
496
if val == "" {
445
- return RandomSelection{}.Select(pool, req, nil)
497
+ return s.fallback.Select(pool, req, nil)
498
}
499
return hostByHashing(pool, val)
500
}
@@ -455,6 +507,24 @@ func (s *HeaderHashSelection) UnmarshalCaddyfile(d *caddyfile.Dispenser) error {
507
}
508
s.Field = d.Val()
509
}
510
+ for nesting := d.Nesting(); d.NextBlock(nesting); {
511
+ switch d.Val() {
512
+ case "fallback":
513
+ if !d.NextArg() {
514
+ return d.ArgErr()
515
+ }
516
+ if s.FallbackRaw != nil {
517
+ return d.Err("fallback selection policy already specified")
518
+ }
519
+ mod, err := loadFallbackPolicy(d)
520
+ if err != nil {
521
+ return err
522
+ }
523
+ s.FallbackRaw = mod
524
+ default:
525
+ return d.Errf("unrecognized option '%s'", d.Val())
526
+ }
527
+ }
528
return nil
529
}
530
@@ -465,6 +535,10 @@ type CookieHashSelection struct {
535
Name string `json:"name,omitempty"`
536
// Secret to hash (Hmac256) chosen upstream in cookie
537
Secret string `json:"secret,omitempty"`
538
+
539
+ // The fallback policy to use if the cookie is not present. Defaults to `random`.
540
+ FallbackRaw json.RawMessage `json:"fallback,omitempty" caddy:"namespace=http.reverse_proxy.selection_policies inline_key=policy"`
541
+ fallback Selector
542
}
543
544
// CaddyModule returns the Caddy module information.
@@ -475,15 +549,48 @@ func (CookieHashSelection) CaddyModule() caddy.ModuleInfo {
549
}
550
}
551
478
-// Select returns an available host, if any.
479
-func (s CookieHashSelection) Select(pool UpstreamPool, req *http.Request, w http.ResponseWriter) *Upstream {
552
+// Provision sets up the module.
553
+func (s *CookieHashSelection) Provision(ctx caddy.Context) error {
554
if s.Name == "" {
555
s.Name = "lb"
556
}
557
+ if s.FallbackRaw == nil {
558
+ s.FallbackRaw = caddyconfig.JSONModuleObject(RandomSelection{}, "policy", "random", nil)
559
+ }
560
+ mod, err := ctx.LoadModule(s, "FallbackRaw")
561
+ if err != nil {
562
+ return fmt.Errorf("loading fallback selection policy: %s", err)
563
+ }
564
+ s.fallback = mod.(Selector)
565
+ return nil
566
+}
567
+
568
+// Select returns an available host, if any.
569
+func (s CookieHashSelection) Select(pool UpstreamPool, req *http.Request, w http.ResponseWriter) *Upstream {
570
+ // selects a new Host using the fallback policy (typically random)
571
+ // and write a sticky session cookie to the response.
572
+ selectNewHost := func() *Upstream {
573
+ upstream := s.fallback.Select(pool, req, w)
574
+ if upstream == nil {
575
+ return nil
576
+ }
577
+ sha, err := hashCookie(s.Secret, upstream.Dial)
578
+ if err != nil {
579
+ return upstream
580
+ }
581
+ http.SetCookie(w, &http.Cookie{
582
+ Name: s.Name,
583
+ Value: sha,
584
+ Path: "/",
585
+ Secure: false,
586
+ })
587
+ return upstream
588
+ }
589
+
590
cookie, err := req.Cookie(s.Name)
484
- // If there's no cookie, select new random host
591
+ // If there's no cookie, select a host using the fallback policy
592
if err != nil || cookie == nil {
486
- return selectNewHostWithCookieHashSelection(pool, w, s.Secret, s.Name)
593
+ return selectNewHost()
594
}
595
// If the cookie is present, loop over the available upstreams until we find a match
596
cookieValue := cookie.Value
@@ -496,13 +603,15 @@ func (s CookieHashSelection) Select(pool UpstreamPool, req *http.Request, w http
603
return upstream
604
}
605
}
499
- // If there is no matching host, select new random host
500
- return selectNewHostWithCookieHashSelection(pool, w, s.Secret, s.Name)
606
+ // If there is no matching host, select a host using the fallback policy
607
+ return selectNewHost()
608
}
609
610
// UnmarshalCaddyfile sets up the module from Caddyfile tokens. Syntax:
611
//
505
-// lb_policy cookie [<name> [<secret>]]
612
+// lb_policy cookie [<name> [<secret>]] {
613
+// fallback <policy>
614
+// }
615
//
616
// By default name is `lb`
617
func (s *CookieHashSelection) UnmarshalCaddyfile(d *caddyfile.Dispenser) error {
@@ -517,22 +626,25 @@ func (s *CookieHashSelection) UnmarshalCaddyfile(d *caddyfile.Dispenser) error {
626
default:
627
return d.ArgErr()
628
}
520
- return nil
521
-}
522
-
523
-// Select a new Host randomly and add a sticky session cookie
524
-func selectNewHostWithCookieHashSelection(pool []*Upstream, w http.ResponseWriter, cookieSecret string, cookieName string) *Upstream {
525
- randomHost := selectRandomHost(pool)
526
-
527
- if randomHost != nil {
528
- // Hash (HMAC with some key for privacy) the upstream.Dial string as the cookie value
529
- sha, err := hashCookie(cookieSecret, randomHost.Dial)
530
- if err == nil {
531
- // write the cookie.
532
- http.SetCookie(w, &http.Cookie{Name: cookieName, Value: sha, Path: "/", Secure: false})
629
+ for nesting := d.Nesting(); d.NextBlock(nesting); {
630
+ switch d.Val() {
631
+ case "fallback":
632
+ if !d.NextArg() {
633
+ return d.ArgErr()
634
+ }
635
+ if s.FallbackRaw != nil {
636
+ return d.Err("fallback selection policy already specified")
637
+ }
638
+ mod, err := loadFallbackPolicy(d)
639
+ if err != nil {
640
+ return err
641
+ }
642
+ s.FallbackRaw = mod
643
+ default:
644
+ return d.Errf("unrecognized option '%s'", d.Val())
645
}
646
}
535
- return randomHost
647
+ return nil
648
}
649
650
// hashCookie hashes (HMAC 256) some data with the secret
@@ -627,6 +739,20 @@ func hash(s string) uint32 {
739
return h.Sum32()
740
}
741
742
+func loadFallbackPolicy(d *caddyfile.Dispenser) (json.RawMessage, error) {
743
+ name := d.Val()
744
+ modID := "http.reverse_proxy.selection_policies." + name
745
+ unm, err := caddyfile.UnmarshalModule(d, modID)
746
+ if err != nil {
747
+ return nil, err
748
+ }
749
+ sel, ok := unm.(Selector)
750
+ if !ok {
751
+ return nil, d.Errf("module %s (%T) is not a reverseproxy.Selector", modID, unm)
752
+ }
753
+ return caddyconfig.JSONModuleObject(sel, "policy", name, nil), nil
754
+}
755
+
756
// Interface guards
757
var (
758
_ Selector = (*RandomSelection)(nil)
modules/caddyhttp/reverseproxy/selectionpolicies_test.go
+82
-10
@@ -20,6 +20,8 @@ import (
20
"net/http/httptest"
21
"testing"
22
23
+ "github.com/caddyserver/caddy/v2"
24
+ "github.com/caddyserver/caddy/v2/caddyconfig"
25
"github.com/caddyserver/caddy/v2/modules/caddyhttp"
26
)
27
@@ -33,7 +35,7 @@ func testPool() UpstreamPool {
35
36
func TestRoundRobinPolicy(t *testing.T) {
37
pool := testPool()
36
- rrPolicy := new(RoundRobinSelection)
38
+ rrPolicy := RoundRobinSelection{}
39
req, _ := http.NewRequest("GET", "/", nil)
40
41
h := rrPolicy.Select(pool, req, nil)
@@ -74,7 +76,7 @@ func TestRoundRobinPolicy(t *testing.T) {
76
77
func TestLeastConnPolicy(t *testing.T) {
78
pool := testPool()
77
- lcPolicy := new(LeastConnSelection)
79
+ lcPolicy := LeastConnSelection{}
80
req, _ := http.NewRequest("GET", "/", nil)
81
82
pool[0].countRequest(10)
@@ -92,7 +94,7 @@ func TestLeastConnPolicy(t *testing.T) {
94
95
func TestIPHashPolicy(t *testing.T) {
96
pool := testPool()
95
- ipHash := new(IPHashSelection)
97
+ ipHash := IPHashSelection{}
98
req, _ := http.NewRequest("GET", "/", nil)
99
100
// We should be able to predict where every request is routed.
@@ -234,7 +236,7 @@ func TestIPHashPolicy(t *testing.T) {
236
237
func TestClientIPHashPolicy(t *testing.T) {
238
pool := testPool()
237
- ipHash := new(ClientIPHashSelection)
239
+ ipHash := ClientIPHashSelection{}
240
req, _ := http.NewRequest("GET", "/", nil)
241
req = req.WithContext(context.WithValue(req.Context(), caddyhttp.VarsCtxKey, make(map[string]any)))
242
@@ -377,7 +379,7 @@ func TestClientIPHashPolicy(t *testing.T) {
379
380
func TestFirstPolicy(t *testing.T) {
381
pool := testPool()
380
- firstPolicy := new(FirstSelection)
382
+ firstPolicy := FirstSelection{}
383
req := httptest.NewRequest(http.MethodGet, "/", nil)
384
385
h := firstPolicy.Select(pool, req, nil)
@@ -393,8 +395,15 @@ func TestFirstPolicy(t *testing.T) {
395
}
396
397
func TestQueryHashPolicy(t *testing.T) {
396
- pool := testPool()
398
+ ctx, cancel := caddy.NewContext(caddy.Context{Context: context.Background()})
399
+ defer cancel()
400
queryPolicy := QueryHashSelection{Key: "foo"}
401
+ if err := queryPolicy.Provision(ctx); err != nil {
402
+ t.Errorf("Provision error: %v", err)
403
+ t.FailNow()
404
+ }
405
+
406
+ pool := testPool()
407
408
request := httptest.NewRequest(http.MethodGet, "/?foo=1", nil)
409
h := queryPolicy.Select(pool, request, nil)
@@ -463,7 +472,7 @@ func TestQueryHashPolicy(t *testing.T) {
472
473
func TestURIHashPolicy(t *testing.T) {
474
pool := testPool()
466
- uriPolicy := new(URIHashSelection)
475
+ uriPolicy := URIHashSelection{}
476
477
request := httptest.NewRequest(http.MethodGet, "/test", nil)
478
h := uriPolicy.Select(pool, request, nil)
@@ -552,8 +561,7 @@ func TestRandomChoicePolicy(t *testing.T) {
561
pool[2].countRequest(30)
562
563
request := httptest.NewRequest(http.MethodGet, "/test", nil)
555
- randomChoicePolicy := new(RandomChoiceSelection)
556
- randomChoicePolicy.Choose = 2
564
+ randomChoicePolicy := RandomChoiceSelection{Choose: 2}
565
566
h := randomChoicePolicy.Select(pool, request, nil)
567
@@ -568,6 +576,14 @@ func TestRandomChoicePolicy(t *testing.T) {
576
}
577
578
func TestCookieHashPolicy(t *testing.T) {
579
+ ctx, cancel := caddy.NewContext(caddy.Context{Context: context.Background()})
580
+ defer cancel()
581
+ cookieHashPolicy := CookieHashSelection{}
582
+ if err := cookieHashPolicy.Provision(ctx); err != nil {
583
+ t.Errorf("Provision error: %v", err)
584
+ t.FailNow()
585
+ }
586
+
587
pool := testPool()
588
pool[0].Dial = "localhost:8080"
589
pool[1].Dial = "localhost:8081"
@@ -577,7 +593,7 @@ func TestCookieHashPolicy(t *testing.T) {
593
pool[2].setHealthy(false)
594
request := httptest.NewRequest(http.MethodGet, "/test", nil)
595
w := httptest.NewRecorder()
580
- cookieHashPolicy := new(CookieHashSelection)
596
+
597
h := cookieHashPolicy.Select(pool, request, w)
598
cookieServer1 := w.Result().Cookies()[0]
599
if cookieServer1 == nil {
@@ -614,3 +630,59 @@ func TestCookieHashPolicy(t *testing.T) {
630
t.Error("Expected cookieHashPolicy to set a new cookie.")
631
}
632
}
633
+
634
+func TestCookieHashPolicyWithFirstFallback(t *testing.T) {
635
+ ctx, cancel := caddy.NewContext(caddy.Context{Context: context.Background()})
636
+ defer cancel()
637
+ cookieHashPolicy := CookieHashSelection{
638
+ FallbackRaw: caddyconfig.JSONModuleObject(FirstSelection{}, "policy", "first", nil),
639
+ }
640
+ if err := cookieHashPolicy.Provision(ctx); err != nil {
641
+ t.Errorf("Provision error: %v", err)
642
+ t.FailNow()
643
+ }
644
+
645
+ pool := testPool()
646
+ pool[0].Dial = "localhost:8080"
647
+ pool[1].Dial = "localhost:8081"
648
+ pool[2].Dial = "localhost:8082"
649
+ pool[0].setHealthy(true)
650
+ pool[1].setHealthy(true)
651
+ pool[2].setHealthy(true)
652
+ request := httptest.NewRequest(http.MethodGet, "/test", nil)
653
+ w := httptest.NewRecorder()
654
+
655
+ h := cookieHashPolicy.Select(pool, request, w)
656
+ cookieServer1 := w.Result().Cookies()[0]
657
+ if cookieServer1 == nil {
658
+ t.Fatal("cookieHashPolicy should set a cookie")
659
+ }
660
+ if cookieServer1.Name != "lb" {
661
+ t.Error("cookieHashPolicy should set a cookie with name lb")
662
+ }
663
+ if h != pool[0] {
664
+ t.Errorf("Expected cookieHashPolicy host to be the first only available host, got %s", h)
665
+ }
666
+ request = httptest.NewRequest(http.MethodGet, "/test", nil)
667
+ w = httptest.NewRecorder()
668
+ request.AddCookie(cookieServer1)
669
+ h = cookieHashPolicy.Select(pool, request, w)
670
+ if h != pool[0] {
671
+ t.Errorf("Expected cookieHashPolicy host to stick to the first host (matching cookie), got %s", h)
672
+ }
673
+ s := w.Result().Cookies()
674
+ if len(s) != 0 {
675
+ t.Error("Expected cookieHashPolicy to not set a new cookie.")
676
+ }
677
+ pool[0].setHealthy(false)
678
+ request = httptest.NewRequest(http.MethodGet, "/test", nil)
679
+ w = httptest.NewRecorder()
680
+ request.AddCookie(cookieServer1)
681
+ h = cookieHashPolicy.Select(pool, request, w)
682
+ if h != pool[1] {
683
+ t.Errorf("Expected cookieHashPolicy to select the next first available host, got %s", h)
684
+ }
685
+ if w.Result().Cookies() == nil {
686
+ t.Error("Expected cookieHashPolicy to set a new cookie.")
687
+ }
688
+}