Refactor: Allow More Answer Types to Be Cacheable

makefile

@@ -1,6 +1,9 @@
 all: zdns
 
-zdns:
+generate:
+	go generate ./...
+
+zdns: generate
 	go build -o zdns
 
 clean:
@@ -37,5 +40,5 @@ benchmark: zdns
 
 ci: zdns lint test integration-tests license-check
 
-.PHONY: zdns clean test integration-tests lint ci license-check benchmark
+.PHONY: generate zdns clean test integration-tests lint ci license-check benchmark
 

src/zdns/answers.go

@@ -24,6 +24,12 @@ import (
 	"github.com/miekg/dns"
 )
 
+//go:generate go run answers_generate.go
+
+type WithBaseAnswer interface {
+	BaseAns() *Answer
+}
+
 type Answer struct {
 	TTL     uint32 `json:"ttl" groups:"ttl,normal,long,trace"`
 	Type    string `json:"type,omitempty" groups:"short,normal,long,trace"`

src/zdns/answers_generate.go

@@ -0,0 +1,193 @@
+//go:build ignore
+// +build ignore
+
+/*
+ * ZDNS Copyright 2024 Regents of the University of Michigan
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License"); you may not
+ * use this file except in compliance with the License. You may obtain a copy
+ * of the License at http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+ * implied. See the License for the specific language governing
+ * permissions and limitations under the License.
+ */
+
+/*
+ * BSD 3-Clause License
+
+ * Copyright (c) 2009, The Go Authors. Extensions copyright (c) 2011, Miek Gieben.
+ * All rights reserved.
+
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ *    list of conditions and the following disclaimer.
+
+ * 2. Redistributions in binary form must reproduce the above copyright notice,
+ *    this list of conditions and the following disclaimer in the documentation
+ *    and/or other materials provided with the distribution.
+
+ * 3. Neither the name of the copyright holder nor the names of its
+ *    contributors may be used to endorse or promote products derived from
+ *    this software without specific prior written permission.
+
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+ * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+ * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+ * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+ * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+ * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+ * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+ * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+ * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ */
+
+// answers_generate.go is meant to run with go generate. It will use
+// go/{importer,types} to track down all the DNS answer struct types that embed
+// the Answer struct. Then for each type, it will generate extraction and conversion
+// methods based on the struct tags. The generated source is written to answers_common.go,
+// and is meant to be checked into git.
+package main
+
+import (
+	"bytes"
+	"go/format"
+	"go/types"
+	"log"
+	"os"
+	"text/template"
+
+	"golang.org/x/tools/go/packages"
+)
+
+var packageHdr = `
+// Code generated by "go run answers_generate.go"; DO NOT EDIT.
+/*
+ * ZDNS Copyright 2024 Regents of the University of Michigan
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License"); you may not
+ * use this file except in compliance with the License. You may obtain a copy
+ * of the License at http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+ * implied. See the License for the specific language governing
+ * permissions and limitations under the License.
+ */
+
+/*
+ * BSD 3-Clause License
+
+ * Copyright (c) 2009, The Go Authors. Extensions copyright (c) 2011, Miek Gieben.
+ * All rights reserved.
+
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ *    list of conditions and the following disclaimer.
+
+ * 2. Redistributions in binary form must reproduce the above copyright notice,
+ *    this list of conditions and the following disclaimer in the documentation
+ *    and/or other materials provided with the distribution.
+
+ * 3. Neither the name of the copyright holder nor the names of its
+ *    contributors may be used to endorse or promote products derived from
+ *    this software without specific prior written permission.
+
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+ * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+ * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+ * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+ * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+ * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+ * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+ * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+ * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ */
+
+package zdns
+
+`
+
+var baseAnsFunc = template.Must(template.New("answerHeaderFunc").Parse(`
+func (ans Answer) BaseAns() *Answer { return &ans }
+{{range .}}  func (ans {{.}}) BaseAns() *Answer { return &ans.Answer }
+{{end}}
+`))
+
+func main() {
+	// Import and type-check the package
+	pkg, err := loadModule("github.com/zmap/zdns/src/zdns")
+	fatalIfErr(err)
+	scope := pkg.Scope()
+
+	// Collect all answer types (*X) that embed the Answer struct
+	var answerTypes []string
+	for _, name := range scope.Names() {
+		o := scope.Lookup(name)
+		if o == nil || !o.Exported() {
+			continue
+		}
+		if _, isAnswerType := getAnswerType(o.Type(), scope); isAnswerType {
+			answerTypes = append(answerTypes, o.Name())
+		}
+	}
+
+	b := &bytes.Buffer{}
+	b.WriteString(packageHdr)
+
+	// Generate answerHeaderFunc
+	fatalIfErr(baseAnsFunc.Execute(b, answerTypes))
+
+	// Format the generated code
+	res, err := format.Source(b.Bytes())
+	if err != nil {
+		b.WriteTo(os.Stderr)
+		log.Fatal(err)
+	}
+
+	// Write the result to answers_helper.go
+	f, err := os.Create("answers_helper.go")
+	fatalIfErr(err)
+	defer f.Close()
+	f.Write(res)
+}
+
+// loadModule retrieves package description for a given module.
+func loadModule(name string) (*types.Package, error) {
+	conf := packages.Config{Mode: packages.NeedTypes | packages.NeedTypesInfo}
+	pkgs, err := packages.Load(&conf, name)
+	if err != nil {
+		return nil, err
+	}
+	return pkgs[0].Types, nil
+}
+
+// getAnswerType checks if a type embeds the Answer struct and returns true if it does.
+func getAnswerType(t types.Type, scope *types.Scope) (*types.Struct, bool) {
+	st, ok := t.Underlying().(*types.Struct)
+	if !ok {
+		return nil, false
+	}
+	if st.NumFields() == 0 {
+		return nil, false
+	}
+	if st.Field(0).Type() == scope.Lookup("Answer").Type() {
+		return st, true
+	}
+	return nil, false
+}
+
+func fatalIfErr(err error) {
+	if err != nil {
+		log.Fatal(err)
+	}
+}