From 9b5809d3e82a806e69b1dabd6fc19a32a2f4b95d Mon Sep 17 00:00:00 2001
From: brettlangdon <me@brett.is>
Date: Mon, 18 Jan 2016 10:48:05 -0500
Subject: [PATCH] support multiple zone files

---
 cmd/realm/main.go | 35 +++++++++++++++++++++++------------
 server.go         |  8 ++++----
 zones.go          | 18 ++++++++++++++++++
 3 files changed, 45 insertions(+), 16 deletions(-)
 create mode 100644 zones.go

diff --git a/cmd/realm/main.go b/cmd/realm/main.go
index cf3c09d..63e4ffb 100644
--- a/cmd/realm/main.go
+++ b/cmd/realm/main.go
@@ -9,16 +9,18 @@ import (
 
 func main() {
 	// Setup our CLI app
-	var app *cli.App = cli.NewApp()
+	var app *cli.App
+	app = cli.NewApp()
 	app.Name = "realm"
 	app.Usage = "A simple non-recursive DNS server"
 	app.Version = "0.1.0"
 	app.Author = "Brett Langdon"
 	app.Email = "me@brett.is"
 	app.Flags = []cli.Flag{
-		cli.StringFlag{
+		cli.StringSliceFlag{
 			Name:   "zone, z",
 			EnvVar: "REALM_ZONE",
+			Value:  &cli.StringSlice{},
 			Usage:  "location to DNS zone file [required]",
 		},
 		cli.StringFlag{
@@ -32,24 +34,33 @@ func main() {
 	// This action is called for all commands
 	app.Action = func(c *cli.Context) {
 		// Ensure that a zone filename was provided
-		var filename string = c.String("zone")
-		if filename == "" {
-			log.Fatal("must supply zone file via \"--zone\" flag or \"REALM_ZONE\" environment variable")
+		var filenames []string
+		filenames = c.StringSlice("zone")
+		if len(filenames) == 0 {
+			log.Fatal("must supply at least 1 zone file via \"--zone\" flag or \"REALM_ZONE\" environment variable")
 		}
 
 		// Load and parse the zone file
-		var zone *realm.Zone
+		var zones realm.Zones
+		zones = make(realm.Zones, 0)
+
 		var err error
-		log.Printf("parsing zone file \"%s\"\n", filename)
-		zone, err = realm.ParseZone(filename)
-		if err != nil {
-			log.Fatal(err)
+		for _, filename := range filenames {
+			log.Printf("parsing zone file \"%s\"\n", filename)
+			var zone *realm.Zone
+			zone, err = realm.ParseZone(filename)
+			if err != nil {
+				log.Fatal(err)
+			}
+			zones = append(zones, zone)
 		}
 
 		// Create and start the server
-		var bind string = c.String("bind")
+		var bind string
+		bind = c.String("bind")
 		log.Printf("starting the server on \"%s\"\n", bind)
-		var server *realm.Server = realm.NewServer(bind, zone)
+		var server *realm.Server
+		server = realm.NewServer(bind, zones)
 		log.Fatal(server.ListenAndServe())
 	}
 
diff --git a/server.go b/server.go
index 9454ff7..29f491b 100644
--- a/server.go
+++ b/server.go
@@ -5,13 +5,13 @@ import "github.com/miekg/dns"
 // A Server listens for DNS requests over UDP and responds with answers from the provided Zone.
 type Server struct {
 	server *dns.Server
-	zone   *Zone
+	zones  Zones
 }
 
 // NewServer returns a new initialized *Server that will bind to listen and will look up answers from zone.
-func NewServer(listen string, zone *Zone) *Server {
+func NewServer(listen string, zones Zones) *Server {
 	var s *Server
-	s = &Server{zone: zone}
+	s = &Server{zones: zones}
 	s.server = &dns.Server{
 		Addr:    listen,
 		Net:     "udp",
@@ -37,7 +37,7 @@ func (s *Server) ServeDNS(w dns.ResponseWriter, request *dns.Msg) {
 	// Lookup answers to any of the questions
 	for _, question := range request.Question {
 		var records []dns.RR
-		records = s.zone.Lookup(question.Name, question.Qtype, question.Qclass)
+		records = s.zones.Lookup(question.Name, question.Qtype, question.Qclass)
 		response.Answer = append(response.Answer, records...)
 	}
 
diff --git a/zones.go b/zones.go
new file mode 100644
index 0000000..f4210ee
--- /dev/null
+++ b/zones.go
@@ -0,0 +1,18 @@
+package realm
+
+import (
+	"github.com/miekg/dns"
+)
+
+// Zones is a convenient helper for managing a slice of *Zone.
+type Zones []*Zone
+
+// Lookup will find all answer records from across all *Zone
+func (z Zones) Lookup(name string, reqType uint16, reqClass uint16) []dns.RR {
+	var records []dns.RR
+	records = make([]dns.RR, 0)
+	for _, zone := range z {
+		records = append(records, zone.Lookup(name, reqType, reqClass)...)
+	}
+	return records
+}