update bugfix. netpipes timeout.

Author: blacknon

connect.go

@@ -151,17 +151,23 @@ func (c *Connect) CreateClient(host, port, user string, authMethods []ssh.AuthMe
 	defer cancel()
 
 	// Dial to host:port
-	netConn, err := c.ProxyDialer.DialContext(ctx, "tcp", uri)
-	if err != nil {
-		return
+	netConn, cerr := c.ProxyDialer.DialContext(ctx, "tcp", uri)
+	if cerr != nil {
+		return cerr
 	}
 
+	// Set deadline
+	netConn.SetDeadline(time.Now().Add(time.Duration(timeout) * time.Second))
+
 	// Create new ssh connect
-	sshCon, channel, req, err := ssh.NewClientConn(netConn, uri, config)
-	if err != nil {
-		return
+	sshCon, channel, req, cerr := ssh.NewClientConn(netConn, uri, config)
+	if cerr != nil {
+		return cerr
 	}
 
+	// Reet deadline
+	netConn.SetDeadline(time.Time{})
+
 	// Create *ssh.Client
 	c.Client = ssh.NewClient(sshCon, channel, req)
 
@@ -172,40 +178,36 @@ func (c *Connect) CreateClient(host, port, user string, authMethods []ssh.AuthMe
 func (c *Connect) CreateSession() (session *ssh.Session, err error) {
 	// Create session
 	session, err = c.Client.NewSession()
-
 	return
 }
 
 // SendKeepAlive send packet to session.
 // TODO(blacknon): Interval及びMaxを設定できるようにする(v0.1.1)
 func (c *Connect) SendKeepAlive(session *ssh.Session) {
 	// keep alive interval (default 30 sec)
-	interval := 30
+	interval := 1
 	if c.SendKeepAliveInterval > 0 {
 		interval = c.SendKeepAliveInterval
 	}
 
-	// keep alive max (default 5)
-	max := 5
-	if c.SendKeepAliveMax > 0 {
-		max = c.SendKeepAliveMax
-	}
-
-	// keep alive counter
-	i := 0
 	for {
-		// Send keep alive packet
-		_, err := session.SendRequest("keepalive", true, nil)
-		// _, _, err := c.Client.SendRequest("keepalive", true, nil)
-		if err == nil {
-			i = 0
-		} else {
-			i += 1
-		}
+		// timeout channel
+		tc := make(chan bool, 1)
+
+		go func() {
+			// Send keep alive packet
+			_, err := session.SendRequest("keepalive", true, nil)
+			if err == nil {
+				tc <- true
+			}
+		}()
 
-		// check counter
-		if max <= i {
+		select {
+		case <-tc:
+		case <-time.After(time.Duration(c.ConnectTimeout) * time.Second):
 			session.Close()
+			c.Client.Close()
+			log.Println("keepalive timeout")
 			return
 		}
 

proxy.go

@@ -23,19 +23,31 @@ type ProxyDialer interface {
 }
 
 type ContextDialer struct {
-	dialer proxy.Dialer
+	Dialer proxy.Dialer
+}
+
+func (c *ContextDialer) GetDialer() proxy.Dialer {
+	return c.Dialer
 }
 
 func (c *ContextDialer) Dial(network, addr string) (net.Conn, error) {
-	return c.dialer.Dial(network, addr)
+	return c.Dialer.Dial(network, addr)
 }
 
 func (c *ContextDialer) DialContext(ctx context.Context, network, addr string) (net.Conn, error) {
+	// Simply call the DialContext method if supported
+	if dialerCtx, ok := c.Dialer.(interface {
+		DialContext(context.Context, string, string) (net.Conn, error)
+	}); ok {
+		return dialerCtx.DialContext(ctx, network, addr)
+	}
+
+	// Fallback if DialContext is not supported
 	connChan := make(chan net.Conn, 1)
 	errChan := make(chan error, 1)
 
 	go func() {
-		conn, err := c.dialer.Dial(network, addr)
+		conn, err := c.Dialer.Dial(network, addr)
 		if err != nil {
 			errChan <- err
 			return
@@ -96,7 +108,7 @@ func (p *Proxy) CreateProxyDialer() (proxyContextDialer ProxyDialer, err error)
 		proxyDialer, err = p.CreateProxyCommandProxyDialer()
 	}
 
-	proxyContextDialer = &ContextDialer{dialer: proxyDialer}
+	proxyContextDialer = &ContextDialer{Dialer: proxyDialer}
 
 	return
 }
@@ -158,6 +170,8 @@ func (p *Proxy) CreateProxyCommandProxyDialer() (proxyDialer proxy.Dialer, err e
 
 type NetPipe struct {
 	Command string
+	ctx     context.Context
+	Cmd     *exec.Cmd
 }
 
 func (n *NetPipe) Dial(network, addr string) (con net.Conn, err error) {
@@ -167,15 +181,15 @@ func (n *NetPipe) Dial(network, addr string) (con net.Conn, err error) {
 	// Create net.Pipe(), and set proxyCommand
 	con, srv := net.Pipe()
 
-	cmd := exec.Command("sh", "-c", n.Command)
+	n.Cmd = exec.Command("sh", "-c", n.Command)
 
 	// setup FD
-	cmd.Stdin = srv
-	cmd.Stdout = srv
-	cmd.Stderr = os.Stderr
+	n.Cmd.Stdin = srv
+	n.Cmd.Stdout = srv
+	n.Cmd.Stderr = os.Stderr
 
-	// run proxyCommand
-	err = cmd.Start()
+	// Start the command
+	err = n.Cmd.Start()
 
 	return
 }
@@ -190,6 +204,7 @@ func (n *NetPipe) DialContext(ctx context.Context, network, addr string) (con ne
 			errChan <- err
 			return
 		}
+
 		connChan <- conn
 	}()
 
@@ -199,6 +214,7 @@ func (n *NetPipe) DialContext(ctx context.Context, network, addr string) (con ne
 	case err := <-errChan:
 		return nil, err
 	case <-ctx.Done():
+		n.Cmd.Process.Kill()
 		return nil, ctx.Err()
 	}
 }
@@ -258,29 +274,6 @@ func (s *httpProxy) Dial(network, addr string) (net.Conn, error) {
 	return c, nil
 }
 
-func (s *httpProxy) DialContext(ctx context.Context, network, addr string) (con net.Conn, err error) {
-	connChan := make(chan net.Conn, 1)
-	errChan := make(chan error, 1)
-
-	go func() {
-		conn, err := s.Dial(network, addr)
-		if err != nil {
-			errChan <- err
-			return
-		}
-		connChan <- conn
-	}()
-
-	select {
-	case conn := <-connChan:
-		return conn, nil
-	case err := <-errChan:
-		return nil, err
-	case <-ctx.Done():
-		return nil, ctx.Err()
-	}
-}
-
 // newHttpProxy
 func newHttpProxy(uri *url.URL, forward proxy.Dialer) (proxy.Dialer, error) {
 	s := new(httpProxy)