Fix test

app/dns/dnscommon.go

@@ -31,13 +31,11 @@ type record struct {
 
 // IPRecord is a cacheable item for a resolved domain
 type IPRecord struct {
-	ReqID  uint16
-	IP     []net.Address
-	Expire time.Time
-	RCode  dnsmessage.RCode
-
-	// Truncated is for udp dns to indicates if the response is truncated and needs to be retried
-	Truncated bool
+	ReqID     uint16
+	IP        []net.Address
+	Expire    time.Time
+	RCode     dnsmessage.RCode
+	RawHeader *dnsmessage.Header
 }
 
 func (r *IPRecord) getIPs() ([]net.Address, error) {
@@ -68,7 +66,6 @@ type dnsRequest struct {
 	start   time.Time
 	expire  time.Time
 	msg     *dnsmessage.Message
-	ctx     context.Context
 }
 
 func genEDNS0Options(clientIP net.IP) *dnsmessage.Resource {
@@ -186,7 +183,7 @@ func parseResponse(payload []byte) (*IPRecord, error) {
 		ReqID:     h.ID,
 		RCode:     h.RCode,
 		Expire:    now.Add(time.Second * 600),
-		Truncated: h.Truncated,
+		RawHeader: &h,
 	}
 
 L:

app/dns/dnscommon_test.go

@@ -51,7 +51,7 @@ func Test_parseResponse(t *testing.T) {
 	}{
 		{
 			"empty",
-			&IPRecord{0, []net.Address(nil), time.Time{}, dnsmessage.RCodeSuccess},
+			&IPRecord{0, []net.Address(nil), time.Time{}, dnsmessage.RCodeSuccess, nil},
 			false,
 		},
 		{
@@ -66,12 +66,13 @@ func Test_parseResponse(t *testing.T) {
 				[]net.Address{net.ParseAddress("8.8.8.8"), net.ParseAddress("8.8.4.4")},
 				time.Time{},
 				dnsmessage.RCodeSuccess,
+				nil,
 			},
 			false,
 		},
 		{
 			"aaaa record",
-			&IPRecord{2, []net.Address{net.ParseAddress("2001::123:8888"), net.ParseAddress("2001::123:8844")}, time.Time{}, dnsmessage.RCodeSuccess},
+			&IPRecord{2, []net.Address{net.ParseAddress("2001::123:8888"), net.ParseAddress("2001::123:8844")}, time.Time{}, dnsmessage.RCodeSuccess, nil},
 			false,
 		},
 	}

app/dns/nameserver_udp.go

@@ -27,7 +27,7 @@ type ClassicNameServer struct {
 	name          string
 	address       *net.Destination
 	ips           map[string]*record
-	requests      map[uint16]*dnsRequest
+	requests      map[uint16]*udpDnsRequest
 	pub           *pubsub.Service
 	udpServer     *udp.Dispatcher
 	cleanup       *task.Periodic
@@ -35,6 +35,11 @@ type ClassicNameServer struct {
 	queryStrategy QueryStrategy
 }
 
+type udpDnsRequest struct {
+	dnsRequest
+	ctx context.Context
+}
+
 // NewClassicNameServer creates udp server object for remote resolving.
 func NewClassicNameServer(address net.Destination, dispatcher routing.Dispatcher, queryStrategy QueryStrategy) *ClassicNameServer {
 	// default to 53 if unspecific
@@ -45,7 +50,7 @@ func NewClassicNameServer(address net.Destination, dispatcher routing.Dispatcher
 	s := &ClassicNameServer{
 		address:       &address,
 		ips:           make(map[string]*record),
-		requests:      make(map[uint16]*dnsRequest),
+		requests:      make(map[uint16]*udpDnsRequest),
 		pub:           pubsub.NewService(),
 		name:          strings.ToUpper(address.String()),
 		queryStrategy: queryStrategy,
@@ -101,7 +106,7 @@ func (s *ClassicNameServer) Cleanup() error {
 	}
 
 	if len(s.requests) == 0 {
-		s.requests = make(map[uint16]*dnsRequest)
+		s.requests = make(map[uint16]*udpDnsRequest)
 	}
 
 	return nil
@@ -129,7 +134,7 @@ func (s *ClassicNameServer) HandleResponse(ctx context.Context, packet *udp_prot
 	}
 
 	// if truncated, retry with EDNS0 option(udp payload size: 1350)
-	if ipRec.Truncated {
+	if ipRec.RawHeader.Truncated {
 		// if already has EDNS0 option, no need to retry
 		if ok && len(req.msg.Additionals) == 0 {
 			// copy necessary meta data from original request
@@ -199,7 +204,7 @@ func (s *ClassicNameServer) newReqID() uint16 {
 	return uint16(atomic.AddUint32(&s.reqID, 1))
 }
 
-func (s *ClassicNameServer) addPendingRequest(req *dnsRequest) {
+func (s *ClassicNameServer) addPendingRequest(req *udpDnsRequest) {
 	s.Lock()
 	defer s.Unlock()
 
@@ -214,8 +219,11 @@ func (s *ClassicNameServer) sendQuery(ctx context.Context, domain string, client
 	reqs := buildReqMsgs(domain, option, s.newReqID, genEDNS0Options(clientIP))
 
 	for _, req := range reqs {
-		req.ctx = ctx
-		s.addPendingRequest(req)
+		udpReq := &udpDnsRequest{
+			dnsRequest: *req,
+			ctx:        ctx,
+		}
+		s.addPendingRequest(udpReq)
 		b, _ := dns.PackMessage(req.msg)
 		s.udpServer.Dispatch(toDnsContext(ctx, s.address.String()), *s.address, b)
 	}