From bf363f061d03f33a8301466c9c0e2a89d75542c0 Mon Sep 17 00:00:00 2001
From: Matthew Holt <mholt@users.noreply.github.com>
Date: Mon, 4 Nov 2019 12:54:46 -0700
Subject: [PATCH] reverse_proxy: Add UnmarshalCaddyfile for random_choose
 selection policy

Also allow caddy.Duration to be given integer values which are treated
like regular time.Duration values (nanoseconds).

Fixes #2856
---
 caddy.go                                      | 16 ++++++++++-----
 .../reverseproxy/selectionpolicies.go         | 20 +++++++++++++++++++
 2 files changed, 31 insertions(+), 5 deletions(-)

diff --git a/caddy.go b/caddy.go
index 7aab8d43..b4acde31 100644
--- a/caddy.go
+++ b/caddy.go
@@ -437,12 +437,18 @@ type Duration time.Duration
 
 // UnmarshalJSON satisfies json.Unmarshaler.
 func (d *Duration) UnmarshalJSON(b []byte) error {
-	dd, err := time.ParseDuration(strings.Trim(string(b), `"`))
-	if err != nil {
-		return err
+	if len(b) == 0 {
+		return io.EOF
 	}
-	*d = Duration(dd)
-	return nil
+	var dur time.Duration
+	var err error
+	if b[0] == byte('"') && b[len(b)-1] == byte('"') {
+		dur, err = time.ParseDuration(strings.Trim(string(b), `"`))
+	} else {
+		err = json.Unmarshal(b, &dur)
+	}
+	*d = Duration(dur)
+	return err
 }
 
 // GoModule returns the build info of this Caddy
diff --git a/modules/caddyhttp/reverseproxy/selectionpolicies.go b/modules/caddyhttp/reverseproxy/selectionpolicies.go
index 5bb2d62d..e34dd9b3 100644
--- a/modules/caddyhttp/reverseproxy/selectionpolicies.go
+++ b/modules/caddyhttp/reverseproxy/selectionpolicies.go
@@ -20,10 +20,12 @@ import (
 	weakrand "math/rand"
 	"net"
 	"net/http"
+	"strconv"
 	"sync/atomic"
 	"time"
 
 	"github.com/caddyserver/caddy/v2"
+	"github.com/caddyserver/caddy/v2/caddyconfig/caddyfile"
 )
 
 func init() {
@@ -87,6 +89,22 @@ func (RandomChoiceSelection) CaddyModule() caddy.ModuleInfo {
 	}
 }
 
+// UnmarshalCaddyfile sets up the module from Caddyfile tokens.
+func (r *RandomChoiceSelection) UnmarshalCaddyfile(d *caddyfile.Dispenser) error {
+	for d.Next() {
+		if !d.NextArg() {
+			return d.ArgErr()
+		}
+		chooseStr := d.Val()
+		choose, err := strconv.Atoi(chooseStr)
+		if err != nil {
+			return d.Errf("invalid choice value '%s': %v", chooseStr, err)
+		}
+		r.Choose = choose
+	}
+	return nil
+}
+
 // Provision sets up r.
 func (r *RandomChoiceSelection) Provision(ctx caddy.Context) error {
 	if r.Choose == 0 {
@@ -350,4 +368,6 @@ var (
 
 	_ caddy.Validator   = (*RandomChoiceSelection)(nil)
 	_ caddy.Provisioner = (*RandomChoiceSelection)(nil)
+
+	_ caddyfile.Unmarshaler = (*RandomChoiceSelection)(nil)
 )