From 6a5285ed0ee56cd7cbdbf90ed678edc0184a8654 Mon Sep 17 00:00:00 2001
From: zjj <jj@yuzao.org>
Date: Thu, 21 Nov 2019 14:26:41 +0800
Subject: [PATCH] Set timeout for custom dialer. (#1035)

Fixes #1034.
---
 AUTHORS      | 1 +
 connector.go | 8 +++++++-
 2 files changed, 8 insertions(+), 1 deletion(-)

diff --git a/AUTHORS b/AUTHORS
index 9765b5348..7d647012d 100644
--- a/AUTHORS
+++ b/AUTHORS
@@ -44,6 +44,7 @@ James Harr <james.harr at gmail.com>
 Jeff Hodges <jeff at somethingsimilar.com>
 Jeffrey Charles <jeffreycharles at gmail.com>
 Jerome Meyer <jxmeyer at gmail.com>
+Jiajia Zhong <zhong2plus at gmail.com>
 Jian Zhen <zhenjl at gmail.com>
 Joshua Prunier <joshua.prunier at gmail.com>
 Julien Lefevre <julien.lefevr at gmail.com>
diff --git a/connector.go b/connector.go
index eac0f01aa..d567b4e4f 100644
--- a/connector.go
+++ b/connector.go
@@ -37,7 +37,13 @@ func (c *connector) Connect(ctx context.Context) (driver.Conn, error) {
 	dial, ok := dials[mc.cfg.Net]
 	dialsLock.RUnlock()
 	if ok {
-		mc.netConn, err = dial(ctx, mc.cfg.Addr)
+		dctx := ctx
+		if mc.cfg.Timeout > 0 {
+			var cancel context.CancelFunc
+			dctx, cancel = context.WithTimeout(ctx, c.cfg.Timeout)
+			defer cancel()
+		}
+		mc.netConn, err = dial(dctx, mc.cfg.Addr)
 	} else {
 		nd := net.Dialer{Timeout: mc.cfg.Timeout}
 		mc.netConn, err = nd.DialContext(ctx, mc.cfg.Net, mc.cfg.Addr)