diff --git a/rpc_server.go b/rpc_server.go index 449ba6cc..064809d2 100644 --- a/rpc_server.go +++ b/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 { diff --git a/server_test.go b/server_test.go index 1aaa0b8d..befb8b00 100644 --- a/server_test.go +++ b/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) }