diff --git a/ping/destination.go b/ping/destination.go index 60decb4..6105024 100644 --- a/ping/destination.go +++ b/ping/destination.go @@ -51,11 +51,11 @@ func ConnectDestination( ) switch runtime.GOOS { case "darwin", "ios", "windows": - conn, err = Connect(ctx, false, controlFunc, destination) + conn, err = Connect(ctx, false, controlFunc, destination, timeout) default: - conn, err = Connect(ctx, true, controlFunc, destination) + conn, err = Connect(ctx, true, controlFunc, destination, timeout) if errors.Is(err, os.ErrPermission) { - conn, err = Connect(ctx, false, controlFunc, destination) + conn, err = Connect(ctx, false, controlFunc, destination, timeout) } } if err != nil { diff --git a/ping/ping.go b/ping/ping.go index eab977b..248987c 100644 --- a/ping/ping.go +++ b/ping/ping.go @@ -30,22 +30,22 @@ type Conn struct { readMsg func(b, oob []byte) (n, oobn int, addr netip.Addr, err error) } -func Connect(ctx context.Context, privileged bool, controlFunc control.Func, destination netip.Addr) (*Conn, error) { +func Connect(ctx context.Context, privileged bool, controlFunc control.Func, destination netip.Addr, idleTimeout time.Duration) (*Conn, error) { c := &Conn{ ctx: ctx, privileged: privileged, destination: destination, } - err := c.connect(controlFunc) + err := c.connect(controlFunc, idleTimeout) if err != nil { return nil, err } return c, nil } -func (c *Conn) connect(controlFunc control.Func) (err error) { +func (c *Conn) connect(controlFunc control.Func, idleTimeout time.Duration) (err error) { if c.isLinuxUnprivileged() { - c.conn, err = newUnprivilegedConn(c.ctx, controlFunc, c.destination) + c.conn, err = newUnprivilegedConn(c.ctx, controlFunc, c.destination, idleTimeout) } else { c.conn, err = connect(c.privileged, controlFunc, c.destination) } diff --git a/ping/ping_test.go b/ping/ping_test.go index 7ec291a..5a04be1 100644 --- a/ping/ping_test.go +++ b/ping/ping_test.go @@ -72,7 +72,7 @@ func TestPing(t *testing.T) { } func testPingIPv4ReadIP(t *testing.T, privileged bool, addr string) { - conn, err := ping.Connect(context.Background(), privileged, nil, netip.MustParseAddr(addr)) + conn, err := ping.Connect(context.Background(), privileged, nil, netip.MustParseAddr(addr), 0) if runtime.GOOS == "linux" && err != nil && err.Error() == "socket(): permission denied" { t.SkipNow() } @@ -105,7 +105,7 @@ func testPingIPv4ReadIP(t *testing.T, privileged bool, addr string) { } func testPingIPv4ReadICMP(t *testing.T, privileged bool, addr string) { - conn, err := ping.Connect(context.Background(), privileged, nil, netip.MustParseAddr(addr)) + conn, err := ping.Connect(context.Background(), privileged, nil, netip.MustParseAddr(addr), 0) if runtime.GOOS == "linux" && err != nil && err.Error() == "socket(): permission denied" { t.SkipNow() } @@ -137,7 +137,7 @@ func testPingIPv4ReadICMP(t *testing.T, privileged bool, addr string) { } func testPingIPv6ReadIP(t *testing.T, privileged bool, addr string) { - conn, err := ping.Connect(context.Background(), privileged, nil, netip.MustParseAddr(addr)) + conn, err := ping.Connect(context.Background(), privileged, nil, netip.MustParseAddr(addr), 0) if runtime.GOOS == "linux" && err != nil && err.Error() == "socket(): permission denied" { t.SkipNow() } @@ -169,7 +169,7 @@ func testPingIPv6ReadIP(t *testing.T, privileged bool, addr string) { } func testPingIPv6ReadICMP(t *testing.T, privileged bool, addr string) { - conn, err := ping.Connect(context.Background(), privileged, nil, netip.MustParseAddr(addr)) + conn, err := ping.Connect(context.Background(), privileged, nil, netip.MustParseAddr(addr), 0) if runtime.GOOS == "linux" && err != nil && err.Error() == "socket(): permission denied" { t.SkipNow() } diff --git a/ping/socket_linux_unprivileged.go b/ping/socket_linux_unprivileged.go index 79fd682..c0026ad 100644 --- a/ping/socket_linux_unprivileged.go +++ b/ping/socket_linux_unprivileged.go @@ -21,6 +21,7 @@ type UnprivilegedConn struct { cancel context.CancelFunc controlFunc control.Func destination netip.Addr + idleTimeout time.Duration receiveChan chan *unprivilegedResponse readDeadline pipe.Deadline mappingAccess sync.Mutex @@ -33,7 +34,7 @@ type unprivilegedResponse struct { Addr netip.Addr } -func newUnprivilegedConn(ctx context.Context, controlFunc control.Func, destination netip.Addr) (net.Conn, error) { +func newUnprivilegedConn(ctx context.Context, controlFunc control.Func, destination netip.Addr, idleTimeout time.Duration) (net.Conn, error) { conn, err := connect(false, controlFunc, destination) if err != nil { return nil, err @@ -45,6 +46,7 @@ func newUnprivilegedConn(ctx context.Context, controlFunc control.Func, destinat cancel: cancel, controlFunc: controlFunc, destination: destination, + idleTimeout: idleTimeout, receiveChan: make(chan *unprivilegedResponse), readDeadline: pipe.MakeDeadline(), mapping: make(map[uint16]net.Conn), @@ -116,6 +118,12 @@ func (c *UnprivilegedConn) Write(b []byte) (n int, err error) { func (c *UnprivilegedConn) fetchResponse(conn *net.UDPConn, identifier uint16) { defer c.removeConn(conn, identifier) for { + if c.idleTimeout > 0 { + err := conn.SetReadDeadline(time.Now().Add(c.idleTimeout)) + if err != nil { + return + } + } buffer := buf.NewPacket() cmsgBuffer := buf.NewSize(1024) n, oobN, _, addr, err := conn.ReadMsgUDPAddrPort(buffer.FreeBytes(), cmsgBuffer.FreeBytes())