mirror of
https://github.com/coredns/coredns.git
synced 2026-01-03 07:21:18 -05:00
Dnstap plugin refactoring (#1257)
This commit is contained in:
Uladzimir Trehubenka
committed by
Miek Gieben
Miek Gieben
parent
06006fac56
commit
6d6e1357b9
@@ -2,67 +2,100 @@ package dnstapio
|
||||
|
||||
import (
|
||||
"log"
|
||||
"net"
|
||||
"time"
|
||||
|
||||
tap "github.com/dnstap/golang-dnstap"
|
||||
fs "github.com/farsightsec/golang-framestream"
|
||||
"github.com/golang/protobuf/proto"
|
||||
)
|
||||
|
||||
// DnstapIO wraps the dnstap I/O routine.
|
||||
type DnstapIO struct {
|
||||
protocol Protocol
|
||||
queue chan tap.Dnstap
|
||||
stop chan bool
|
||||
const (
|
||||
tcpTimeout = 4 * time.Second
|
||||
flushTimeout = 1 * time.Second
|
||||
queueSize = 1000
|
||||
)
|
||||
|
||||
type dnstapIO struct {
|
||||
enc *fs.Encoder
|
||||
conn net.Conn
|
||||
queue chan tap.Dnstap
|
||||
}
|
||||
|
||||
// Protocol is either `out.TCP` or `out.Socket`.
|
||||
type Protocol interface {
|
||||
// Write takes a single frame at once.
|
||||
Write([]byte) (int, error)
|
||||
|
||||
Close() error
|
||||
// New returns a new and initialized DnstapIO.
|
||||
func New() DnstapIO {
|
||||
return &dnstapIO{queue: make(chan tap.Dnstap, queueSize)}
|
||||
}
|
||||
|
||||
// New dnstap I/O routine from Protocol.
|
||||
func New(w Protocol) *DnstapIO {
|
||||
dio := DnstapIO{}
|
||||
dio.protocol = w
|
||||
dio.queue = make(chan tap.Dnstap, 10)
|
||||
dio.stop = make(chan bool)
|
||||
// DnstapIO interface
|
||||
type DnstapIO interface {
|
||||
Connect(endpoint string, socket bool) error
|
||||
Dnstap(payload tap.Dnstap)
|
||||
Close()
|
||||
}
|
||||
|
||||
// Connect connects to the dnstop endpoint.
|
||||
func (dio *dnstapIO) Connect(endpoint string, socket bool) error {
|
||||
var err error
|
||||
if socket {
|
||||
dio.conn, err = net.Dial("unix", endpoint)
|
||||
} else {
|
||||
dio.conn, err = net.DialTimeout("tcp", endpoint, tcpTimeout)
|
||||
}
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
dio.enc, err = fs.NewEncoder(dio.conn, &fs.EncoderOptions{
|
||||
ContentType: []byte("protobuf:dnstap.Dnstap"),
|
||||
Bidirectional: true,
|
||||
})
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
go dio.serve()
|
||||
return &dio
|
||||
return nil
|
||||
}
|
||||
|
||||
// Dnstap enqueues the payload for log.
|
||||
func (dio *DnstapIO) Dnstap(payload tap.Dnstap) {
|
||||
func (dio *dnstapIO) Dnstap(payload tap.Dnstap) {
|
||||
select {
|
||||
case dio.queue <- payload:
|
||||
default:
|
||||
log.Println("[WARN] Dnstap payload dropped.")
|
||||
}
|
||||
}
|
||||
|
||||
func (dio *DnstapIO) serve() {
|
||||
for {
|
||||
select {
|
||||
case payload := <-dio.queue:
|
||||
frame, err := proto.Marshal(&payload)
|
||||
if err == nil {
|
||||
dio.protocol.Write(frame)
|
||||
} else {
|
||||
log.Printf("[ERROR] Invalid dnstap payload dropped: %s\n", err)
|
||||
}
|
||||
case <-dio.stop:
|
||||
close(dio.queue)
|
||||
dio.stop <- true
|
||||
return
|
||||
}
|
||||
log.Printf("[ERROR] Dnstap payload dropped")
|
||||
}
|
||||
}
|
||||
|
||||
// Close waits until the I/O routine is finished to return.
|
||||
func (dio DnstapIO) Close() error {
|
||||
dio.stop <- true
|
||||
<-dio.stop
|
||||
close(dio.stop)
|
||||
return dio.protocol.Close()
|
||||
func (dio *dnstapIO) Close() {
|
||||
close(dio.queue)
|
||||
}
|
||||
|
||||
func (dio *dnstapIO) serve() {
|
||||
timeout := time.After(flushTimeout)
|
||||
for {
|
||||
select {
|
||||
case payload, ok := <-dio.queue:
|
||||
if !ok {
|
||||
dio.enc.Close()
|
||||
dio.conn.Close()
|
||||
return
|
||||
}
|
||||
frame, err := proto.Marshal(&payload)
|
||||
if err != nil {
|
||||
log.Printf("[ERROR] Invalid dnstap payload dropped: %s", err)
|
||||
continue
|
||||
}
|
||||
_, err = dio.enc.Write(frame)
|
||||
if err != nil {
|
||||
log.Printf("[ERROR] Cannot write dnstap payload: %s", err)
|
||||
continue
|
||||
}
|
||||
case <-timeout:
|
||||
err := dio.enc.Flush()
|
||||
if err != nil {
|
||||
log.Printf("[ERROR] Cannot flush dnstap payloads: %s", err)
|
||||
}
|
||||
timeout = time.After(flushTimeout)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,78 +1,155 @@
|
||||
package dnstapio
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"io/ioutil"
|
||||
"log"
|
||||
"net"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
tap "github.com/dnstap/golang-dnstap"
|
||||
fs "github.com/farsightsec/golang-framestream"
|
||||
)
|
||||
|
||||
func init() {
|
||||
log.SetOutput(ioutil.Discard)
|
||||
func accept(t *testing.T, l net.Listener, count int) {
|
||||
server, err := l.Accept()
|
||||
if err != nil {
|
||||
t.Fatalf("server accept: %s", err)
|
||||
return
|
||||
}
|
||||
|
||||
dec, err := fs.NewDecoder(server, &fs.DecoderOptions{
|
||||
ContentType: []byte("protobuf:dnstap.Dnstap"),
|
||||
Bidirectional: true,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("server decoder: %s", err)
|
||||
return
|
||||
}
|
||||
|
||||
for i := 0; i < count; i++ {
|
||||
if _, err := dec.Decode(); err != nil {
|
||||
t.Errorf("server decode: %s", err)
|
||||
}
|
||||
}
|
||||
|
||||
if err := server.Close(); err != nil {
|
||||
t.Error(err)
|
||||
}
|
||||
}
|
||||
|
||||
type buf struct {
|
||||
*bytes.Buffer
|
||||
cost time.Duration
|
||||
const endpointTCP = "localhost:0"
|
||||
|
||||
func TestTCP(t *testing.T) {
|
||||
dio := New()
|
||||
|
||||
err := dio.Connect(endpointTCP, false)
|
||||
if err == nil {
|
||||
t.Fatal("Not listening but no error")
|
||||
}
|
||||
|
||||
// Start TCP listener
|
||||
l, err := net.Listen("tcp", endpointTCP)
|
||||
if err != nil {
|
||||
t.Fatalf("Cannot start listener: %s", err)
|
||||
}
|
||||
defer l.Close()
|
||||
|
||||
var wg sync.WaitGroup
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
accept(t, l, 1)
|
||||
wg.Done()
|
||||
}()
|
||||
|
||||
err = dio.Connect(l.Addr().String(), false)
|
||||
if err != nil {
|
||||
t.Fatalf("Cannot connect to listener: %s", err)
|
||||
}
|
||||
|
||||
msg := tap.Dnstap_MESSAGE
|
||||
dio.Dnstap(tap.Dnstap{Type: &msg})
|
||||
|
||||
wg.Wait()
|
||||
|
||||
dio.Close()
|
||||
}
|
||||
|
||||
func (b buf) Write(frame []byte) (int, error) {
|
||||
time.Sleep(b.cost)
|
||||
return b.Buffer.Write(frame)
|
||||
}
|
||||
const endpointSocket = "dnstap.sock"
|
||||
|
||||
func (b buf) Close() error {
|
||||
return nil
|
||||
func TestSocket(t *testing.T) {
|
||||
dio := New()
|
||||
|
||||
err := dio.Connect(endpointSocket, true)
|
||||
if err == nil {
|
||||
t.Fatal("Not listening but no error")
|
||||
}
|
||||
|
||||
// Start Socket listener
|
||||
l, err := net.Listen("unix", endpointSocket)
|
||||
if err != nil {
|
||||
t.Fatalf("Cannot start listener: %s", err)
|
||||
}
|
||||
defer l.Close()
|
||||
|
||||
var wg sync.WaitGroup
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
accept(t, l, 1)
|
||||
wg.Done()
|
||||
}()
|
||||
|
||||
err = dio.Connect(endpointSocket, true)
|
||||
if err != nil {
|
||||
t.Fatalf("Cannot connect to listener: %s", err)
|
||||
}
|
||||
|
||||
msg := tap.Dnstap_MESSAGE
|
||||
dio.Dnstap(tap.Dnstap{Type: &msg})
|
||||
|
||||
wg.Wait()
|
||||
|
||||
dio.Close()
|
||||
}
|
||||
|
||||
func TestRace(t *testing.T) {
|
||||
b := buf{&bytes.Buffer{}, 100 * time.Millisecond}
|
||||
dio := New(b)
|
||||
wg := &sync.WaitGroup{}
|
||||
wg.Add(10)
|
||||
for i := 0; i < 10; i++ {
|
||||
timeout := time.After(time.Second)
|
||||
go func() {
|
||||
for {
|
||||
select {
|
||||
case <-timeout:
|
||||
wg.Done()
|
||||
return
|
||||
default:
|
||||
time.Sleep(50 * time.Millisecond)
|
||||
t := tap.Dnstap_MESSAGE
|
||||
dio.Dnstap(tap.Dnstap{Type: &t})
|
||||
}
|
||||
}
|
||||
}()
|
||||
}
|
||||
wg.Wait()
|
||||
}
|
||||
count := 10
|
||||
dio := New()
|
||||
|
||||
func TestClose(t *testing.T) {
|
||||
done := make(chan bool)
|
||||
var dio *DnstapIO
|
||||
go func() {
|
||||
b := buf{&bytes.Buffer{}, 0}
|
||||
dio = New(b)
|
||||
dio.Close()
|
||||
close(done)
|
||||
}()
|
||||
select {
|
||||
case <-done:
|
||||
case <-time.After(time.Second):
|
||||
t.Fatal("Not closing.")
|
||||
err := dio.Connect(endpointTCP, false)
|
||||
if err == nil {
|
||||
t.Fatal("Not listening but no error")
|
||||
}
|
||||
func() {
|
||||
defer func() {
|
||||
if err := recover(); err == nil {
|
||||
t.Fatal("Send on closed channel.")
|
||||
}
|
||||
}()
|
||||
dio.Dnstap(tap.Dnstap{})
|
||||
|
||||
// Start TCP listener
|
||||
l, err := net.Listen("tcp", endpointTCP)
|
||||
if err != nil {
|
||||
t.Fatalf("Cannot start listener: %s", err)
|
||||
}
|
||||
defer l.Close()
|
||||
|
||||
var wg sync.WaitGroup
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
accept(t, l, count)
|
||||
wg.Done()
|
||||
}()
|
||||
|
||||
err = dio.Connect(l.Addr().String(), false)
|
||||
if err != nil {
|
||||
t.Fatalf("Cannot connect to listener: %s", err)
|
||||
}
|
||||
|
||||
msg := tap.Dnstap_MESSAGE
|
||||
wg.Add(count)
|
||||
for i := 0; i < count; i++ {
|
||||
go func(i byte) {
|
||||
time.Sleep(50 * time.Millisecond)
|
||||
dio.Dnstap(tap.Dnstap{Type: &msg, Extra: []byte{i}})
|
||||
wg.Done()
|
||||
}(byte(i))
|
||||
}
|
||||
|
||||
wg.Wait()
|
||||
|
||||
dio.Close()
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user