Move log mode to Logger

1 files changed, 18 insertions, 0 deletions

diff --git a/log/logger.go b/log/logger.go
index 7fa4001..1a5b2d0 100644
--- a/log/logger.go
+++ b/log/logger.go
@@ -10,9 +10,19 @@ import (
 	"github.com/mpolden/zdns/sql"
 )
 
+const (
+	// ModeDiscard disables logging of DNS requests.
+	ModeDiscard = iota
+	// ModeAll logs all DNS requests.
+	ModeAll
+	// ModeHijacked only logs hijacked DNS requests.
+	ModeHijacked
+)
+
 // Logger wraps a standard log.Logger and an optional log database.
 type Logger struct {
 	*log.Logger
+	mode  int
 	queue chan Entry
 	db    *sql.Client
 	wg    sync.WaitGroup
@@ -23,6 +33,7 @@ type Logger struct {
 // RecordOptions configures recording of DNS requests.
 type RecordOptions struct {
 	Database string
+	Mode     int
 	TTL      time.Duration
 }
 
@@ -47,6 +58,7 @@ func newLogger(w io.Writer, prefix string, options RecordOptions, interval time.
 		Logger: log.New(w, prefix, 0),
 		queue:  make(chan Entry, 100),
 		now:    time.Now,
+		mode:   options.Mode,
 	}
 	var err error
 	if options.Database != "" {
@@ -97,6 +109,12 @@ func (l *Logger) Record(remoteAddr net.IP, hijacked bool, qtype uint16, question
 	if l.db == nil {
 		return
 	}
+	if l.mode == ModeDiscard {
+		return
+	}
+	if l.mode == ModeHijacked && !hijacked {
+		return
+	}
 	l.queue <- Entry{
 		Time:       l.now(),
 		RemoteAddr: remoteAddr,