Merge pull request #217 from hashicorp/fix-rpc-serve-close-chan

RPCServer: close chan when Serve returns

rpc_server.go

@@ -42,6 +42,8 @@ func (s *RPCServer) Config() string { return "" }
 
 // ServerProtocol impl.
 func (s *RPCServer) Serve(lis net.Listener) {
+	defer s.done()
+
 	for {
 		conn, err := lis.Accept()
 		if err != nil {
@@ -82,7 +84,7 @@ func (s *RPCServer) ServeConn(conn io.ReadWriteCloser) {
 
 	// Connect the stdstreams (in, out, err)
 	stdstream := make([]net.Conn, 2)
-	for i, _ := range stdstream {
+	for i := range stdstream {
 		stdstream[i], err = mux.Accept()
 		if err != nil {
 			mux.Close()
@@ -133,13 +135,15 @@ type controlServer struct {
 // Ping can be called to verify the connection (and likely the binary)
 // is still alive to a plugin.
 func (c *controlServer) Ping(
-	null bool, response *struct{}) error {
+	null bool, response *struct{},
+) error {
 	*response = struct{}{}
 	return nil
 }
 
 func (c *controlServer) Quit(
-	null bool, response *struct{}) error {
+	null bool, response *struct{},
+) error {
 	// End the server
 	c.server.done()
 
@@ -156,7 +160,8 @@ type dispenseServer struct {
 }
 
 func (d *dispenseServer) Dispense(
-	name string, response *uint32) error {
+	name string, response *uint32,
+) error {
 	// Find the function to create this implementation
 	p, ok := d.plugins[name]
 	if !ok {

server_test.go

@@ -151,6 +151,37 @@ func TestServer_testMode_AutoMTLS(t *testing.T) {
 	<-closeCh
 }
 
+func TestServer_RPC(t *testing.T) {
+	closeCh := make(chan struct{})
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	// make a server, but we don't need to attach to it
+	ch := make(chan *ReattachConfig, 1)
+	go Serve(&ServeConfig{
+		HandshakeConfig: testHandshake,
+		Plugins:         testPluginMap,
+		Test: &ServeTestConfig{
+			Context:          ctx,
+			CloseCh:          closeCh,
+			ReattachConfigCh: ch,
+		},
+	})
+
+	// Wait for the server
+	select {
+	case cfg := <-ch:
+		if cfg == nil {
+			t.Fatal("attach config should not be nil")
+		}
+	case <-time.After(2000 * time.Millisecond):
+		t.Fatal("should've received reattach")
+	}
+
+	cancel()
+	<-closeCh
+}
+
 func TestRmListener_impl(t *testing.T) {
 	var _ net.Listener = new(rmListener)
 }