-
Notifications
You must be signed in to change notification settings - Fork 2.1k
/
multirecorder.go
41 lines (35 loc) · 987 Bytes
/
multirecorder.go
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
package dnstest
import (
"time"
"github.com/miekg/dns"
)
// MultiRecorder is a type of ResponseWriter that captures all messages written to it.
type MultiRecorder struct {
Len int
Msgs []*dns.Msg
Start time.Time
dns.ResponseWriter
}
// NewMultiRecorder makes and returns a new MultiRecorder.
func NewMultiRecorder(w dns.ResponseWriter) *MultiRecorder {
return &MultiRecorder{
ResponseWriter: w,
Msgs: make([]*dns.Msg, 0),
Start: time.Now(),
}
}
// WriteMsg records the message and its length written to it and call the
// underlying ResponseWriter's WriteMsg method.
func (r *MultiRecorder) WriteMsg(res *dns.Msg) error {
r.Len += res.Len()
r.Msgs = append(r.Msgs, res)
return r.ResponseWriter.WriteMsg(res)
}
// Write is a wrapper that records the length of the messages that get written to it.
func (r *MultiRecorder) Write(buf []byte) (int, error) {
n, err := r.ResponseWriter.Write(buf)
if err == nil {
r.Len += n
}
return n, err
}