github.com/thajeztah/cli@v0.0.0-20240223162942-dc6bfac81a8b/cli-plugins/socket/socket.go (about)

     1  package socket
     2  
     3  import (
     4  	"crypto/rand"
     5  	"encoding/hex"
     6  	"errors"
     7  	"io"
     8  	"net"
     9  	"os"
    10  )
    11  
    12  // EnvKey represents the well-known environment variable used to pass the plugin being
    13  // executed the socket name it should listen on to coordinate with the host CLI.
    14  const EnvKey = "DOCKER_CLI_PLUGIN_SOCKET"
    15  
    16  // SetupConn sets up a Unix socket listener, establishes a goroutine to handle connections
    17  // and update the conn pointer, and returns the listener for the socket (which the caller
    18  // is responsible for closing when it's no longer needed).
    19  func SetupConn(conn **net.UnixConn) (*net.UnixListener, error) {
    20  	listener, err := listen("docker_cli_" + randomID())
    21  	if err != nil {
    22  		return nil, err
    23  	}
    24  
    25  	accept(listener, conn)
    26  
    27  	return listener, nil
    28  }
    29  
    30  func randomID() string {
    31  	b := make([]byte, 16)
    32  	if _, err := rand.Read(b); err != nil {
    33  		panic(err) // This shouldn't happen
    34  	}
    35  	return hex.EncodeToString(b)
    36  }
    37  
    38  func accept(listener *net.UnixListener, conn **net.UnixConn) {
    39  	go func() {
    40  		for {
    41  			// ignore error here, if we failed to accept a connection,
    42  			// conn is nil and we fallback to previous behavior
    43  			*conn, _ = listener.AcceptUnix()
    44  			// perform any platform-specific actions on accept (e.g. unlink non-abstract sockets)
    45  			onAccept(*conn, listener)
    46  		}
    47  	}()
    48  }
    49  
    50  // ConnectAndWait connects to the socket passed via well-known env var,
    51  // if present, and attempts to read from it until it receives an EOF, at which
    52  // point cb is called.
    53  func ConnectAndWait(cb func()) {
    54  	socketAddr, ok := os.LookupEnv(EnvKey)
    55  	if !ok {
    56  		// if a plugin compiled against a more recent version of docker/cli
    57  		// is executed by an older CLI binary, ignore missing environment
    58  		// variable and behave as usual
    59  		return
    60  	}
    61  	addr, err := net.ResolveUnixAddr("unix", socketAddr)
    62  	if err != nil {
    63  		return
    64  	}
    65  	conn, err := net.DialUnix("unix", nil, addr)
    66  	if err != nil {
    67  		return
    68  	}
    69  
    70  	go func() {
    71  		b := make([]byte, 1)
    72  		for {
    73  			_, err := conn.Read(b)
    74  			if errors.Is(err, io.EOF) {
    75  				cb()
    76  				return
    77  			}
    78  		}
    79  	}()
    80  }