Use an interface method for unpacking records (#884)

* Use an interface method for unpacking records

* Eliminate err var declaration from unpack functions

* Remove pointless r.Data assignment in PrivateRR.unpack

dns.go

@@ -43,6 +43,12 @@ type RR interface {
 
 	// pack packs an RR into wire format.
 	pack(msg []byte, off int, compression compressionMap, compress bool) (headerEnd int, off1 int, err error)
+
+	// unpack unpacks an RR from wire format.
+	//
+	// This will only be called on a new and empty RR type with only the header populated. It
+	// will only be called if the record's RDATA is non-empty.
+	unpack(msg []byte, off int) (off1 int, err error)
 }
 
 // RR_Header is the header all DNS resource records share.
@@ -81,6 +87,10 @@ func (h *RR_Header) len(off int, compression map[string]struct{}) int {
 	return l
 }
 
+func (h *RR_Header) unpack(msg []byte, off int) (int, error) {
+	panic("dns: internal error: unpack should never be called on RR_Header")
+}
+
 // ToRFC3597 converts a known RR to the unknown RR representation from RFC 3597.
 func (rr *RFC3597) ToRFC3597(r RR) error {
 	buf := make([]byte, Len(r)*2)
@@ -90,14 +100,17 @@ func (rr *RFC3597) ToRFC3597(r RR) error {
 	}
 	buf = buf[:off]
 
-	hdr := *r.Header()
-	hdr.Rdlength = uint16(off - headerEnd)
+	*rr = RFC3597{Hdr: *r.Header()}
+	rr.Hdr.Rdlength = uint16(off - headerEnd)
 
-	rfc3597, _, err := unpackRFC3597(hdr, buf, headerEnd)
+	if noRdata(rr.Hdr) {
+		return nil
+	}
+
+	_, err = rr.unpack(buf, headerEnd)
 	if err != nil {
 		return err
 	}
 
-	*rr = *rfc3597.(*RFC3597)
 	return nil
 }

msg.go

@@ -661,17 +661,28 @@ func UnpackRR(msg []byte, off int) (rr RR, off1 int, err error) {
 // UnpackRRWithHeader unpacks the record type specific payload given an existing
 // RR_Header.
 func UnpackRRWithHeader(h RR_Header, msg []byte, off int) (rr RR, off1 int, err error) {
+	if newFn, ok := TypeToRR[h.Rrtype]; ok {
+		rr = newFn()
+		*rr.Header() = h
+	} else {
+		rr = &RFC3597{Hdr: h}
+	}
+
+	if noRdata(h) {
+		return rr, off, nil
+	}
+
 	end := off + int(h.Rdlength)
 
-	if fn, known := typeToUnpack[h.Rrtype]; !known {
-		rr, off, err = unpackRFC3597(h, msg, off)
-	} else {
-		rr, off, err = fn(h, msg, off)
+	off, err = rr.unpack(msg, off)
+	if err != nil {
+		return nil, end, err
 	}
 	if off != end {
 		return &h, end, &Error{err: "bad rdlength"}
 	}
-	return rr, off, err
+
+	return rr, off, nil
 }
 
 // unpackRRslice unpacks msg[off:] into an []RR.

msg_generate.go

@@ -184,15 +184,8 @@ if rr.%s != "-" {
 		o := scope.Lookup(name)
 		st, _ := getTypeStruct(o.Type(), scope)
 
-		fmt.Fprintf(b, "func unpack%s(h RR_Header, msg []byte, off int) (RR, int, error) {\n", name)
-		fmt.Fprintf(b, "rr := new(%s)\n", name)
-		fmt.Fprint(b, "rr.Hdr = h\n")
-		fmt.Fprint(b, `if noRdata(h) {
-return rr, off, nil
-	}
-var err error
-_ = err
-rdStart := off
+		fmt.Fprintf(b, "func (rr *%s) unpack(msg []byte, off int) (off1 int, err error) {\n", name)
+		fmt.Fprint(b, `rdStart := off
 _ = rdStart
 
 `)
@@ -200,7 +193,7 @@ _ = rdStart
 			o := func(s string) {
 				fmt.Fprintf(b, s, st.Field(i).Name())
 				fmt.Fprint(b, `if err != nil {
-return rr, off, err
+return off, err
 }
 `)
 			}
@@ -220,7 +213,7 @@ return rr, off, err
 					log.Fatalln(name, st.Field(i).Name(), st.Tag(i))
 				}
 				fmt.Fprint(b, `if err != nil {
-return rr, off, err
+return off, err
 }
 `)
 				continue
@@ -288,22 +281,13 @@ return rr, off, err
 			// If we've hit len(msg) we return without error.
 			if i < st.NumFields()-1 {
 				fmt.Fprintf(b, `if off == len(msg) {
-return rr, off, nil
+return off, nil
 	}
 `)
 			}
 		}
-		fmt.Fprintf(b, "return rr, off, nil }\n\n")
+		fmt.Fprintf(b, "return off, nil }\n\n")
 	}
-	// Generate typeToUnpack map
-	fmt.Fprintln(b, "var typeToUnpack = map[uint16]func(RR_Header, []byte, int) (RR, int, error){")
-	for _, name := range namedTypes {
-		if name == "RFC3597" {
-			continue
-		}
-		fmt.Fprintf(b, "Type%s: unpack%s,\n", name, name)
-	}
-	fmt.Fprintln(b, "}\n")
 
 	// gofmt
 	res, err := format.Source(b.Bytes())

privaterr.go

@@ -83,6 +83,16 @@ func (r *PrivateRR) pack(msg []byte, off int, compression compressionMap, compre
 	return headerEnd, off, nil
 }
 
+func (r *PrivateRR) unpack(msg []byte, off int) (int, error) {
+	off1, err := r.Data.Unpack(msg[off:])
+	off += off1
+	if err != nil {
+		return off, err
+	}
+
+	return off, nil
+}
+
 // PrivateHandle registers a private resource record type. It requires
 // string and numeric representation of private RR type and generator function as argument.
 func PrivateHandle(rtypestr string, rtype uint16, generator func() PrivateRdata) {
@@ -92,23 +102,6 @@ func PrivateHandle(rtypestr string, rtype uint16, generator func() PrivateRdata)
 	TypeToString[rtype] = rtypestr
 	StringToType[rtypestr] = rtype
 
-	typeToUnpack[rtype] = func(h RR_Header, msg []byte, off int) (RR, int, error) {
-		if noRdata(h) {
-			return &h, off, nil
-		}
-		var err error
-
-		rr := mkPrivateRR(h.Rrtype)
-		rr.Hdr = h
-
-		off1, err := rr.Data.Unpack(msg[off:])
-		off += off1
-		if err != nil {
-			return rr, off, err
-		}
-		return rr, off, err
-	}
-
 	setPrivateRR := func(h RR_Header, c *zlexer, o, f string) (RR, *ParseError) {
 		rr := mkPrivateRR(h.Rrtype)
 		rr.Hdr = h
@@ -146,6 +139,5 @@ func PrivateHandleRemove(rtype uint16) {
 		delete(TypeToString, rtype)
 		delete(typeToparserFunc, rtype)
 		delete(StringToType, rtypestr)
-		delete(typeToUnpack, rtype)
 	}
 }