From 1d9ab7803ac29113d29bdb35a29830a4105e1977 Mon Sep 17 00:00:00 2001
From: Nick Santos <nick@tilt.dev>
Date: Mon, 6 Jun 2022 12:23:53 -0400
Subject: [PATCH] cli: set timeout connection ping on sockets as well

Note that this does not fully fix the referenced issue, but
at least makes sure that API clients don't hang forever on
the initialization step.

See: https://github.com/docker/cli/issues/3652
Signed-off-by: Nick Santos <nick.santos@docker.com>
Signed-off-by: Sebastiaan van Stijn <github@gone.nl>
---
 cli/command/cli.go      | 14 +++++++++--
 cli/command/cli_test.go | 53 +++++++++++++++++++++++++++++++++++++++++
 2 files changed, 65 insertions(+), 2 deletions(-)

diff --git a/cli/command/cli.go b/cli/command/cli.go
index afa7cd3387..ba86491aef 100644
--- a/cli/command/cli.go
+++ b/cli/command/cli.go
@@ -35,6 +35,8 @@ import (
 	notaryclient "github.com/theupdateframework/notary/client"
 )
 
+const defaultInitTimeout = 2 * time.Second
+
 // Streams is an interface which exposes the standard input and output streams
 type Streams interface {
 	In() *streams.In
@@ -77,6 +79,7 @@ type DockerCli struct {
 	currentContext     string
 	dockerEndpoint     docker.Endpoint
 	contextStoreConfig store.Config
+	initTimeout        time.Duration
 }
 
 // DefaultVersion returns api.defaultVersion.
@@ -313,13 +316,20 @@ func resolveDefaultDockerEndpoint(opts *cliflags.CommonOptions) (docker.Endpoint
 	}, nil
 }
 
+func (cli *DockerCli) getInitTimeout() time.Duration {
+	if cli.initTimeout != 0 {
+		return cli.initTimeout
+	}
+	return defaultInitTimeout
+}
+
 func (cli *DockerCli) initializeFromClient() {
 	ctx := context.Background()
-	if strings.HasPrefix(cli.DockerEndpoint().Host, "tcp://") {
+	if !strings.HasPrefix(cli.DockerEndpoint().Host, "ssh://") {
 		// @FIXME context.WithTimeout doesn't work with connhelper / ssh connections
 		// time="2020-04-10T10:16:26Z" level=warning msg="commandConn.CloseWrite: commandconn: failed to wait: signal: killed"
 		var cancel func()
-		ctx, cancel = context.WithTimeout(ctx, 2*time.Second)
+		ctx, cancel = context.WithTimeout(ctx, cli.getInitTimeout())
 		defer cancel()
 	}
 
diff --git a/cli/command/cli_test.go b/cli/command/cli_test.go
index ba157ebe4b..73e17d7e4c 100644
--- a/cli/command/cli_test.go
+++ b/cli/command/cli_test.go
@@ -5,12 +5,15 @@ import (
 	"context"
 	"fmt"
 	"io"
+	"net"
 	"net/http"
 	"net/http/httptest"
 	"os"
+	"path/filepath"
 	"runtime"
 	"strings"
 	"testing"
+	"time"
 
 	"github.com/docker/cli/cli/config"
 	"github.com/docker/cli/cli/config/configfile"
@@ -165,6 +168,56 @@ func TestInitializeFromClient(t *testing.T) {
 	}
 }
 
+// Makes sure we don't hang forever on the initial connection.
+// https://github.com/docker/cli/issues/3652
+func TestInitializeFromClientHangs(t *testing.T) {
+	dir := t.TempDir()
+	socket := filepath.Join(dir, "my.sock")
+	l, err := net.Listen("unix", socket)
+	assert.NilError(t, err)
+
+	receiveReqCh := make(chan bool)
+	timeoutCtx, cancel := context.WithTimeout(context.Background(), time.Second)
+	defer cancel()
+
+	// Simulate a server that hangs on connections.
+	ts := httptest.NewUnstartedServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		select {
+		case <-timeoutCtx.Done():
+		case receiveReqCh <- true: // Blocks until someone receives on the channel.
+		}
+		_, _ = w.Write([]byte("OK"))
+	}))
+	ts.Listener = l
+	ts.Start()
+	defer ts.Close()
+
+	opts := &flags.CommonOptions{Hosts: []string{fmt.Sprintf("unix://%s", socket)}}
+	configFile := &configfile.ConfigFile{}
+	apiClient, err := NewAPIClientFromFlags(opts, configFile)
+	assert.NilError(t, err)
+
+	initializedCh := make(chan bool)
+
+	go func() {
+		cli := &DockerCli{client: apiClient, initTimeout: time.Millisecond}
+		cli.Initialize(flags.NewClientOptions())
+		close(initializedCh)
+	}()
+
+	select {
+	case <-timeoutCtx.Done():
+		t.Fatal("timeout waiting for initialization to complete")
+	case <-initializedCh:
+	}
+
+	select {
+	case <-timeoutCtx.Done():
+		t.Fatal("server never received an init request")
+	case <-receiveReqCh:
+	}
+}
+
 // The CLI no longer disables/hides experimental CLI features, however, we need
 // to verify that existing configuration files do not break
 func TestExperimentalCLI(t *testing.T) {