Plugin conn refactor

This changes things to rely on a plugin server that manages all
connections made to the server.

An optional handler can be passed into the server when the caller wants
to do extra things with the connection.

It is the caller's responsability to close the server.
When the server is closed, first all existing connections are closed
(and new connections are prevented).

Now the signal loop only needs to close the server and not deal with
`net.Conn`'s directly (or double-indirects as the case was before this
change).

Signed-off-by: Brian Goff <[email protected]>

Author: cpuguy83

cli-plugins/socket/socket.go

@@ -7,24 +7,106 @@ import (
 	"io"
 	"net"
 	"os"
+	"runtime"
+	"sync"
 )
 
 // EnvKey represents the well-known environment variable used to pass the plugin being
 // executed the socket name it should listen on to coordinate with the host CLI.
 const EnvKey = "DOCKER_CLI_PLUGIN_SOCKET"
 
-// SetupConn sets up a Unix socket listener, establishes a goroutine to handle connections
-// and update the conn pointer, and returns the listener for the socket (which the caller
-// is responsible for closing when it's no longer needed).
-func SetupConn(conn **net.UnixConn) (*net.UnixListener, error) {
-	listener, err := listen("docker_cli_" + randomID())
+// NewPluginServer creates a plugin server that listens on a new Unix domain socket.
+// `h` is called for each new connection to the socket in a goroutine.
+func NewPluginServer(h func(net.Conn)) (*PluginServer, error) {
+	l, err := listen("docker_cli_" + randomID())
 	if err != nil {
 		return nil, err
 	}
 
-	accept(listener, conn)
+	if h == nil {
+		h = func(net.Conn) {}
+	}
+
+	pl := &PluginServer{
+		l: l,
+		h: h,
+	}
+
+	unlinkOnce := sync.OnceFunc(func() {
+		unlink(l)
+	})
+
+	go func() {
+		defer pl.Close()
+		for {
+			err := pl.accept()
+			unlinkOnce()
+			if err != nil {
+				return
+			}
+		}
+	}()
+
+	return pl, nil
+}
+
+type PluginServer struct {
+	mu     sync.Mutex
+	conns  []net.Conn
+	l      *net.UnixListener
+	h      func(net.Conn)
+	closed bool
+}
+
+func (l *PluginServer) accept() error {
+	conn, err := l.l.Accept()
+	if err != nil {
+		return err
+	}
+
+	l.mu.Lock()
+	defer l.mu.Unlock()
+
+	if l.closed {
+		// handle potential race condition between Close and Accept
+		conn.Close()
+		return errors.New("plugin server is closed")
+	}
+
+	l.conns = append(l.conns, conn)
+
+	go l.h(conn)
+	return nil
+}
+
+func (l *PluginServer) Addr() net.Addr {
+	return l.l.Addr()
+}
+
+// Close ensures that the server is no longer accepting new connections and closes all existing connections.
+// Existing connections will receive [io.EOF].
+func (l *PluginServer) Close() error {
+	// close connections first to ensure the connections get io.EOF instead of a connection reset.
+	l.closeAllConns()
+
+	// Try to ensure that any active connections have a chance to receive io.EOF
+	runtime.Gosched()
+
+	return l.l.Close()
+}
 
-	return listener, nil
+func (l *PluginServer) closeAllConns() {
+	l.mu.Lock()
+	defer l.mu.Unlock()
+
+	// Prevent new connections from being accepted
+	l.closed = true
+
+	for _, conn := range l.conns {
+		conn.Close()
+	}
+
+	l.conns = nil
 }
 
 func randomID() string {
@@ -35,18 +117,6 @@ func randomID() string {
 	return hex.EncodeToString(b)
 }
 
-func accept(listener *net.UnixListener, conn **net.UnixConn) {
-	go func() {
-		for {
-			// ignore error here, if we failed to accept a connection,
-			// conn is nil and we fallback to previous behavior
-			*conn, _ = listener.AcceptUnix()
-			// perform any platform-specific actions on accept (e.g. unlink non-abstract sockets)
-			onAccept(*conn, listener)
-		}
-	}()
-}
-
 // ConnectAndWait connects to the socket passed via well-known env var,
 // if present, and attempts to read from it until it receives an EOF, at which
 // point cb is called.

cli-plugins/socket/socket_darwin.go

@@ -14,6 +14,6 @@ func listen(socketname string) (*net.UnixListener, error) {
 	})
 }
 
