From aa764512ec186120c28bac7cf7f0d45be8b3eba1 Mon Sep 17 00:00:00 2001
From: Ning Bo <ning.bo9@zte.com.cn>
Date: Fri, 8 May 2020 17:07:23 +0800
Subject: [PATCH] vsock: add DialTimeout

Signed-off-by: Ning Bo <ning.bo9@zte.com.cn>
---
 conn_linux.go             | 18 ++++++++++++++++++
 integration_linux_test.go | 35 +++++++++++++++++++++++++++++++++++
 vsock.go                  | 14 ++++++++++++++
 vsock_others.go           |  3 ++-
 4 files changed, 69 insertions(+), 1 deletion(-)

diff --git a/conn_linux.go b/conn_linux.go
index d9558cd..9685653 100644
--- a/conn_linux.go
+++ b/conn_linux.go
@@ -3,6 +3,8 @@
 package vsock
 
 import (
+	"time"
+
 	"golang.org/x/sys/unix"
 )
 
@@ -32,6 +34,22 @@ func dial(cid, port uint32) (*Conn, error) {
 	return dialLinux(cfd, cid, port)
 }
 
+// dialTimeout acts like dial but takes a timeout.
+func dialTimeout(cid, port uint32, timeout time.Duration) (*Conn, error) {
+	cfd, err := newConnFD()
+	if err != nil {
+		return nil, err
+	}
+
+	tv := unix.NsecToTimeval(timeout.Nanoseconds())
+	err = unix.SetsockoptTimeval(cfd.fd, unix.AF_VSOCK, unix.SO_VM_SOCKETS_CONNECT_TIMEOUT, &tv)
+	if err != nil {
+		return nil, err
+	}
+
+	return dialLinux(cfd, cid, port)
+}
+
 // dialLinux is the entry point for tests on Linux.
 func dialLinux(cfd connFD, cid, port uint32) (c *Conn, err error) {
 	defer func() {
diff --git a/integration_linux_test.go b/integration_linux_test.go
index ac12420..bdb211c 100644
--- a/integration_linux_test.go
+++ b/integration_linux_test.go
@@ -207,3 +207,38 @@ func makeLocalPipe(
 		}
 	}
 }
+
+func makeVSockPipeWithTimeout(timeout time.Duration) nettest.MakePipe {
+	return makeLocalPipe(
+		func() (net.Listener, error) { return vsock.Listen(0) },
+		func(addr net.Addr) (net.Conn, error) {
+			a := addr.(*vsock.Addr)
+			return vsock.DialTimeout(a.ContextID, a.Port, timeout)
+		},
+	)
+}
+
+func TestIntegrationConnTimeout(t *testing.T) {
+	vsutil.SkipHostIntegration(t)
+
+	timer := time.AfterFunc(10*time.Second, func() {
+		panic("test took too long")
+	})
+	defer timer.Stop()
+
+	mp0 := makeVSockPipeWithTimeout(time.Microsecond)
+	_, _, stop0, err0 := mp0()
+	if err0 == nil {
+		defer stop0()
+	}
+	if !os.IsTimeout(err0) {
+		t.Fatalf("connect with a very short dial timeout expect `ETIMEDOUT`, but got %v", err0)
+	}
+
+	mp1 := makeVSockPipeWithTimeout(time.Second * 2)
+	_, _, stop1, err1 := mp1()
+	if err1 != nil {
+		t.Fatalf("failed to make pipe: %v", err1)
+	}
+	defer stop1()
+}
diff --git a/vsock.go b/vsock.go
index 004263e..0f42c08 100644
--- a/vsock.go
+++ b/vsock.go
@@ -149,6 +149,20 @@ func Dial(contextID, port uint32) (*Conn, error) {
 	return c, nil
 }
 
+// DialTimeout acts like Dial but takes a timeout.
+func DialTimeout(contextID, port uint32, timeout time.Duration) (*Conn, error) {
+	c, err := dialTimeout(contextID, port, timeout)
+	if err != nil {
+		// No local address, but we have a remote address we can return.
+		return nil, opError(opDial, err, nil, &Addr{
+			ContextID: contextID,
+			Port:      port,
+		})
+	}
+
+	return c, nil
+}
+
 var _ net.Conn = &Conn{}
 var _ syscall.Conn = &Conn{}
 
diff --git a/vsock_others.go b/vsock_others.go
index a246de9..b78f54e 100644
--- a/vsock_others.go
+++ b/vsock_others.go
@@ -26,7 +26,8 @@ func (*listener) Addr() net.Addr                { return nil }
 func (*listener) Close() error                  { return errUnimplemented }
 func (*listener) SetDeadline(_ time.Time) error { return errUnimplemented }
 
-func dial(_, _ uint32) (*Conn, error) { return nil, errUnimplemented }
+func dial(_, _ uint32) (*Conn, error)                         { return nil, errUnimplemented }
+func dialTimeout(_, _ uint32, _ time.Duration) (*Conn, error) { return nil, errUnimplemented }
 
 type connFD struct{}