Add CheckAddrWithOptions method with IPv4/IPv6 and socket mark support

Author: Lan-ce-lot

checker_linux.go

@@ -141,16 +141,25 @@ func (c *Checker) CheckAddr(addr string, timeout time.Duration) (err error) {
 
 // CheckAddrZeroLinger is like CheckAddr with an extra parameter indicating whether to enable zero linger.
 func (c *Checker) CheckAddrZeroLinger(addr string, timeout time.Duration, zeroLinger bool) error {
+	opts := DefaultOptions().WithTimeout(timeout).WithZeroLinger(zeroLinger)
+	return c.CheckAddrWithOptions(addr, opts)
+}
+
+// CheckAddrWithOptions performs a TCP check with given address and options.
+// A successful check will result in nil error.
+// ErrTimeout is returned if timeout.
+// Note: timeout includes domain resolving.
+func (c *Checker) CheckAddrWithOptions(addr string, opts Options) error {
 	// Set deadline
-	deadline := time.Now().Add(timeout)
+	deadline := time.Now().Add(opts.Timeout)
 
-	// Parse address
-	rAddr, family, err := parseSockAddr(addr)
+	// Parse address with specified network
+	rAddr, family, err := parseSockAddrWithNetwork(addr, opts.Network)
 	if err != nil {
 		return err
 	}
 	// Create socket with options set
-	fd, err := createSocketZeroLinger(family, zeroLinger)
+	fd, err := createSocketWithOptions(family, opts.ZeroLinger, opts.Mark)
 	if err != nil {
 		return err
 	}

checker_nonlinux.go

@@ -42,9 +42,17 @@ func (c *Checker) CheckAddr(addr string, timeout time.Duration) error {
 
 // CheckAddrZeroLinger is CheckerAddr with a zeroLinger parameter.
 func (c *Checker) CheckAddrZeroLinger(addr string, timeout time.Duration, zeroLinger bool) error {
-	conn, err := net.DialTimeout("tcp", addr, timeout)
+	opts := DefaultOptions().WithTimeout(timeout).WithZeroLinger(zeroLinger)
+	return c.CheckAddrWithOptions(addr, opts)
+}
+
+// CheckAddrWithOptions performs a TCP check with given address and options.
+// NOTE: zeroLinger is ignored on non-POSIX operating systems because
+// net.TCPConn.SetLinger is only implemented in src/net/sockopt_posix.go.
+func (c *Checker) CheckAddrWithOptions(addr string, opts Options) error {
+	conn, err := net.DialTimeout(opts.Network, addr, opts.Timeout)
 	if conn != nil {
-		if zeroLinger {
+		if opts.ZeroLinger {
 			// Simply ignore the error since this is a fake implementation.
 			_ = conn.(*net.TCPConn).SetLinger(0)
 		}

options.go

@@ -0,0 +1,62 @@
+package tcp
+
+import (
+	"time"
+)
+
+// Options contains configuration for TCP connectivity checks
+type Options struct {
+	// Timeout specifies the maximum duration for the check operation
+	Timeout time.Duration
+
+	// Network specifies the network type for address resolution and connection
+	// Supported values:
+	//   "tcp"  - Try IPv4 first, then IPv6 (default behavior)
+	//   "tcp4" - IPv4 only
+	//   "tcp6" - IPv6 only
+	Network string
+
+	// ZeroLinger indicates whether to set SO_LINGER with zero timeout
+	// This forces the connection to be reset immediately when closed
+	ZeroLinger bool
+
+	// Mark sets the SO_MARK socket option (Linux only)
+	// This is useful for traffic marking and routing policies
+	// Value of 0 means no mark is set
+	Mark int
+}
+
+// DefaultOptions returns Options with default values
+func DefaultOptions() Options {
+	return Options{
+		Timeout:    time.Second * 3,
+		Network:    "tcp",
+		ZeroLinger: true,
+		Mark:       0, // No mark by default
+	}
+}
+
+// WithTimeout sets the timeout for the operation
+func (o Options) WithTimeout(timeout time.Duration) Options {
+	o.Timeout = timeout
+	return o
+}
+
+// WithNetwork sets the network type (tcp, tcp4, tcp6)
+func (o Options) WithNetwork(network string) Options {
+	o.Network = network
+	return o
+}
+
+// WithZeroLinger sets the zero linger option
+func (o Options) WithZeroLinger(zeroLinger bool) Options {
+	o.ZeroLinger = zeroLinger
+	return o
+}
+
+// WithMark sets the SO_MARK socket option (Linux only)
+// This is useful for traffic marking and routing policies
+func (o Options) WithMark(mark int) Options {
+	o.Mark = mark
+	return o
+}

options_test.go

@@ -0,0 +1,147 @@
+package tcp
+
+import (
+	"context"
+	"testing"
+	"time"
+)
+
+// setupTestChecker creates a checker with CheckingLoop running and waits until ready
+func setupTestChecker(t *testing.T) (*Checker, context.CancelFunc) {
+	checker := NewChecker()
+	ctx, cancel := context.WithCancel(context.Background())
+	go func() {
+		_ = checker.CheckingLoop(ctx)
+	}()
+	<-checker.WaitReady()
+	return checker, cancel
+}
+
+// testWithChecker runs a test function with a prepared checker and test server
+func testWithChecker(t *testing.T, testFunc func(*testing.T, *Checker, string)) {
+	checker, cancel := setupTestChecker(t)
+	defer cancel()
+
+	testAddr, stopServer := StartTestServer()
+	defer stopServer()
+
+	testFunc(t, checker, testAddr)
+}
+
+func TestOptions(t *testing.T) {
+	// Test default options
+	opts := DefaultOptions()
+	if opts.Network != "tcp" {
+		t.Errorf("expected default network to be 'tcp', got %s", opts.Network)
+	}
+	if opts.Timeout != 3*time.Second {
+		t.Errorf("expected default timeout to be 3s, got %v", opts.Timeout)
+	}
+	if !opts.ZeroLinger {
+		t.Error("expected default ZeroLinger to be true")
+	}
+	if opts.Mark != 0 {
+		t.Errorf("expected default Mark to be 0, got %d", opts.Mark)
+	}
+
+	// Test fluent interface
+	customOpts := DefaultOptions().
+		WithTimeout(5 * time.Second).
+		WithNetwork("tcp6").
+		WithZeroLinger(false).
+		WithMark(100)
+
+	if customOpts.Network != "tcp6" {
+		t.Errorf("expected network to be 'tcp6', got %s", customOpts.Network)
+	}
+	if customOpts.Timeout != 5*time.Second {
+		t.Errorf("expected timeout to be 5s, got %v", customOpts.Timeout)
+	}
+	if customOpts.ZeroLinger {
+		t.Error("expected ZeroLinger to be false")
+	}
+	if customOpts.Mark != 100 {
+		t.Errorf("expected Mark to be 100, got %d", customOpts.Mark)
+	}
+}
+
+func TestCheckerWithOptions(t *testing.T) {
+	testWithChecker(t, func(t *testing.T, checker *Checker, testAddr string) {
+		// Test basic functionality with default options
+		opts := DefaultOptions().WithTimeout(2 * time.Second)
+		err := checker.CheckAddrWithOptions(testAddr, opts)
+		if err != nil {
+			t.Errorf("Connection to test server failed: %v", err)
+		}
+
+		// Test IPv4 specific
+		opts4 := DefaultOptions().WithTimeout(2 * time.Second).WithNetwork("tcp4")
+		err = checker.CheckAddrWithOptions(testAddr, opts4)
+		if err != nil {
+			t.Errorf("IPv4 connection to test server failed: %v", err)
+		}
+	})
+}
+
+func TestIPv6Support(t *testing.T) {
+	checker, cancel := setupTestChecker(t)
+	defer cancel()
+
+	// Try to start IPv6 server
+	testAddr6, stopServer6, err := StartTestServerIPv6()
+	if err != nil {
+		t.Skipf("Skipping IPv6 test: %v", err)
+		return
+	}
+	defer stopServer6()
+
+	// Test IPv6 connection with tcp6 network
+	opts6 := DefaultOptions().WithTimeout(2 * time.Second).WithNetwork("tcp6")
+	err = checker.CheckAddrWithOptions(testAddr6, opts6)
+	if err != nil {
+		t.Errorf("IPv6 connection to test server failed: %v", err)
+	}
+
+	// Test IPv6 connection with tcp network (should also work)
+	optsGeneral := DefaultOptions().WithTimeout(2 * time.Second).WithNetwork("tcp")
+	err = checker.CheckAddrWithOptions(testAddr6, optsGeneral)
+	if err != nil {
+		t.Errorf("General TCP connection to IPv6 test server failed: %v", err)
+	}
+
+	// Test that tcp4 fails on IPv6 address (should fail)
+	opts4 := DefaultOptions().WithTimeout(1 * time.Second).WithNetwork("tcp4")
+	err = checker.CheckAddrWithOptions(testAddr6, opts4)
+	if err == nil {
+		t.Error("Expected tcp4 connection to IPv6 address to fail, but it succeeded")
+	} else {
+		t.Logf("tcp4 connection to IPv6 address failed as expected: %v", err)
+	}
+}
+
+func TestMarkOption(t *testing.T) {
+	testWithChecker(t, func(t *testing.T, checker *Checker, testAddr string) {
+		// Test with mark set (only meaningful on Linux)
+		opts := DefaultOptions().WithTimeout(2 * time.Second).WithMark(42)
+		err := checker.CheckAddrWithOptions(testAddr, opts)
+		if err != nil {
+			t.Errorf("Connection with mark to test server failed: %v", err)
+		}
+	})
+}
+
+func TestBackwardCompatibility(t *testing.T) {
+	testWithChecker(t, func(t *testing.T, checker *Checker, testAddr string) {
+		// Old API should still work
+		err := checker.CheckAddr(testAddr, 2*time.Second)
+		if err != nil {
+			t.Errorf("Old API connection to test server failed: %v", err)
+		}
+
+		// CheckAddrZeroLinger should use new implementation internally
+		err = checker.CheckAddrZeroLinger(testAddr, 2*time.Second, true)
+		if err != nil {
+			t.Errorf("CheckAddrZeroLinger connection to test server failed: %v", err)
+		}
+	})
+}

socket_linux.go

@@ -12,15 +12,28 @@ import (
 
 const maxEpollEvents = 32
 
-// createSocket creates a socket with necessary options set.
-func createSocketZeroLinger(family int, zeroLinger bool) (fd int, err error) {
+// createSocketWithOptions creates a socket with specified options
+func createSocketWithOptions(family int, zeroLinger bool, mark int) (fd int, err error) {
 	// Create socket
 	fd, err = _createNonBlockingSocket(family)
-	if err == nil {
-		if zeroLinger {
-			err = _setZeroLinger(fd)
+	if err != nil {
+		return
+	}
+
+	if zeroLinger {
+		if err = _setZeroLinger(fd); err != nil {
+			_ = unix.Close(fd)
+			return
 		}
 	}
+
+	if mark != 0 {
+		if err = _setMark(fd, mark); err != nil {
+			_ = unix.Close(fd)
+			return
+		}
+	}
+
 	return
 }
 
@@ -62,6 +75,11 @@ func _setZeroLinger(fd int) error {
 	return unix.SetsockoptLinger(fd, unix.SOL_SOCKET, unix.SO_LINGER, &zeroLinger)
 }
 
+// setMark sets SO_MARK for given fd (Linux only)
+func _setMark(fd int, mark int) error {
+	return unix.SetsockoptInt(fd, unix.SOL_SOCKET, unix.SO_MARK, mark)
+}
+
 func createPoller() (fd int, err error) {
 	fd, err = unix.EpollCreate1(unix.EPOLL_CLOEXEC)
 	if err != nil {

socket_unix.go

@@ -10,7 +10,12 @@ import (
 
 // parseSockAddr resolves given addr to unix.Sockaddr
 func parseSockAddr(addr string) (sAddr unix.Sockaddr, family int, err error) {
-	tAddr, err := net.ResolveTCPAddr("tcp", addr)
+	return parseSockAddrWithNetwork(addr, "tcp")
+}
+
+// parseSockAddrWithNetwork resolves given addr to unix.Sockaddr with specified network
+func parseSockAddrWithNetwork(addr, network string) (sAddr unix.Sockaddr, family int, err error) {
+	tAddr, err := net.ResolveTCPAddr(network, addr)
 	if err != nil {
 		return
 	}

test_server_test.go

@@ -2,6 +2,7 @@ package tcp
 
 import (
 	"context"
+	"net"
 	"net/http"
 	"net/http/httptest"
 )
@@ -12,3 +13,29 @@ func StartTestServer() (string, context.CancelFunc) {
 	addr := ts.Listener.Addr().String()
 	return addr, ts.Close
 }
+
+// StartTestServerIPv6 starts a test HTTP server on IPv6 loopback and returns its address and a cancel function
+func StartTestServerIPv6() (string, context.CancelFunc, error) {
+	// Create a listener on IPv6 loopback address
+	listener, err := net.Listen("tcp6", "[::1]:0")
+	if err != nil {
+		return "", nil, err
+	}
+
+	// Create HTTP server with the IPv6 listener
+	server := &http.Server{
+		Handler: http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {}),
+	}
+
+	// Start server in background
+	go func() {
+		_ = server.Serve(listener)
+	}()
+
+	addr := listener.Addr().String()
+	cancel := func() {
+		_ = server.Close()
+	}
+
+	return addr, cancel, nil
+}