From 14342d7cdecef5c350a0161205dc9f2261f51e4a Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Micha=C5=82=20Matczuk?= <michal@scylladb.com>
Date: Tue, 28 Nov 2017 22:04:47 +0100
Subject: [PATCH] keepalive: server listen keepalive fix plus log message fixes

---
 server.go   | 29 +++++++++++++++--------------
 tcpproxy.go |  6 +++---
 2 files changed, 18 insertions(+), 17 deletions(-)

diff --git a/server.go b/server.go
index a61ba01..b3f7bd6 100644
--- a/server.go
+++ b/server.go
@@ -149,7 +149,7 @@ func (s *Server) Start() {
 
 			s.logger.Log(
 				"level", 0,
-				"msg", "accept control connection failed",
+				"msg", "accept of control connection failed",
 				"addr", addr,
 				"err", err,
 			)
@@ -158,8 +158,8 @@ func (s *Server) Start() {
 
 		if err := keepAlive(conn); err != nil {
 			s.logger.Log(
-				"level", 1,
-				"msg", "could not enable TCP keepalive for control connection",
+				"level", 0,
+				"msg", "TCP keepalive for control connection failed",
 				"addr", addr,
 				"err", err,
 			)
@@ -440,19 +440,9 @@ func (s *Server) listen(l net.Listener, identifier id.ID) {
 				return
 			}
 
-			if err = keepAlive(conn); err != nil {
-				s.logger.Log(
-					"level", 1,
-					"msg", "could not enable TCP keepalive for tunnel connection",
-					"identifier", identifier,
-					"addr", addr,
-					"err", err,
-				)
-			}
-
 			s.logger.Log(
 				"level", 0,
-				"msg", "accept connection failed",
+				"msg", "accept of connection failed",
 				"identifier", identifier,
 				"addr", addr,
 				"err", err,
@@ -465,6 +455,17 @@ func (s *Server) listen(l net.Listener, identifier id.ID) {
 			ForwardedHost:  l.Addr().String(),
 			ForwardedProto: l.Addr().Network(),
 		}
+
+		if err := keepAlive(conn); err != nil {
+			s.logger.Log(
+				"level", 1,
+				"msg", "TCP keepalive for tunneled connection failed",
+				"identifier", identifier,
+				"ctrlMsg", msg,
+				"err", err,
+			)
+		}
+
 		go func() {
 			if err := s.proxyConn(identifier, conn, msg); err != nil {
 				s.logger.Log(
diff --git a/tcpproxy.go b/tcpproxy.go
index f887750..613de9b 100644
--- a/tcpproxy.go
+++ b/tcpproxy.go
@@ -91,10 +91,10 @@ func (p *TCPProxy) Proxy(w io.Writer, r io.ReadCloser, msg *proto.ControlMessage
 	}
 	defer local.Close()
 
-	if err = keepAlive(local); err != nil {
+	if err := keepAlive(local); err != nil {
 		p.logger.Log(
-			"level", 0,
-			"msg", "could not enable TCP keepalive for local tunnel connection",
+			"level", 1,
+			"msg", "TCP keepalive for tunneled connection failed",
 			"target", target,
 			"ctrlMsg", msg,
 			"err", err,