From dc4fcacd8c0be367159b36be3a087ce613899c27 Mon Sep 17 00:00:00 2001
From: Andrey Smirnov <smirnov.andrey@gmail.com>
Date: Thu, 15 Oct 2020 16:34:41 +0300
Subject: [PATCH] Add support for TCP_USER_TIMEOUT setting

See https://blog.cloudflare.com/when-tcp-sockets-refuse-to-die/ for
technical details.

Signed-off-by: Andrey Smirnov <smirnov.andrey@gmail.com>
---
 tcp_linux.go     | 74 ++++++++++++++++++++++++++++++++++++++++++++++++
 tcp_other.go     | 37 ++++++++++++++++++++++++
 tcpproxy.go      |  9 ++++++
 tcpproxy_test.go |  1 +
 4 files changed, 121 insertions(+)
 create mode 100644 tcp_linux.go
 create mode 100644 tcp_other.go

diff --git a/tcp_linux.go b/tcp_linux.go
new file mode 100644
index 0000000..f54ef04
--- /dev/null
+++ b/tcp_linux.go
@@ -0,0 +1,74 @@
+// +build linux !appengine
+
+/*
+ *
+ * Copyright 2018 gRPC authors.
+ *
+ * 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 tcpproxy
+
+import (
+	"fmt"
+	"net"
+	"syscall"
+	"time"
+
+	"golang.org/x/sys/unix"
+)
+
+// SetTCPUserTimeout sets the TCP user timeout on a connection's socket
+func SetTCPUserTimeout(conn net.Conn, timeout time.Duration) error {
+	tcpconn, ok := conn.(*net.TCPConn)
+	if !ok {
+		// not a TCP connection. exit early
+		return nil
+	}
+	rawConn, err := tcpconn.SyscallConn()
+	if err != nil {
+		return fmt.Errorf("error getting raw connection: %v", err)
+	}
+	err = rawConn.Control(func(fd uintptr) {
+		err = syscall.SetsockoptInt(int(fd), syscall.IPPROTO_TCP, unix.TCP_USER_TIMEOUT, int(timeout/time.Millisecond))
+	})
+	if err != nil {
+		return fmt.Errorf("error setting option on socket: %v", err)
+	}
+
+	return nil
+}
+
+// GetTCPUserTimeout gets the TCP user timeout on a connection's socket
+func GetTCPUserTimeout(conn net.Conn) (opt int, err error) {
+	tcpconn, ok := conn.(*net.TCPConn)
+	if !ok {
+		err = fmt.Errorf("conn is not *net.TCPConn. got %T", conn)
+		return
+	}
+	rawConn, err := tcpconn.SyscallConn()
+	if err != nil {
+		err = fmt.Errorf("error getting raw connection: %v", err)
+		return
+	}
+	err = rawConn.Control(func(fd uintptr) {
+		opt, err = syscall.GetsockoptInt(int(fd), syscall.IPPROTO_TCP, unix.TCP_USER_TIMEOUT)
+	})
+	if err != nil {
+		err = fmt.Errorf("error getting option on socket: %v", err)
+		return
+	}
+
+	return
+}
diff --git a/tcp_other.go b/tcp_other.go
new file mode 100644
index 0000000..b8ba680
--- /dev/null
+++ b/tcp_other.go
@@ -0,0 +1,37 @@
+// +build !linux appengine
+
+/*
+ *
+ * Copyright 2018 gRPC authors.
+ *
+ * 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 tcpproxy
+
+import (
+	"net"
+	"time"
+)
+
+// SetTCPUserTimeout is a no-op function under non-linux or appengine environments
+func SetTCPUserTimeout(conn net.Conn, timeout time.Duration) error {
+	return nil
+}
+
+// GetTCPUserTimeout is a no-op function under non-linux or appengine environments
+// a negative return value indicates the operation is not supported
+func GetTCPUserTimeout(conn net.Conn) (int, error) {
+	return -1, nil
+}
diff --git a/tcpproxy.go b/tcpproxy.go
index 9826d94..87dea37 100644
--- a/tcpproxy.go
+++ b/tcpproxy.go
@@ -318,6 +318,10 @@ type DialProxy struct {
 	// If negative, the timeout is disabled.
 	DialTimeout time.Duration
 
+	// TCPUserTimeout optionally specifies a TCP_USER_TIMEOUT (only on Linux).
+	// If zero, TCP_USER_TIMEOUT is not set.
+	TCPUserTimeout time.Duration
+
 	// DialContext optionally specifies an alternate dial function
 	// for TCP targets. If nil, the standard
 	// net.Dialer.DialContext method is used.
@@ -383,6 +387,11 @@ func (dp *DialProxy) HandleConn(src net.Conn) {
 		}
 	}
 
+	if dp.TCPUserTimeout > 0 {
+		SetTCPUserTimeout(src, dp.TCPUserTimeout)
+		SetTCPUserTimeout(dst, dp.TCPUserTimeout)
+	}
+
 	errc := make(chan error, 1)
 	go proxyCopy(errc, src, dst)
 	go proxyCopy(errc, dst, src)
diff --git a/tcpproxy_test.go b/tcpproxy_test.go
index 5d75cc3..0d81280 100644
--- a/tcpproxy_test.go
+++ b/tcpproxy_test.go
@@ -297,6 +297,7 @@ func TestProxyPROXYOut(t *testing.T) {
 	p.AddRoute(testFrontAddr, &DialProxy{
 		Addr:                 back.Addr().String(),
 		ProxyProtocolVersion: 1,
+		TCPUserTimeout:       time.Second,
 	})
 	if err := p.Start(); err != nil {
 		t.Fatal(err)