blob: d1e466ec34df4f8a7a1bd14a773311f049f751f5 (
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
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
|
package middleware
import (
"time"
"github.com/miekg/dns"
)
// ResponseRecorder 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 ResponseRecorder struct {
dns.ResponseWriter
rcode int
size int
msg *dns.Msg
start time.Time
}
// NewResponseRecorder makes and returns a new responseRecorder,
// which captures the DNS rcode from the ResponseWriter
// and also the length of the response message written through it.
func NewResponseRecorder(w dns.ResponseWriter) *ResponseRecorder {
return &ResponseRecorder{
ResponseWriter: w,
rcode: 0,
msg: nil,
start: time.Now(),
}
}
// WriteMsg records the status code and calls the
// underlying ResponseWriter's WriteMsg method.
func (r *ResponseRecorder) 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.size += res.Len()
r.msg = res
return r.ResponseWriter.WriteMsg(res)
}
// Write is a wrapper that records the size of the message that gets written.
func (r *ResponseRecorder) Write(buf []byte) (int, error) {
n, err := r.ResponseWriter.Write(buf)
if err == nil {
r.size += n
}
return n, err
}
// Size returns the size.
func (r *ResponseRecorder) Size() int { return r.size }
// Rcode returns the rcode.
func (r *ResponseRecorder) Rcode() string { return RcodeToString(r.rcode) }
// Start returns the start time of the ResponseRecorder.
func (r *ResponseRecorder) Start() time.Time { return r.start }
// Msg returns the written message from the ResponseRecorder.
func (r *ResponseRecorder) Msg() *dns.Msg { return r.msg }
// Hijack implements dns.Hijacker. It simply wraps the underlying
// ResponseWriter's Hijack method if there is one, or returns an error.
func (r *ResponseRecorder) Hijack() {
r.ResponseWriter.Hijack()
return
}
|
