From 8fe504241d6e73f7188cec8841ce8d5cf03843b6 Mon Sep 17 00:00:00 2001
From: James Tucker <james@tailscale.com>
Date: Sun, 18 Feb 2024 18:56:43 -0800
Subject: [PATCH] net/ktimeout: add a package to set TCP user timeout

Setting a user timeout will be a more practical tuning knob for a number
of endpoints, this provides a way to set it.

Updates tailscale/corp#17587

Signed-off-by: James Tucker <james@tailscale.com>
---
 net/ktimeout/ktimeout.go            | 36 ++++++++++++++++++++++
 net/ktimeout/ktimeout_default.go    | 15 ++++++++++
 net/ktimeout/ktimeout_linux.go      | 15 ++++++++++
 net/ktimeout/ktimeout_linux_test.go | 46 +++++++++++++++++++++++++++++
 net/ktimeout/ktimeout_test.go       | 24 +++++++++++++++
 5 files changed, 136 insertions(+)
 create mode 100644 net/ktimeout/ktimeout.go
 create mode 100644 net/ktimeout/ktimeout_default.go
 create mode 100644 net/ktimeout/ktimeout_linux.go
 create mode 100644 net/ktimeout/ktimeout_linux_test.go
 create mode 100644 net/ktimeout/ktimeout_test.go

diff --git a/net/ktimeout/ktimeout.go b/net/ktimeout/ktimeout.go
new file mode 100644
index 000000000..7cd439143
--- /dev/null
+++ b/net/ktimeout/ktimeout.go
@@ -0,0 +1,36 @@
+// Copyright (c) Tailscale Inc & AUTHORS
+// SPDX-License-Identifier: BSD-3-Clause
+
+// Package ktimeout configures kernel TCP stack timeouts via the provided
+// control functions. Platform support varies; on unsupported platforms control
+// functions may be entirely no-ops.
+package ktimeout
+
+import (
+	"fmt"
+	"syscall"
+	"time"
+)
+
+// UserTimeout returns a control function that sets the TCP user timeout
+// (TCP_USER_TIMEOUT on linux). A user timeout specifies the maximum age of
+// unacknowledged data on the connection (either in buffer, or sent but not
+// acknowledged) before the connection is terminated. This timer has no effect
+// on limiting the lifetime of idle connections. This may be entirely local to
+// the network stack or may also apply RFC 5482 options to packets.
+func UserTimeout(timeout time.Duration) func(network, address string, c syscall.RawConn) error {
+	return func(network, address string, c syscall.RawConn) error {
+		switch network {
+		case "tcp", "tcp4", "tcp6":
+		default:
+			return fmt.Errorf("ktimeout.UserTimeout: unsupported network: %s", network)
+		}
+		var err error
+		if e := c.Control(func(fd uintptr) {
+			err = SetUserTimeout(fd, timeout)
+		}); e != nil {
+			return e
+		}
+		return err
+	}
+}
diff --git a/net/ktimeout/ktimeout_default.go b/net/ktimeout/ktimeout_default.go
new file mode 100644
index 000000000..f1b11661b
--- /dev/null
+++ b/net/ktimeout/ktimeout_default.go
@@ -0,0 +1,15 @@
+// Copyright (c) Tailscale Inc & AUTHORS
+// SPDX-License-Identifier: BSD-3-Clause
+
+//go:build !linux
+
+package ktimeout
+
+import (
+	"time"
+)
+
+// SetUserTimeout is a no-op on this platform.
+func SetUserTimeout(fd uintptr, timeout time.Duration) error {
+	return nil
+}
diff --git a/net/ktimeout/ktimeout_linux.go b/net/ktimeout/ktimeout_linux.go
new file mode 100644
index 000000000..84286b647
--- /dev/null
+++ b/net/ktimeout/ktimeout_linux.go
@@ -0,0 +1,15 @@
+// Copyright (c) Tailscale Inc & AUTHORS
+// SPDX-License-Identifier: BSD-3-Clause
+
+package ktimeout
+
+import (
+	"time"
+
+	"golang.org/x/sys/unix"
+)
+
+// SetUserTimeout sets the TCP_USER_TIMEOUT option on the given file descriptor.
+func SetUserTimeout(fd uintptr, timeout time.Duration) error {
+	return unix.SetsockoptInt(int(fd), unix.SOL_TCP, unix.TCP_USER_TIMEOUT, int(timeout/time.Millisecond))
+}
diff --git a/net/ktimeout/ktimeout_linux_test.go b/net/ktimeout/ktimeout_linux_test.go
new file mode 100644
index 000000000..a367bfd4a
--- /dev/null
+++ b/net/ktimeout/ktimeout_linux_test.go
@@ -0,0 +1,46 @@
+// Copyright (c) Tailscale Inc & AUTHORS
+// SPDX-License-Identifier: BSD-3-Clause
+
+package ktimeout
+
+import (
+	"net"
+	"testing"
+	"time"
+
+	"golang.org/x/net/nettest"
+	"golang.org/x/sys/unix"
+	"tailscale.com/util/must"
+)
+
+func TestSetUserTimeout(t *testing.T) {
+	l := must.Get(nettest.NewLocalListener("tcp"))
+	defer l.Close()
+
+	var err error
+	if e := must.Get(l.(*net.TCPListener).SyscallConn()).Control(func(fd uintptr) {
+		err = SetUserTimeout(fd, 0)
+	}); e != nil {
+		t.Fatal(e)
+	}
+	if err != nil {
+		t.Fatal(err)
+	}
+	v := must.Get(unix.GetsockoptInt(int(must.Get(l.(*net.TCPListener).File()).Fd()), unix.SOL_TCP, unix.TCP_USER_TIMEOUT))
+	if v != 0 {
+		t.Errorf("TCP_USER_TIMEOUT: got %v; want 0", v)
+	}
+
+	if e := must.Get(l.(*net.TCPListener).SyscallConn()).Control(func(fd uintptr) {
+		err = SetUserTimeout(fd, 30*time.Second)
+	}); e != nil {
+		t.Fatal(e)
+	}
+	if err != nil {
+		t.Fatal(err)
+	}
+	v = must.Get(unix.GetsockoptInt(int(must.Get(l.(*net.TCPListener).File()).Fd()), unix.SOL_TCP, unix.TCP_USER_TIMEOUT))
+	if v != 30000 {
+		t.Errorf("TCP_USER_TIMEOUT: got %v; want 30000", v)
+	}
+}
diff --git a/net/ktimeout/ktimeout_test.go b/net/ktimeout/ktimeout_test.go
new file mode 100644
index 000000000..7befa3b1a
--- /dev/null
+++ b/net/ktimeout/ktimeout_test.go
@@ -0,0 +1,24 @@
+// Copyright (c) Tailscale Inc & AUTHORS
+// SPDX-License-Identifier: BSD-3-Clause
+
+package ktimeout
+
+import (
+	"context"
+	"fmt"
+	"net"
+	"time"
+)
+
+func ExampleUserTimeout() {
+	lc := net.ListenConfig{
+		Control: UserTimeout(30 * time.Second),
+	}
+	l, err := lc.Listen(context.TODO(), "tcp", "127.0.0.1:0")
+	if err != nil {
+		fmt.Printf("error: %v", err)
+		return
+	}
+	l.Close()
+	// Output:
+}