Codebase list golang-github-dnstap-golang-dnstap / debian/0.2.0-2 dnstap / fileoutput.go
debian/0.2.0-2

Tree @debian/0.2.0-2 (Download .tar.gz)

fileoutput.go @debian/0.2.0-2raw · history · blame

/*
 * 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/dnstap/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) {
	if formatter == nil {
		if filename == "-" || filename == "" {
			o = dnstap.NewTextOutput(os.Stdout, dnstap.TextFormat)
			return
		}
		o, err = dnstap.NewFrameStreamOutputFromFilename(filename)
	} else {
		if filename == "-" || filename == "" {
			if doAppend {
				return nil, errors.New("cannot append to stdout (-)")
			}
			o = dnstap.NewTextOutput(os.Stdout, formatter)
			return
		}
		o, err = dnstap.NewTextOutputFromFilename(filename, formatter, doAppend)
	}
	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)
		}
	}
}