forked from dnstap/golang-dnstap
-
Notifications
You must be signed in to change notification settings - Fork 0
/
fileoutput.go
139 lines (129 loc) · 3.42 KB
/
fileoutput.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
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
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
/*
* Copyright (c) 2019 by Farsight Security, Inc.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package main
import (
"errors"
"fmt"
"os"
"os/signal"
"syscall"
dnstap "github.com/45hur/golang-dnstap"
)
// Output channel buffer size value from main dnstap package.
const outputChannelSize = 32
//
// A fileOutput implements a dnstap.Output which writes frames to a file
// and closes and reopens the file on SIGHUP.
//
// Data frames are written in binary fstrm format unless a text formatting
// function (dnstp.TextFormatFunc) is given or the filename is blank or "-".
// In the latter case, data is written in compact (quiet) text format unless
// an alternate text format is given on the assumption that stdout is a terminal.
//
type fileOutput struct {
formatter dnstap.TextFormatFunc
filename string
doAppend bool
output dnstap.Output
data chan []byte
done chan struct{}
}
func openOutputFile(filename string, formatter dnstap.TextFormatFunc, doAppend bool) (o dnstap.Output, err error) {
var fso *dnstap.FrameStreamOutput
var to *dnstap.TextOutput
if formatter == nil {
if filename == "-" || filename == "" {
to = dnstap.NewTextOutput(os.Stdout, dnstap.TextFormat)
to.SetLogger(logger)
return to, nil
}
fso, err = dnstap.NewFrameStreamOutputFromFilename(filename)
if err == nil {
fso.SetLogger(logger)
return fso, nil
}
} else {
if filename == "-" || filename == "" {
if doAppend {
return nil, errors.New("cannot append to stdout (-)")
}
to = dnstap.NewTextOutput(os.Stdout, formatter)
to.SetLogger(logger)
return to, nil
}
to, err = dnstap.NewTextOutputFromFilename(filename, formatter, doAppend)
if err == nil {
to.SetLogger(logger)
}
return to, nil
}
return
}
func newFileOutput(filename string, formatter dnstap.TextFormatFunc, doAppend bool) (*fileOutput, error) {
o, err := openOutputFile(filename, formatter, doAppend)
if err != nil {
return nil, err
}
return &fileOutput{
formatter: formatter,
filename: filename,
doAppend: doAppend,
output: o,
data: make(chan []byte, outputChannelSize),
done: make(chan struct{}),
}, nil
}
func (fo *fileOutput) GetOutputChannel() chan []byte {
return fo.data
}
func (fo *fileOutput) Close() {
close(fo.data)
<-fo.done
}
func (fo *fileOutput) RunOutputLoop() {
sigch := make(chan os.Signal, 1)
signal.Notify(sigch, os.Interrupt, syscall.SIGHUP)
o := fo.output
go o.RunOutputLoop()
defer func() {
o.Close()
close(fo.done)
}()
for {
select {
case b, ok := <-fo.data:
if !ok {
return
}
o.GetOutputChannel() <- b
case sig := <-sigch:
if sig == syscall.SIGHUP {
o.Close()
newo, err := openOutputFile(fo.filename, fo.formatter, fo.doAppend)
if err != nil {
fmt.Fprintf(os.Stderr,
"dnstap: Error: failed to reopen %s: %v\n",
fo.filename, err)
os.Exit(1)
}
o = newo
go o.RunOutputLoop()
continue
}
os.Exit(0)
}
}
}