ServeConfig supports a manual listener

client_test.go

@@ -56,6 +56,39 @@ func TestClient(t *testing.T) {
 	}
 }
 
+func TestClient_manualListener(t *testing.T) {
+	process := helperProcess("test-grpc-manual-listener")
+	c := NewClient(&ClientConfig{
+		Cmd:              process,
+		HandshakeConfig:  testHandshake,
+		Plugins:          testGRPCPluginMap,
+		AllowedProtocols: []Protocol{ProtocolGRPC},
+	})
+	defer c.Kill()
+
+	// Test that it parses the proper address
+	addr, err := c.Start()
+	if err != nil {
+		t.Fatalf("err should be nil, got %s", err)
+	}
+
+	if addr.Network() != "tcp" {
+		t.Fatalf("bad: %#v", addr)
+	}
+
+	if addr.String() != "127.0.0.1:1234" {
+		t.Fatalf("bad: %#v", addr)
+	}
+
+	// Test that it exits properly if killed
+	c.Kill()
+
+	// Test that it knows it is exited
+	if !c.Exited() {
+		t.Fatal("should say client has exited")
+	}
+}
+
 // This tests a bug where Kill would start
 func TestClient_killStart(t *testing.T) {
 	// Create a temporary dir to store the result file

plugin_test.go

@@ -8,6 +8,7 @@ import (
 	"io"
 	"io/ioutil"
 	"log"
+	"net"
 	"net/rpc"
 	"os"
 	"os/exec"
@@ -549,6 +550,22 @@ func TestHelperProcess(*testing.T) {
 			TLSProvider:     helperTLSProvider,
 		})
 
+		// Shouldn't reach here but make sure we exit anyways
+		os.Exit(0)
+	case "test-grpc-manual-listener":
+		ln, err := net.Listen("tcp", "127.0.0.1:1234")
+		if err != nil {
+			panic(err)
+		}
+		defer ln.Close()
+
+		Serve(&ServeConfig{
+			HandshakeConfig: testHandshake,
+			Plugins:         testGRPCPluginMap,
+			GRPCServer:      DefaultGRPCServer,
+			Listener:        ln,
+		})
+
 		// Shouldn't reach here but make sure we exit anyways
 		os.Exit(0)
 	case "test-interface":

server.go

@@ -83,6 +83,21 @@ type ServeConfig struct {
 	// Logger is used to pass a logger into the server. If none is provided the
 	// server will create a default logger.
 	Logger hclog.Logger
+
+	// Listener is the listener that the plugin server will listen for
+	// plugin connections. THIS DOES NOT NORMALLY NEED TO BE SET. If this
+	// isn't set, the plugin chooses a listener. This is exposed in case you
+	// want to carefully control how a plugin is served.
+	//
+	// If TLSProvider is set, this listener will be wrapped with a TLS
+	// listener. If you want to manually control TLS you should set
+	// TLSProvider to nil but be aware that the client side will need to be
+	// manually made aware of the certificate used.
+	//
+	// Serve will take ownership of this listener and close it when it is
+	// complete. The caller should NOT close this listener once `Serve` is
+	// called.
+	Listener net.Listener
 }
 
 // protocolVersion determines the protocol version and plugin set to be used by
@@ -171,13 +186,28 @@ func protocolVersion(opts *ServeConfig) (int, Protocol, PluginSet) {
 //
 // This is the method that plugins should call in their main() functions.
 func Serve(opts *ServeConfig) {
+	// We use this to trigger an `os.Exit` so that we can execute our other
+	// deferred functions.
+	exitCode := -1
+	defer func() {
+		if exitCode >= 0 {
+			os.Exit(exitCode)
+		}
+	}()
+
+	// If our listener is not nil, then we want to close that on exit.
+	if opts.Listener != nil {
+		defer opts.Listener.Close()
+	}
+
 	// Validate the handshake config
 	if opts.MagicCookieKey == "" || opts.MagicCookieValue == "" {
 		fmt.Fprintf(os.Stderr,
 			"Misconfigured ServeConfig given to serve this plugin: no magic cookie\n"+
 				"key or value was set. Please notify the plugin author and report\n"+
 				"this as a bug.\n")
-		os.Exit(1)
+		exitCode = 1
+		return
 	}
 
 	// First check the cookie
@@ -186,7 +216,8 @@ func Serve(opts *ServeConfig) {
 			"This binary is a plugin. These are not meant to be executed directly.\n"+
 				"Please execute the program that consumes these plugins, which will\n"+
 				"load any plugins automatically\n")
-		os.Exit(1)
+		exitCode = 1
+		return
 	}
 
 	// negotiate the version and plugins
@@ -219,18 +250,21 @@ func Serve(opts *ServeConfig) {
 		os.Exit(1)
 	}
 
-	// Register a listener so we can accept a connection
-	listener, err := serverListener()
-	if err != nil {
-		logger.Error("plugin init error", "error", err)
-		return
-	}
+	listener := opts.Listener
+	if listener == nil {
+		// Register a listener so we can accept a connection
+		listener, err = serverListener()
+		if err != nil {
+			logger.Error("plugin init error", "error", err)
+			return
+		}
 
-	// Close the listener on return. We wrap this in a func() on purpose
-	// because the "listener" reference may change to TLS.
-	defer func() {
-		listener.Close()
-	}()
+		// Close the listener on return. We wrap this in a func() on purpose
+		// because the "listener" reference may change to TLS.
+		defer func() {
+			listener.Close()
+		}()
+	}
 
 	var tlsConfig *tls.Config
 	if opts.TLSProvider != nil {
@@ -349,6 +383,11 @@ func Serve(opts *ServeConfig) {
 
 	// Accept connections and wait for completion
 	go server.Serve(listener)
+
+	// Note that given the documentation of Serve we should probably be
+	// setting exitCode = 0 and using os.Exit here. That's how it used to
+	// work before extracting this library. However, for years we've done
+	// this so we'll keep this functionality.
 	<-doneCh
 }