-func onAccept(conn *net.UnixConn, listener *net.UnixListener) {
+func unlink(listener *net.UnixListener) {
 	syscall.Unlink(listener.Addr().String())
 }

cli-plugins/socket/socket_nodarwin.go

@@ -13,7 +13,7 @@ func listen(socketname string) (*net.UnixListener, error) {
 	})
 }
 
-func onAccept(conn *net.UnixConn, listener *net.UnixListener) {
+func unlink(listener *net.UnixListener) {
 	// do nothing
 	// while on darwin and OpenBSD we would unlink here;
 	// on non-darwin the socket is abstract and not present on the filesystem

cli-plugins/socket/socket_openbsd.go

@@ -14,6 +14,6 @@ func listen(socketname string) (*net.UnixListener, error) {
 	})
 }
 
-func onAccept(conn *net.UnixConn, listener *net.UnixListener) {
+func unlink(listener *net.UnixListener) {
 	syscall.Unlink(listener.Addr().String())
 }

cli-plugins/socket/socket_test.go

@@ -1,66 +1,128 @@
 package socket
 
 import (
+	"errors"
+	"io"
 	"io/fs"
 	"net"
 	"os"
 	"runtime"
 	"strings"
+	"sync/atomic"
 	"testing"
 	"time"
 
 	"gotest.tools/v3/assert"
 	"gotest.tools/v3/poll"
 )
 
-func TestSetupConn(t *testing.T) {
-	t.Run("updates conn when connected", func(t *testing.T) {
-		var conn *net.UnixConn
-		listener, err := SetupConn(&conn)
+func TestPluginServer(t *testing.T) {
+	t.Run("connection closes with EOF when server closes", func(t *testing.T) {
+		called := make(chan struct{})
+		srv, err := NewPluginServer(func(_ net.Conn) { close(called) })
 		assert.NilError(t, err)
-		assert.Check(t, listener != nil, "returned nil listener but no error")
-		addr, err := net.ResolveUnixAddr("unix", listener.Addr().String())
+		assert.Assert(t, srv != nil, "returned nil listener but no error")
+
+		addr, err := net.ResolveUnixAddr("unix", srv.Addr().String())
 		assert.NilError(t, err, "failed to resolve listener address")
 
-		_, err = net.DialUnix("unix", nil, addr)
+		conn, err := net.DialUnix("unix", nil, addr)
 		assert.NilError(t, err, "failed to dial returned listener")
+		defer conn.Close()
+
+		done := make(chan error, 1)
+		go func() {
+			_, err := conn.Read(make([]byte, 1))
+			done <- err
+		}()
+
+		select {
+		case <-called:
+		case <-time.After(10 * time.Millisecond):
+			t.Fatal("handler not called")
+		}
+
+		srv.Close()
 
-		pollConnNotNil(t, &conn)
+		select {
+		case err := <-done:
+			if !errors.Is(err, io.EOF) {
+				t.Fatalf("exepcted EOF error, got: %v", err)
+			}
+		case <-time.After(10 * time.Millisecond):
+		}
 	})
 
 	t.Run("allows reconnects", func(t *testing.T) {
-		var conn *net.UnixConn
-		listener, err := SetupConn(&conn)
+		var calls int32
+		h := func(_ net.Conn) {
+			atomic.AddInt32(&calls, 1)
+		}
+
+		srv, err := NewPluginServer(h)
 		assert.NilError(t, err)
-		assert.Check(t, listener != nil, "returned nil listener but no error")
-		addr, err := net.ResolveUnixAddr("unix", listener.Addr().String())
+		defer srv.Close()
+
+		addr := srv.Addr().(*net.UnixAddr)
+
+		assert.Check(t, addr != nil, "returned nil listener but no error")
+
+		_, err = net.ResolveUnixAddr("unix", addr.String())
 		assert.NilError(t, err, "failed to resolve listener address")
 
+		waitForCalls := func(n int) {
+			poll.WaitOn(t, func(t poll.LogT) poll.Result {
+				if atomic.LoadInt32(&calls) == int32(n) {
+					return poll.Success()
+				}
+				return poll.Continue("waiting for handler to be called")
+			})
+		}
+
 		otherConn, err := net.DialUnix("unix", nil, addr)
 		assert.NilError(t, err, "failed to dial returned listener")
-
 		otherConn.Close()
 
-		_, err = net.DialUnix("unix", nil, addr)
+		waitForCalls(1)
+
+		conn, err := net.DialUnix("unix", nil, addr)
 		assert.NilError(t, err, "failed to redial listener")
+		defer conn.Close()
+		waitForCalls(2)
+
+		// and again but don't close the existing connection
+		conn2, err := net.DialUnix("unix", nil, addr)
+		assert.NilError(t, err, "failed to redial listener")
+		defer conn2.Close()
+		waitForCalls(3)
+
+		srv.Close()
+
+		// now make sure we get EOF on the existing connections
+		buf := make([]byte, 1)
+		_, err = conn.Read(buf)
+		assert.ErrorIs(t, err, io.EOF, "expected EOF error, got: %v", err)
+
+		_, err = conn2.Read(buf)
+		assert.ErrorIs(t, err, io.EOF, "expected EOF error, got: %v", err)
 	})
 
 	t.Run("does not leak sockets to local directory", func(t *testing.T) {
-		var conn *net.UnixConn
-		listener, err := SetupConn(&conn)
+		srv, err := NewPluginServer(nil)
 		assert.NilError(t, err)
-		assert.Check(t, listener != nil, "returned nil listener but no error")
-		checkDirNoPluginSocket(t)
+		assert.Check(t, srv != nil, "returned nil server but no error")
+		checkDirNoNewPluginServer(t)
 
-		addr, err := net.ResolveUnixAddr("unix", listener.Addr().String())
+		_, err = net.ResolveUnixAddr("unix", srv.Addr().String())
 		assert.NilError(t, err, "failed to resolve listener address")
-		_, err = net.DialUnix("unix", nil, addr)
+
+		_, err = net.DialUnix("unix", nil, srv.Addr().(*net.UnixAddr))
 		assert.NilError(t, err, "failed to dial returned listener")
-		checkDirNoPluginSocket(t)
+		checkDirNoNewPluginServer(t)
 	})
 }
 
-func checkDirNoPluginSocket(t *testing.T) {
+func checkDirNoNewPluginServer(t *testing.T) {
 	t.Helper()
 
 	files, err := os.ReadDir(".")
@@ -78,18 +140,24 @@ func checkDirNoPluginSocket(t *testing.T) {
 
 func TestConnectAndWait(t *testing.T) {
 	t.Run("calls cancel func on EOF", func(t *testing.T) {
-		var conn *net.UnixConn
-		listener, err := SetupConn(&conn)
+		srv, err := NewPluginServer(nil)
 		assert.NilError(t, err, "failed to setup listener")
+		defer srv.Close()
 
 		done := make(chan struct{})
-		t.Setenv(EnvKey, listener.Addr().String())
+		t.Setenv(EnvKey, srv.Addr().String())
 		cancelFunc := func() {
 			done <- struct{}{}
 		}
 		ConnectAndWait(cancelFunc)
-		pollConnNotNil(t, &conn)
-		conn.Close()
+
+		select {
+		case <-done:
+			t.Fatal("unexpectedly done")
+		default:
+		}
+
+		srv.Close()
 
 		select {
 		case <-done:
@@ -101,17 +169,19 @@ func TestConnectAndWait(t *testing.T) {
 	// TODO: this test cannot be executed with `t.Parallel()`, due to
 	// relying on goroutine numbers to ensure correct behaviour
 	t.Run("connect goroutine exits after EOF", func(t *testing.T) {
-		var conn *net.UnixConn
-		listener, err := SetupConn(&conn)
+		srv, err := NewPluginServer(nil)
 		assert.NilError(t, err, "failed to setup listener")
-		t.Setenv(EnvKey, listener.Addr().String())
+
+		defer srv.Close()
+
+		t.Setenv(EnvKey, srv.Addr().String())
 		numGoroutines := runtime.NumGoroutine()
 
 		ConnectAndWait(func() {})
 		assert.Equal(t, runtime.NumGoroutine(), numGoroutines+1)
 
-		pollConnNotNil(t, &conn)
-		conn.Close()
+		srv.Close()
+
 		poll.WaitOn(t, func(t poll.LogT) poll.Result {
 			if runtime.NumGoroutine() > numGoroutines+1 {
 				return poll.Continue("waiting for connect goroutine to exit")
@@ -120,14 +190,3 @@ func TestConnectAndWait(t *testing.T) {
 		}, poll.WithDelay(1*time.Millisecond), poll.WithTimeout(10*time.Millisecond))
 	})
 }
-
-func pollConnNotNil(t *testing.T, conn **net.UnixConn) {
-	t.Helper()
-
-	poll.WaitOn(t, func(t poll.LogT) poll.Result {
-		if *conn == nil {
-			return poll.Continue("waiting for conn to not be nil")
-		}
-		return poll.Success()
-	}, poll.WithDelay(1*time.Millisecond), poll.WithTimeout(10*time.Millisecond))
-}