aboutsummaryrefslogtreecommitdiff
path: root/plugin/pkg/dnstest/recorder.go
blob: 1da063e144d28826aa37e4aa4be0d3b77b98878f (plain) (blame)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
// Package dnstest allows for easy testing of DNS client against a test server.
package dnstest

import (
	"time"

	"github.com/miekg/dns"
)

// Recorder is a type of ResponseWriter that captures
// the rcode code written to it and also the size of the message
// written in the response. A rcode code does not have
// to be written, however, in which case 0 must be assumed.
// It is best to have the constructor initialize this type
// with that default status code.
type Recorder struct {
	dns.ResponseWriter
	Rcode int
	Len   int
	Msg   *dns.Msg
	Start time.Time
}

// NewRecorder makes and returns a new Recorder,
// which captures the DNS rcode from the ResponseWriter
// and also the length of the response message written through it.
func NewRecorder(w dns.ResponseWriter) *Recorder {
	return &Recorder{
		ResponseWriter: w,
		Rcode:          0,
		Msg:            nil,
		Start:          time.Now(),
	}
}

// WriteMsg records the status code and calls the
// underlying ResponseWriter's WriteMsg method.
func (r *Recorder) WriteMsg(res *dns.Msg) error {
	r.Rcode = res.Rcode
	// We may get called multiple times (axfr for instance).
	// Save the last message, but add the sizes.
	r.Len += res.Len()
	r.Msg = res
	return r.ResponseWriter.WriteMsg(res)
}

// Write is a wrapper that records the length of the message that gets written.
func (r *Recorder) Write(buf []byte) (int, error) {
	n, err := r.ResponseWriter.Write(buf)
	if err == nil {
		r.Len += n
	}
	return n, err
}