Fix muxer.Attach (#1493)

The muxer's attach method is broken when setup with a nil host. It
panics with a nil pointer exception. This replaces the previous
implementation with a straight dispatch call to subsidiary servers.

x/muxer/muxer.go

@@ -16,7 +16,9 @@ package muxer
 
 import (
 	"context"
+	"errors"
 	"fmt"
+	"io"
 	"sync"
 
 	"github.com/golang/glog"
@@ -59,7 +61,7 @@ type server = rpc.ResourceProviderServer
 type muxer struct {
 	rpc.UnimplementedResourceProviderServer
 
-	host *provider.HostClient
+	host hostClient
 
 	dispatchTable dispatchTable
 
@@ -70,6 +72,12 @@ type muxer struct {
 	getMappingByKey map[string]MultiMappingHandler
 }
 
+// An interface to make *provider.HostClient test-able.
+type hostClient interface {
+	io.Closer
+	Log(context.Context, diag.Severity, urn.URN, string) error
+}
+
 type GetMappingArgs interface {
 	Fetch() []GetMappingResponse
 }
@@ -394,19 +402,24 @@ func (m *muxer) GetPluginInfo(ctx context.Context, e *emptypb.Empty) (*rpc.Plugi
 }
 
 func (m *muxer) Attach(ctx context.Context, req *rpc.PluginAttach) (*emptypb.Empty, error) {
-	host, err := provider.NewHostClient(req.GetAddress())
-	if err != nil {
-		return nil, err
+	attach := make([]func() error, len(m.servers))
+	for i, s := range m.servers {
+		s := s
+		attach[i] = func() error {
+			_, err := s.Attach(ctx, req)
+			return err
+		}
 	}
+
+	var closeErr error
 	if m.host != nil {
-		if err := m.host.Close(); err != nil {
-			return nil, err
-		}
+		closeErr = m.host.Close()
 	}
-	// Here we override the underlying host. This should replace the host instance of
-	// each subsidiary provider.
-	*m.host = *host
-	return &emptypb.Empty{}, nil
+
+	var err error
+	m.host, err = provider.NewHostClient(req.GetAddress())
+
+	return &emptypb.Empty{}, errors.Join(append(asyncJoin(attach), err, closeErr)...)
 }
 
 type getMappingArgs struct {

x/muxer/muxer_test.go

@@ -0,0 +1,89 @@
+// Copyright 2016-2023, Pulumi Corporation.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package muxer
+
+import (
+	"context"
+	"fmt"
+	"testing"
+
+	"github.com/pulumi/pulumi/sdk/v3/go/common/diag"
+	urn "github.com/pulumi/pulumi/sdk/v3/go/common/resource"
+	pulumirpc "github.com/pulumi/pulumi/sdk/v3/proto/go"
+	"github.com/stretchr/testify/assert"
+	emptypb "google.golang.org/protobuf/types/known/emptypb"
+)
+
+func TestAttach(t *testing.T) {
+	req := &pulumirpc.PluginAttach{Address: "test"}
+	ctx := context.Background()
+
+	t.Run("empty", func(t *testing.T) {
+		h := &host{}
+		m := &muxer{host: h}
+		_, err := m.Attach(ctx, req)
+		assert.NoError(t, err)
+		assert.NotZero(t, m.host)
+		assert.True(t, h.closed)
+	})
+
+	t.Run("dispatch", func(t *testing.T) {
+		h := &host{}
+		m := &muxer{
+			host: h,
+			servers: []server{
+				&attach{t: t, expected: "test"},
+				&attach{t: t, expected: "test"},
+			}}
+		_, err := m.Attach(ctx, req)
+		assert.NoError(t, err)
+		for i, s := range m.servers {
+			assert.Equalf(t, 1, s.(*attach).called, "i = %d", i)
+		}
+		assert.NotZero(t, m.host)
+		assert.True(t, h.closed)
+	})
+}
+
+type attach struct {
+	pulumirpc.UnimplementedResourceProviderServer
+
+	t        *testing.T
+	expected string
+	called   int
+}
+
+func (s *attach) Attach(ctx context.Context, req *pulumirpc.PluginAttach) (*emptypb.Empty, error) {
+	assert.Equal(s.t, s.expected, req.Address)
+	s.called++
+	return &emptypb.Empty{}, nil
+}
+
+type host struct{ closed bool }
+
+func (h *host) Close() error {
+	if h.closed {
+		return fmt.Errorf("host already closed")
+	}
+	h.closed = true
+	return nil
+}
+
+func (h *host) Log(context.Context, diag.Severity, urn.URN, string) error {
+	if h.closed {
+		return fmt.Errorf("cannot log against a closed host")
+	}
+	return nil
+}