Push more flag handling to main instead of referring to flags in other modules

Author: csstaub

main.go

@@ -98,12 +98,13 @@ var exitFunc = os.Exit
 
 // Context groups listening context data together
 type Context struct {
-	watcher    chan bool
-	status     *statusHandler
-	statusHTTP *http.Server
-	dial       func() (net.Conn, error)
-	metrics    *sqmetrics.SquareMetrics
-	cert       *certificate
+	watcher         chan bool
+	status          *statusHandler
+	statusHTTP      *http.Server
+	shutdownTimeout time.Duration
+	dial            func() (net.Conn, error)
+	metrics         *sqmetrics.SquareMetrics
+	cert            *certificate
 }
 
 // Dialer is an interface for dialers (either net.Dialer, or http_dialer.HttpTunnel)
@@ -280,7 +281,7 @@ func run(args []string) error {
 		logger.Printf("using target address %s", *serverForwardAddress)
 
 		status := newStatusHandler(dial)
-		context := &Context{watcher, status, nil, dial, metrics, cert}
+		context := &Context{watcher, status, nil, *shutdownTimeout, dial, metrics, cert}
 
 		// Start listening
 		err = serverListen(context)
@@ -309,7 +310,7 @@ func run(args []string) error {
 		}
 
 		status := newStatusHandler(dial)
-		context := &Context{watcher, status, nil, dial, metrics, cert}
+		context := &Context{watcher, status, nil, *shutdownTimeout, dial, metrics, cert}
 
 		// Start listening
 		err = clientListen(context)
@@ -328,7 +329,7 @@ func run(args []string) error {
 // connections. This is useful for the purpose of replacing certificates
 // in-place without having to take downtime, e.g. if a certificate is expiring.
 func serverListen(context *Context) error {
-	config, err := buildConfig(*caBundlePath)
+	config, err := buildConfig(*enabledCipherSuites, *caBundlePath)
 	if err != nil {
 		logger.Printf("error trying to read CA bundle: %s", err)
 		return err
@@ -353,10 +354,11 @@ func serverListen(context *Context) error {
 	}
 
 	proxy := &proxy{
-		quit:     0,
-		listener: tls.NewListener(listener, config),
-		handlers: &sync.WaitGroup{},
-		dial:     context.dial,
+		quit:           0,
+		listener:       tls.NewListener(listener, config),
+		handlers:       &sync.WaitGroup{},
+		connectTimeout: *timeoutDuration,
+		dial:           context.dial,
 	}
 
 	if *statusAddress != "" {
@@ -399,10 +401,11 @@ func clientListen(context *Context) error {
 	}
 
 	proxy := &proxy{
-		quit:     0,
-		listener: listener,
-		handlers: &sync.WaitGroup{},
-		dial:     context.dial,
+		quit:           0,
+		listener:       listener,
+		handlers:       &sync.WaitGroup{},
+		connectTimeout: *timeoutDuration,
+		dial:           context.dial,
 	}
 
 	if *statusAddress != "" {
@@ -437,7 +440,7 @@ func (context *Context) serveStatus() error {
 		mux.Handle("/debug/pprof/trace", http.HandlerFunc(pprof.Trace))
 	}
 
-	config, err := buildConfig(*caBundlePath)
+	config, err := buildConfig(*enabledCipherSuites, *caBundlePath)
 	if err != nil {
 		return err
 	}
@@ -495,7 +498,7 @@ func serverBackendDialer() (func() (net.Conn, error), error) {
 
 // Get backend dialer function in client mode (connecting to a TLS port)
 func clientBackendDialer(cert *certificate, network, address, host string) (func() (net.Conn, error), error) {
-	config, err := buildConfig(*caBundlePath)
+	config, err := buildConfig(*enabledCipherSuites, *caBundlePath)
 	if err != nil {
 		return nil, err
 	}
@@ -522,7 +525,7 @@ func clientBackendDialer(cert *certificate, network, address, host string) (func
 		logger.Printf("using HTTP(S) CONNECT proxy %s", (*clientConnectProxy).String())
 
 		// Use HTTP CONNECT proxy to connect to target.
-		proxyConfig, err := buildConfig(*caBundlePath)
+		proxyConfig, err := buildConfig(*enabledCipherSuites, *caBundlePath)
 		if err != nil {
 			return nil, err
 		}

net.go

@@ -29,10 +29,11 @@ import (
 )
 
 type proxy struct {
-	quit     int32
-	listener net.Listener
-	handlers *sync.WaitGroup
-	dial     func() (net.Conn, error)
+	quit           int32
+	listener       net.Listener
+	handlers       *sync.WaitGroup
+	connectTimeout time.Duration
+	dial           func() (net.Conn, error)
 }
 
 var (
@@ -70,7 +71,7 @@ func (p *proxy) accept() {
 			defer conn.Close()
 			defer openCounter.Dec(1)
 
-			err := forceHandshake(conn)
+			err := forceHandshake(p.connectTimeout, conn)
 			if err != nil {
 				errorCounter.Inc(1)
 				logger.Printf("error on TLS handshake from %s: %s", conn.RemoteAddr(), err)
@@ -96,13 +97,13 @@ func (p *proxy) accept() {
 // Otherwise, unauthenticated clients would be able to open connections
 // and leave them hanging forever. Going through the handshake verifies
 // that clients have a valid client cert and are allowed to talk to us.
-func forceHandshake(conn net.Conn) error {
+func forceHandshake(timeout time.Duration, conn net.Conn) error {
 	if tlsConn, ok := conn.(*tls.Conn); ok {
 		startTime := time.Now()
 		defer handshakeTimer.UpdateSince(startTime)
 
 		// Set deadline to avoid blocking forever
-		err := tlsConn.SetDeadline(time.Now().Add(*timeoutDuration))
+		err := tlsConn.SetDeadline(time.Now().Add(timeout))
 		if err != nil {
 			return err
 		}

signals.go

@@ -58,7 +58,7 @@ func (context *Context) signalHandler(proxy *proxy, closeables []io.Closer) {
 				}
 
 				// Force-exit after timeout
-				time.AfterFunc(*shutdownTimeout, func() {
+				time.AfterFunc(context.shutdownTimeout, func() {
 					// Graceful shutdown timeout reached. If we can't drain connections
 					// to exit gracefully after this timeout, let's just exit.
 					logger.Printf("graceful shutdown timeout: forcing exit")

tls.go

@@ -225,7 +225,7 @@ func dialWithDialer(dialer Dialer, timeout time.Duration, network, addr string,
 }
 
 // buildConfig reads command-line options and builds a tls.Config
-func buildConfig(caBundlePath string) (*tls.Config, error) {
+func buildConfig(enabledCipherSuites string, caBundlePath string) (*tls.Config, error) {
 	ca, err := caBundle(caBundlePath)
 	if err != nil {
 		return nil, err
@@ -236,7 +236,7 @@ func buildConfig(caBundlePath string) (*tls.Config, error) {
 	// * We list AES-128 ahead of AES-256 for performance reasons.
 
 	suites := []uint16{}
-	for _, suite := range strings.Split(*enabledCipherSuites, ",") {
+	for _, suite := range strings.Split(enabledCipherSuites, ",") {
 		ciphers, ok := cipherSuites[strings.TrimSpace(suite)]
 		if !ok {
 			return nil, fmt.Errorf("invalid cipher suite '%s' selected", suite)

tls_test.go

@@ -180,18 +180,16 @@ func TestBuildConfig(t *testing.T) {
 	defer os.Remove(tmpCaBundle.Name())
 	defer os.Remove(tmpKeystoreNoPrivKey.Name())
 
-	*enabledCipherSuites = ""
-	conf, err := buildConfig(tmpCaBundle.Name())
+	conf, err := buildConfig("", tmpCaBundle.Name())
 	assert.NotNil(t, err, "should fail to build config with no cipher suites")
 
-	*enabledCipherSuites = "AES,CHACHA"
-	conf, err = buildConfig(tmpCaBundle.Name())
+	conf, err = buildConfig("AES,CHACHA", tmpCaBundle.Name())
 	assert.Nil(t, err, "should be able to build TLS config")
 	assert.NotNil(t, conf.RootCAs, "config must have CA certs")
 	assert.NotNil(t, conf.ClientCAs, "config must have CA certs")
 	assert.True(t, conf.MinVersion == tls.VersionTLS12, "must have correct TLS min version")
 
-	conf, err = buildConfig("does-not-exist")
+	conf, err = buildConfig("AES", "does-not-exist")
 	assert.Nil(t, conf, "conf with invalid params should be nil")
 	assert.NotNil(t, err, "should reject invalid CA cert bundle")
 
@@ -222,21 +220,17 @@ func TestCipherSuitePreference(t *testing.T) {
 	tmpCaBundle.Sync()
 	defer os.Remove(tmpCaBundle.Name())
 
-	*enabledCipherSuites = "XYZ"
-	conf, err := buildConfig(tmpCaBundle.Name())
+	conf, err := buildConfig("XYZ", tmpCaBundle.Name())
 	assert.NotNil(t, err, "should not be able to build TLS config with invalid cipher suite option")
 
-	*enabledCipherSuites = ""
-	conf, err = buildConfig(tmpCaBundle.Name())
+	conf, err = buildConfig("", tmpCaBundle.Name())
 	assert.NotNil(t, err, "should not be able to build TLS config wihout cipher suite selection")
 
-	*enabledCipherSuites = "CHACHA,AES"
-	conf, err = buildConfig(tmpCaBundle.Name())
+	conf, err = buildConfig("CHACHA,AES", tmpCaBundle.Name())
 	assert.Nil(t, err, "should be able to build TLS config")
 	assert.True(t, conf.CipherSuites[0] == tls.TLS_ECDHE_ECDSA_WITH_CHACHA20_POLY1305, "expecting ChaCha20")
 
-	*enabledCipherSuites = "AES,CHACHA"
-	conf, err = buildConfig(tmpCaBundle.Name())
+	conf, err = buildConfig("AES,CHACHA", tmpCaBundle.Name())
 	assert.Nil(t, err, "should be able to build TLS config")
 	assert.True(t, conf.CipherSuites[0] == tls.TLS_ECDHE_ECDSA_WITH_AES_128_GCM_SHA256, "expecting AES")
 }
@@ -262,7 +256,7 @@ func TestBuildConfigSystemRoots(t *testing.T) {
 		t.SkipNow()
 		return
 	}
-	conf, err := buildConfig("")
+	conf, err := buildConfig("AES", "")
 	assert.Nil(t, err, "should be able to build TLS config")
 	assert.NotNil(t, conf.RootCAs, "config must have CA certs")
 	assert.NotNil(t, conf.ClientCAs, "config must have CA certs")