fix(sdk): race condition issues on stub resolver implementation

Author: joe0BAB

plugin/resolver.go

@@ -4,8 +4,10 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"time"
 
 	"connectrpc.com/connect"
+	"github.com/sirupsen/logrus"
 	"google.golang.org/protobuf/proto"
 
 	resolverv1 "github.com/docker/secrets-engine/pkg/api/resolver/v1"
@@ -16,10 +18,20 @@ import (
 var _ resolverv1connect.ResolverServiceHandler = &resolverService{}
 
 type resolverService struct {
-	resolver secrets.Resolver
+	resolver            secrets.Resolver
+	setupCompleted      chan struct{}
+	registrationTimeout time.Duration
 }
 
 func (r *resolverService) GetSecret(ctx context.Context, c *connect.Request[resolverv1.GetSecretRequest]) (*connect.Response[resolverv1.GetSecretResponse], error) {
+	logrus.Debugf("GetSecret request (ID %q)", c.Msg.GetSecretId())
+	select {
+	case <-r.setupCompleted:
+	case <-ctx.Done():
+		return nil, connect.NewError(connect.CodeInternal, errors.New("context cancelled while waiting for registration"))
+	case <-time.After(r.registrationTimeout):
+		return nil, connect.NewError(connect.CodeDeadlineExceeded, fmt.Errorf("registration incomplete (timeout after %s)", r.registrationTimeout))
+	}
 	msgID := c.Msg.GetSecretId()
 	id, err := secrets.ParseID(msgID)
 	if err != nil {
@@ -33,6 +45,9 @@ func (r *resolverService) GetSecret(ctx context.Context, c *connect.Request[reso
 		}
 		return nil, connect.NewError(connect.CodeInternal, fmt.Errorf("failed to get secret %q: %w", msgID, err))
 	}
+	if envelope.ID != id {
+		return nil, connect.NewError(connect.CodeInternal, fmt.Errorf("secret ID mismatch: expected %q, got %q", id, envelope.ID))
+	}
 	return connect.NewResponse(resolverv1.GetSecretResponse_builder{
 		SecretId:    proto.String(envelope.ID.String()),
 		SecretValue: proto.String(string(envelope.Value)),

plugin/resolver_test.go

@@ -0,0 +1,149 @@
+package plugin
+
+import (
+	"context"
+	"errors"
+	"testing"
+	"time"
+
+	"connectrpc.com/connect"
+	"github.com/stretchr/testify/assert"
+	"google.golang.org/protobuf/proto"
+
+	resolverv1 "github.com/docker/secrets-engine/pkg/api/resolver/v1"
+	"github.com/docker/secrets-engine/pkg/secrets"
+)
+
+const (
+	mockSecretValue = "mockSecretValue"
+	mockSecretID    = secrets.ID("mockSecretID")
+)
+
+type mockResolver struct {
+	t     *testing.T
+	id    secrets.ID
+	value string
+	err   error
+}
+
+func newMockResolver(t *testing.T, options ...mockResolverOption) *mockResolver {
+	resolver := &mockResolver{
+		t:     t,
+		id:    mockSecretID,
+		value: mockSecretValue,
+	}
+	for _, opt := range options {
+		resolver = opt(resolver)
+	}
+	return resolver
+}
+
+type mockResolverOption func(*mockResolver) *mockResolver
+
+func withMockResolverID(id secrets.ID) mockResolverOption {
+	return func(m *mockResolver) *mockResolver {
+		m.id = id
+		return m
+	}
+}
+
+func withMockResolverError(err error) mockResolverOption {
+	return func(m *mockResolver) *mockResolver {
+		m.err = err
+		return m
+	}
+}
+
+func (m mockResolver) GetSecret(_ context.Context, request secrets.Request) (secrets.Envelope, error) {
+	if request.ID != m.id {
+		return secrets.Envelope{}, errors.New("unexpected secret ID")
+	}
+	if m.err != nil {
+		return secrets.Envelope{}, m.err
+	}
+	return secrets.Envelope{
+		ID:    m.id,
+		Value: []byte(m.value),
+	}, nil
+}
+
+func TestResolverService_GetSecret(t *testing.T) {
+	t.Parallel()
+	tests := []struct {
+		name string
+		test func(t *testing.T)
+	}{
+		{
+			name: "does not resolve secrets before setup completed",
+			test: func(t *testing.T) {
+				s := &resolverService{resolver: newMockResolver(t)}
+				_, err := s.GetSecret(t.Context(), newGetSecretRequest(mockSecretID))
+				assert.ErrorContains(t, err, "registration incomplete (timeout ")
+			},
+		},
+		{
+			name: "returns an error if request secret ID is invalid",
+			test: func(t *testing.T) {
+				done := make(chan struct{})
+				close(done)
+				s := &resolverService{resolver: newMockResolver(t), setupCompleted: done, registrationTimeout: 10 * time.Second}
+				_, err := s.GetSecret(t.Context(), newGetSecretRequest("/"))
+				assert.ErrorContains(t, err, "invalid secret ID")
+			},
+		},
+		{
+			name: "secret not found",
+			test: func(t *testing.T) {
+				done := make(chan struct{})
+				close(done)
+				s := &resolverService{resolver: newMockResolver(t, withMockResolverError(secrets.ErrNotFound)), setupCompleted: done, registrationTimeout: 10 * time.Second}
+				_, err := s.GetSecret(t.Context(), newGetSecretRequest(mockSecretID))
+				assert.ErrorIs(t, err, secrets.ErrNotFound)
+			},
+		},
+		{
+			name: "error fetching secret",
+			test: func(t *testing.T) {
+				done := make(chan struct{})
+				close(done)
+				s := &resolverService{resolver: newMockResolver(t, withMockResolverError(errors.New("foo"))), setupCompleted: done, registrationTimeout: 10 * time.Second}
+				_, err := s.GetSecret(t.Context(), newGetSecretRequest(mockSecretID))
+				assert.ErrorContains(t, err, "foo")
+			},
+		},
+		{
+			name: "returns wrong ID",
+			test: func(t *testing.T) {
+				done := make(chan struct{})
+				close(done)
+				s := &resolverService{resolver: newMockResolver(t, withMockResolverID("wrongID")), setupCompleted: done, registrationTimeout: 10 * time.Second}
+				_, err := s.GetSecret(t.Context(), newGetSecretRequest(mockSecretID))
+				assert.ErrorContains(t, err, "secret ID mismatch")
+			},
+		},
+		{
+			name: "returns secret value",
+			test: func(t *testing.T) {
+				done := make(chan struct{})
+				close(done)
+				s := &resolverService{resolver: newMockResolver(t), setupCompleted: done, registrationTimeout: 10 * time.Second}
+				resp, err := s.GetSecret(t.Context(), newGetSecretRequest(mockSecretID))
+				assert.NoError(t, err)
+				assert.Equal(t, mockSecretID.String(), resp.Msg.GetSecretId())
+				assert.Equal(t, mockSecretValue, resp.Msg.GetSecretValue())
+			},
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+
+			tt.test(t)
+		})
+	}
+}
+
+func newGetSecretRequest(secretID secrets.ID) *connect.Request[resolverv1.GetSecretRequest] {
+	return connect.NewRequest(resolverv1.GetSecretRequest_builder{
+		SecretId: proto.String(string(secretID)),
+	}.Build())
+}

plugin/stub.go

@@ -93,7 +93,8 @@ func setup(ctx context.Context, conn net.Conn, name string, p Plugin, timeout ti
 		_, _ = w.Write([]byte("ok"))
 	})
 	httpMux.Handle(resolverv1connect.NewPluginServiceHandler(&pluginService{p.Shutdown}))
-	httpMux.Handle(resolverv1connect.NewResolverServiceHandler(&resolverService{p}))
+	setupCompleted := make(chan struct{})
+	httpMux.Handle(resolverv1connect.NewResolverServiceHandler(&resolverService{p, setupCompleted, timeout}))
 	ipc, c, err := ipc.NewPluginIPC(conn, httpMux, func(err error) {
 		if errors.Is(err, io.EOF) {
 			logrus.Infof("Plugin runtime stopped, plugin %s is shutting down...", name)
@@ -114,6 +115,7 @@ func setup(ctx context.Context, conn net.Conn, name string, p Plugin, timeout ti
 		return nil, fmt.Errorf("failed to configure plugin %q: %w", name, err)
 	}
 	logrus.Infof("Started plugin %s...", name)
+	close(setupCompleted)
 	return ipc, nil
 }