2 files changed, 19 insertions, 6 deletions
diff --git a/middleware/federation/setup.go b/middleware/federation/setup.go
index b5a9469ab..de37fce2b 100644
--- a/middleware/federation/setup.go
+++ b/middleware/federation/setup.go
@@ -49,17 +49,25 @@ func federationParse(c *caddy.Controller) (*Federation, error) {
 
 	for c.Next() {
 		// federation [zones..]
-		origins := make([]string, len(c.ServerBlockKeys))
-		copy(origins, c.ServerBlockKeys)
+		zones := c.RemainingArgs()
+		origins := []string{}
+		if len(zones) > 0 {
+			origins = make([]string, len(zones))
+			copy(origins, zones)
+		} else {
+			origins = make([]string, len(c.ServerBlockKeys))
+			copy(origins, c.ServerBlockKeys)
+		}
 
 		for c.NextBlock() {
 			x := c.Val()
-			switch c.Val() {
+			switch x {
 			default:
 				args := c.RemainingArgs()
-				if len(args) != 1 {
-					return fed, fmt.Errorf("need two arguments for federation: %q", args)
+				if x := len(args); x != 1 {
+					return fed, fmt.Errorf("need two arguments for federation, got %d", x)
 				}
+
 				fed.f[x] = dns.Fqdn(args[0])
 			}
 		}
diff --git a/middleware/federation/setup_test.go b/middleware/federation/setup_test.go
index b0d09419a..e85b01772 100644
--- a/middleware/federation/setup_test.go
+++ b/middleware/federation/setup_test.go
@@ -13,6 +13,7 @@ func TestSetup(t *testing.T) {
 		expectedLen      int
 		expectedNameZone []string // contains only entry for now
 	}{
+		// ok
 		{`federation {
 			prod prod.example.org
 		}`, false, 1, []string{"prod", "prod.example.org."}},
@@ -25,11 +26,15 @@ func TestSetup(t *testing.T) {
 			staging staging.example.org
 			prod prod.example.org
 		}`, false, 2, []string{"staging", "staging.example.org."}},
+		{`federation example.com {
+			staging staging.example.org
+			prod prod.example.org
+		}`, false, 2, []string{"staging", "staging.example.org."}},
 		// errors
 		{`federation {
 		}`, true, 0, []string{}},
 		{`federation {
-			staging 
+			staging
 		}`, true, 0, []string{}},
 	}
 	for i, test := range tests {