From 94ae8b15e19184a36f04246e14dc13d78817d75a Mon Sep 17 00:00:00 2001
From: Cedric Staub <cs@squareup.com>
Date: Thu, 7 Mar 2019 13:32:45 -0800
Subject: [PATCH] Add --quiet=all option

---
 main.go        | 20 +++++++++++++++++---
 main_test.go   | 12 +++++++++++-
 proxy/proxy.go |  2 +-
 3 files changed, 29 insertions(+), 5 deletions(-)

diff --git a/main.go b/main.go
index e1e462316ce..129f352549e 100644
--- a/main.go
+++ b/main.go
@@ -20,6 +20,7 @@ import (
 	"crypto/tls"
 	"errors"
 	"fmt"
+	"io/ioutil"
 	"log"
 	"net"
 	"net/http"
@@ -112,7 +113,7 @@ var (
 	// Status & logging
 	statusAddress = app.Flag("status", "Enable serving /_status and /_metrics on given HOST:PORT (or unix:SOCKET).").PlaceHolder("ADDR").String()
 	enableProf    = app.Flag("enable-pprof", "Enable serving /debug/pprof endpoints alongside /_status (for profiling).").Bool()
-	quiet         = app.Flag("quiet", "Silence selected log messages (can be 'conns', 'conn-errs' and 'handshake-errs')").Default("").Enums("", "conns", "handshake-errs", "conn-errs")
+	quiet         = app.Flag("quiet", "Silence log messages (can be all, conns, conn-errs, handshake-errs; repeat flag for more than one)").Default("").Enums("", "all", "conns", "handshake-errs", "conn-errs")
 )
 
 func init() {
@@ -157,10 +158,17 @@ type Dialer interface {
 // Global logger instance
 var logger = log.New(os.Stderr, "", log.LstdFlags|log.Lmicroseconds)
 
-func initLogger(syslog bool) (err error) {
+func initLogger(syslog bool, flags []string) (err error) {
 	// If user has indicated request for syslog, override default stderr
 	// logger with a syslog one instead. This can fail, e.g. in containers
 	// that don't have syslog available.
+	for _, flag := range flags {
+		if flag == "all" {
+			// If --quiet=all if passed, disable all logging
+			logger = log.New(ioutil.Discard, "", 0)
+			return
+		}
+	}
 	if syslog {
 		var syslogWriter gsyslog.Syslogger
 		syslogWriter, err = gsyslog.NewLogger(gsyslog.LOG_INFO, "DAEMON", "")
@@ -298,7 +306,7 @@ func run(args []string) error {
 	command := kingpin.MustParse(app.Parse(args))
 
 	// Logger
-	err := initLogger(useSyslog())
+	err := initLogger(useSyslog(), *quiet)
 	if err != nil {
 		fmt.Fprintf(os.Stderr, "error initializing logger: %s\n", err)
 		os.Exit(1)
@@ -681,11 +689,17 @@ func proxyLoggerFlags(flags []string) int {
 	out := proxy.LogEverything
 	for _, flag := range flags {
 		switch flag {
+		case "all":
+			// Disable all proxy logs
+			out = 0
 		case "conns":
+			// Disable connection logs
 			out = out & ^proxy.LogConnections
 		case "conn-errs":
+			// Disable connection errors logs
 			out = out & ^proxy.LogConnectionErrors
 		case "handshake-errs":
+			// Disable handshake error logs
 			out = out & ^proxy.LogHandshakeErrors
 		}
 	}
diff --git a/main_test.go b/main_test.go
index a7c1c983208..8fa8e31ec80 100644
--- a/main_test.go
+++ b/main_test.go
@@ -82,9 +82,19 @@ func TestIntegrationMain(t *testing.T) {
 	}
 }
 
+func TestInitLoggerQuiet(t *testing.T) {
+	originalLogger := logger
+	err := initLogger(false, []string{"all"})
+	assert.Nil(t, err)
+
+	updatedLogger := logger
+	assert.NotEqual(t, originalLogger, updatedLogger, "should have updated logger object")
+	assert.NotNil(t, logger, "logger should never be nil after init")
+}
+
 func TestInitLoggerSyslog(t *testing.T) {
 	originalLogger := logger
-	err := initLogger(true)
+	err := initLogger(true, []string{})
 	updatedLogger := logger
 	if err != nil {
 		// Tests running in containers often don't have access to syslog,
diff --git a/proxy/proxy.go b/proxy/proxy.go
index 16151c2808c..00b1268c4ce 100644
--- a/proxy/proxy.go
+++ b/proxy/proxy.go
@@ -243,8 +243,8 @@ func (p *Proxy) fuse(client, backend net.Conn) {
 
 // Copy data between two connections
 func (p *Proxy) copyData(dst net.Conn, src net.Conn) {
-	defer src.Close()
 	defer dst.Close()
+	defer src.Close()
 
 	_, err := io.Copy(dst, src)