directives_generate.go

//+build ignore

package main

import (
	"bufio"
	"go/format"
	"io/ioutil"
	"log"
	"os"
	"sort"
	"strconv"
	"strings"
)

func main() {
	mi := make(map[string]string, 0)
	md := make(map[int]string, 0)

	file, err := os.Open(middlewareFile)
	fatalIfErr(err)

	defer file.Close()

	scanner := bufio.NewScanner(file)
	for scanner.Scan() {
		line := scanner.Text()
		if strings.HasPrefix(line, "#") {
			continue
		}

		items := strings.Split(line, ":")
		if len(items) != 3 {
			// ignore
			continue
		}
		priority, err := strconv.Atoi(items[0])
		fatalIfErr(err)

		if v, ok := md[priority]; ok {
			log.Fatalf("Duplicate priority '%d', slot already taken by %q", priority, v)
		}
		md[priority] = items[1]
		mi[items[1]] = middlewarePath + items[2] // Default, unless overriden by 3rd arg

		if strings.Contains(items[2], "/") { // External package has been given
			mi[items[1]] = items[2]
		}
	}

	genImports("core/zmiddleware.go", "core", mi)
	genDirectives("core/dnsserver/zdirectives.go", "dnsserver", md)
}

func genImports(file, pack string, mi map[string]string) {
	outs := header + "package " + pack + "\n\n" + "import ("

	if len(mi) > 0 {
		outs += "\n"
	}

	outs += "// Include all middleware.\n"
	for _, v := range mi {
		outs += `_ "` + v + `"` + "\n"
	}
	outs += ")\n"

	res, err := format.Source([]byte(outs))
	fatalIfErr(err)

	err = ioutil.WriteFile(file, res, 0644)
	fatalIfErr(err)
}

func genDirectives(file, pack string, md map[int]string) {

	outs := header + "package " + pack + "\n\n"
	outs += `
// Directives are registered in the order they should be
// executed.
//
// Ordering is VERY important. Every middleware will
// feel the effects of all other middleware below
// (after) them during a request, but they must not
// care what middleware above them are doing.

var directives = []string{
`

	var orders []int
	for k := range md {
		orders = append(orders, k)
	}
	sort.Ints(orders)

	for _, k := range orders {
		outs += `"` + md[k] + `",` + "\n"
	}

	outs += "}\n"

	res, err := format.Source([]byte(outs))
	fatalIfErr(err)

	err = ioutil.WriteFile(file, res, 0644)
	fatalIfErr(err)
}

func fatalIfErr(err error) {
	if err != nil {
		log.Fatal(err)
	}
}

const (
	middlewarePath = "github.com/coredns/coredns/middleware/"
	middlewareFile = "middleware.cfg"
	header         = "// generated by directives_generate.go; DO NOT EDIT\n\n"
)