fix: Race when writing to a closed pipe

kylecarbs · kylecarbs · commit 8e61caced3d9 · 2022-05-31T15:02:55.000Z
This is such an intermittent race it's difficult to track,
but regardless this is an improvement to the code.
diff --git a/provisionersdk/serve.go b/provisionersdk/serve.go
@@ -31,9 +31,9 @@ func Serve(ctx context.Context, server proto.DRPCProvisionerServer, options *Ser
 	if options.Listener == nil {
 		config := yamux.DefaultConfig()
 		config.LogOutput = io.Discard
-		stdio, err := yamux.Server(readWriteCloser{
-			ReadCloser: os.Stdin,
-			Writer:     os.Stdout,
+		stdio, err := yamux.Server(&readWriteCloser{
+			ReadCloser:  os.Stdin,
+			WriteCloser: os.Stdout,
 		}, config)
 		if err != nil {
 			return xerrors.Errorf("create yamux: %w", err)
@@ -54,6 +54,9 @@ func Serve(ctx context.Context, server proto.DRPCProvisionerServer, options *Ser
 	// short-lived processes that can be executed concurrently.
 	err = srv.Serve(ctx, options.Listener)
 	if err != nil {
+		if errors.Is(err, io.EOF) {
+			return nil
+		}
 		if errors.Is(err, context.Canceled) {
 			return nil
 		}
diff --git a/provisionersdk/transport.go b/provisionersdk/transport.go
@@ -3,6 +3,7 @@ package provisionersdk
 import (
 	"context"
 	"io"
+	"sync"
 
 	"github.com/hashicorp/yamux"
 	"storj.io/drpc"
@@ -22,16 +23,16 @@ func TransportPipe() (*yamux.Session, *yamux.Session) {
 	yamuxConfig := yamux.DefaultConfig()
 	yamuxConfig.LogOutput = io.Discard
 	client, err := yamux.Client(&readWriteCloser{
-		ReadCloser: clientReader,
-		Writer:     serverWriter,
+		ReadCloser:  clientReader,
+		WriteCloser: serverWriter,
 	}, yamuxConfig)
 	if err != nil {
 		panic(err)
 	}
 
 	server, err := yamux.Server(&readWriteCloser{
-		ReadCloser: serverReader,
-		Writer:     clientWriter,
+		ReadCloser:  serverReader,
+		WriteCloser: clientWriter,
 	}, yamuxConfig)
 	if err != nil {
 		panic(err)
@@ -45,8 +46,19 @@ func Conn(session *yamux.Session) drpc.Conn {
 }
 
 type readWriteCloser struct {
+	closeMutex sync.Mutex
 	io.ReadCloser
-	io.Writer
+	io.WriteCloser
+}
+
+func (c *readWriteCloser) Close() error {
+	c.closeMutex.Lock()
+	defer c.closeMutex.Unlock()
+	err := c.ReadCloser.Close()
+	if err != nil {
+		return err
+	}
+	return c.WriteCloser.Close()
 }
 
 // Allows concurrent requests on a single dRPC connection.
