mirror of
https://github.com/coredns/coredns.git
synced 2025-12-17 07:45:13 -05:00
86
middleware/dnstap/out/socket.go
Normal file
86
middleware/dnstap/out/socket.go
Normal file
@@ -0,0 +1,86 @@
|
||||
package out
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"net"
|
||||
|
||||
fs "github.com/farsightsec/golang-framestream"
|
||||
)
|
||||
|
||||
// Socket is a Frame Streams encoder over a UNIX socket.
|
||||
type Socket struct {
|
||||
path string
|
||||
enc *fs.Encoder
|
||||
conn net.Conn
|
||||
err error
|
||||
}
|
||||
|
||||
func openSocket(s *Socket) error {
|
||||
conn, err := net.Dial("unix", s.path)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
s.conn = conn
|
||||
|
||||
enc, err := fs.NewEncoder(conn, &fs.EncoderOptions{
|
||||
ContentType: []byte("protobuf:dnstap.Dnstap"),
|
||||
Bidirectional: true,
|
||||
})
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
s.enc = enc
|
||||
|
||||
s.err = nil
|
||||
return nil
|
||||
}
|
||||
|
||||
// NewSocket will always return a new Socket.
|
||||
// err if nothing is listening to it, it will attempt to reconnect on the next Write.
|
||||
func NewSocket(path string) (s *Socket, err error) {
|
||||
s = &Socket{path: path}
|
||||
if err = openSocket(s); err != nil {
|
||||
err = fmt.Errorf("open socket: %s", err)
|
||||
s.err = err
|
||||
return
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
// Write a single Frame Streams frame.
|
||||
func (s *Socket) Write(frame []byte) (int, error) {
|
||||
if s.err != nil {
|
||||
// is the dnstap tool listening?
|
||||
if err := openSocket(s); err != nil {
|
||||
return 0, fmt.Errorf("open socket: %s", err)
|
||||
}
|
||||
}
|
||||
n, err := s.enc.Write(frame)
|
||||
if err != nil {
|
||||
// the dnstap command line tool is down
|
||||
s.conn.Close()
|
||||
s.err = err
|
||||
return 0, err
|
||||
}
|
||||
return n, nil
|
||||
|
||||
}
|
||||
|
||||
// Close the socket and flush the remaining frames.
|
||||
func (s *Socket) Close() error {
|
||||
if s.err != nil {
|
||||
// nothing to close
|
||||
return nil
|
||||
}
|
||||
|
||||
defer s.conn.Close()
|
||||
|
||||
if err := s.enc.Flush(); err != nil {
|
||||
return fmt.Errorf("flush: %s", err)
|
||||
}
|
||||
if err := s.enc.Close(); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
94
middleware/dnstap/out/socket_test.go
Normal file
94
middleware/dnstap/out/socket_test.go
Normal file
@@ -0,0 +1,94 @@
|
||||
package out
|
||||
|
||||
import (
|
||||
"net"
|
||||
"testing"
|
||||
|
||||
fs "github.com/farsightsec/golang-framestream"
|
||||
)
|
||||
|
||||
func acceptOne(t *testing.T, l net.Listener) {
|
||||
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
|
||||
}
|
||||
|
||||
if _, err := dec.Decode(); err != nil {
|
||||
t.Errorf("server decode: %s", err)
|
||||
}
|
||||
|
||||
if err := server.Close(); err != nil {
|
||||
t.Error(err)
|
||||
}
|
||||
}
|
||||
func sendOne(socket *Socket) error {
|
||||
if _, err := socket.Write([]byte("frame")); err != nil {
|
||||
return err
|
||||
}
|
||||
if err := socket.enc.Flush(); err != nil {
|
||||
// Would happen during Write in real life.
|
||||
socket.conn.Close()
|
||||
socket.err = err
|
||||
return err
|
||||
}
|
||||
return nil
|
||||
}
|
||||
func TestSocket(t *testing.T) {
|
||||
socket, err := NewSocket("dnstap.sock")
|
||||
if err == nil {
|
||||
t.Fatal("new socket: not listening but no error")
|
||||
return
|
||||
}
|
||||
|
||||
if err := sendOne(socket); err == nil {
|
||||
t.Fatal("not listening but no error")
|
||||
return
|
||||
}
|
||||
|
||||
l, err := net.Listen("unix", "dnstap.sock")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
return
|
||||
}
|
||||
|
||||
wait := make(chan bool)
|
||||
go func() {
|
||||
acceptOne(t, l)
|
||||
wait <- true
|
||||
}()
|
||||
|
||||
if err := sendOne(socket); err != nil {
|
||||
t.Fatalf("send one: %s", err)
|
||||
return
|
||||
}
|
||||
|
||||
<-wait
|
||||
if err := sendOne(socket); err == nil {
|
||||
panic("must fail")
|
||||
}
|
||||
|
||||
go func() {
|
||||
acceptOne(t, l)
|
||||
wait <- true
|
||||
}()
|
||||
|
||||
if err := sendOne(socket); err != nil {
|
||||
t.Fatalf("send one: %s", err)
|
||||
return
|
||||
}
|
||||
|
||||
<-wait
|
||||
if err := l.Close(); err != nil {
|
||||
t.Error(err)
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user