diff --git a/client_test.go b/client_test.go index 7c821545..018d9cbf 100644 --- a/client_test.go +++ b/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 diff --git a/plugin_test.go b/plugin_test.go index 7ce24aa6..b66b4051 100644 --- a/plugin_test.go +++ b/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": diff --git a/server.go b/server.go index b2371148..0e244c8f 100644 --- a/server.go +++ b/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 }