ping: Add timeout to destinations

Author: nekohasekai

ping/destination.go

@@ -6,6 +6,7 @@ import (
 	"net/netip"
 	"os"
 	"runtime"
+	"time"
 
 	"github.com/sagernet/sing-tun"
 	"github.com/sagernet/sing/common/buf"
@@ -17,13 +18,21 @@ import (
 var _ tun.DirectRouteDestination = (*Destination)(nil)
 
 type Destination struct {
+	conn         *Conn
 	ctx          context.Context
 	logger       logger.ContextLogger
 	routeContext tun.DirectRouteContext
-	conn         *Conn
+	timeout      time.Duration
 }
 
-func ConnectDestination(ctx context.Context, logger logger.ContextLogger, controlFunc control.Func, address netip.Addr, routeContext tun.DirectRouteContext) (tun.DirectRouteDestination, error) {
+func ConnectDestination(
+	ctx context.Context,
+	logger logger.ContextLogger,
+	controlFunc control.Func,
+	address netip.Addr,
+	routeContext tun.DirectRouteContext,
+	timeout time.Duration,
+) (tun.DirectRouteDestination, error) {
 	var (
 		conn *Conn
 		err  error
@@ -41,19 +50,25 @@ func ConnectDestination(ctx context.Context, logger logger.ContextLogger, contro
 		return nil, err
 	}
 	d := &Destination{
+		conn:         conn,
 		ctx:          ctx,
 		logger:       logger,
 		routeContext: routeContext,
-		conn:         conn,
+		timeout:      timeout,
 	}
 	go d.loopRead()
 	return d, nil
 }
 
 func (d *Destination) loopRead() {
+	defer d.Close()
 	for {
 		buffer := buf.NewPacket()
-		err := d.conn.ReadIP(buffer)
+		err := d.conn.SetReadDeadline(time.Now().Add(d.timeout))
+		if err != nil {
+			d.logger.ErrorContext(d.ctx, E.Cause(err, "set read deadline for ICMP conn"))
+		}
+		err = d.conn.ReadIP(buffer)
 		if err != nil {
 			buffer.Release()
 			if !E.IsClosed(err) {
@@ -76,3 +91,11 @@ func (d *Destination) WritePacket(packet *buf.Buffer) error {
 func (d *Destination) Close() error {
 	return d.conn.Close()
 }
+
+func (d *Destination) IsClosed() bool {
+	_, err := d.conn.conn.Write([]byte{})
+	if err != nil {
+		return false
+	}
+	return true
+}

ping/destination_gvisor.go

@@ -5,11 +5,13 @@ package ping
 import (
 	"context"
 	"net/netip"
+	"time"
 
 	"github.com/sagernet/gvisor/pkg/tcpip"
 	"github.com/sagernet/gvisor/pkg/tcpip/adapters/gonet"
 	"github.com/sagernet/gvisor/pkg/tcpip/header"
 	"github.com/sagernet/gvisor/pkg/tcpip/stack"
+	"github.com/sagernet/gvisor/pkg/tcpip/transport"
 	"github.com/sagernet/gvisor/pkg/waiter"
 	"github.com/sagernet/sing-tun"
 	"github.com/sagernet/sing/common"
@@ -23,8 +25,10 @@ var _ tun.DirectRouteDestination = (*GVisorDestination)(nil)
 type GVisorDestination struct {
 	ctx      context.Context
 	logger   logger.ContextLogger
+	endpoint tcpip.Endpoint
 	conn     *gonet.TCPConn
 	rewriter *Rewriter
+	timeout  time.Duration
 }
 
 func ConnectGVisor(
@@ -33,6 +37,7 @@ func ConnectGVisor(
 	routeContext tun.DirectRouteContext,
 	stack *stack.Stack,
 	bindAddress4, bindAddress6 netip.Addr,
+	timeout time.Duration,
 ) (*GVisorDestination, error) {
 	var (
 		bindAddress tcpip.Address
@@ -76,16 +81,23 @@ func ConnectGVisor(
 	destination := &GVisorDestination{
 		ctx:      ctx,
 		logger:   logger,
+		endpoint: endpoint,
 		conn:     gonet.NewTCPConn(&wq, endpoint),
 		rewriter: rewriter,
+		timeout:  timeout,
 	}
 	go destination.loopRead()
 	return destination, nil
 }
 
 func (d *GVisorDestination) loopRead() {
+	defer d.endpoint.Close()
 	for {
 		buffer := buf.NewPacket()
+		err := d.conn.SetReadDeadline(time.Now().Add(d.timeout))
+		if err != nil {
+			d.logger.ErrorContext(d.ctx, E.Cause(err, "set read deadline for ICMP conn"))
+		}
 		n, err := d.conn.Read(buffer.FreeBytes())
 		if err != nil {
 			buffer.Release()
@@ -111,3 +123,7 @@ func (d *GVisorDestination) WritePacket(packet *buf.Buffer) error {
 func (d *GVisorDestination) Close() error {
 	return d.conn.Close()
 }
+
+func (d *GVisorDestination) IsClosed() bool {
+	return transport.DatagramEndpointState(d.endpoint.State()) == transport.DatagramEndpointStateClosed
+}

ping/socket_linux_unprivileged.go

@@ -74,6 +74,9 @@ func (c *UnprivilegedConn) ReadMsg(b []byte, oob []byte) (n, oobn int, addr neti
 }
 
 func (c *UnprivilegedConn) Write(b []byte) (n int, err error) {
+	if len(b) == 0 {
+		return
+	}
 	conn, err := connect(false, c.controlFunc, c.destination)
 	if err != nil {
 		return

route_direct.go

@@ -13,6 +13,7 @@ import (
 type DirectRouteDestination interface {
 	WritePacket(packet *buf.Buffer) error
 	Close() error
+	IsClosed() bool
 }
 
 type DirectRouteSession struct {
@@ -28,6 +29,9 @@ type DirectRouteMapping struct {
 
 func NewDirectRouteMapping(timeout time.Duration) *DirectRouteMapping {
 	mapping := common.Must1(freelru.NewSharded[DirectRouteSession, DirectRouteDestination](1024, maphash.NewHasher[DirectRouteSession]().Hash32))
+	mapping.SetHealthCheck(func(session DirectRouteSession, destination DirectRouteDestination) bool {
+		return !destination.IsClosed()
+	})
 	mapping.SetOnEvict(func(session DirectRouteSession, action DirectRouteDestination) {
 		action.Close()
 	})