Merge pull request #35 from smartcontractkit/fix/potential-write-dead…

…lock

Fix potential deadlock in client & server writes

client.go

@@ -280,7 +280,7 @@ func (cc *ClientConn) handleMessageRequest(r *message.Request) {
 		cc.conn.mu.RUnlock()
 		cc.mu.RUnlock()
 
-		if err := tr.Write(replyMsg); err != nil {
+		if err := tr.Write(ctx, replyMsg); err != nil {
 			cc.dopts.logger.Errorf("error writing to transport: %s", err)
 		}
 	}
@@ -376,7 +376,7 @@ func (cc *ClientConn) Invoke(ctx context.Context, method string, args interface{
 	cc.conn.mu.RUnlock()
 	cc.mu.RUnlock()
 
-	if err := tr.Write(reqB); err != nil {
+	if err := tr.Write(ctx, reqB); err != nil {
 		return err
 	}
 

internal/transport/transport.go

@@ -36,7 +36,7 @@ type ClientTransport interface {
 	Read() <-chan []byte
 
 	// Write sends a message to the stream.
-	Write(msg []byte) error
+	Write(ctx context.Context, msg []byte) error
 
 	// Close tears down this transport. Once it returns, the transport
 	// should not be accessed any more.
@@ -73,7 +73,7 @@ type ServerTransport interface {
 	Read() <-chan []byte
 
 	// Write sends a message to the stream.
-	Write(msg []byte) error
+	Write(ctx context.Context, msg []byte) error
 
 	// Close tears down the transport. Once it is called, the transport
 	// should not be accessed any more.

internal/transport/websocket_client.go

@@ -78,10 +78,17 @@ func (c *WebsocketClient) Read() <-chan []byte {
 }
 
 // Write writes a message the websocket connection.
-func (c *WebsocketClient) Write(msg []byte) error {
-	c.write <- msg
-
-	return nil
+func (c *WebsocketClient) Write(ctx context.Context, msg []byte) error {
+	select {
+	case <-c.done:
+		return fmt.Errorf("[wsrpc] could not write message, websocket is closed")
+	case <-c.interrupt:
+		return fmt.Errorf("[wsrpc] could not write message, transport is closed")
+	case <-ctx.Done():
+		return fmt.Errorf("[wsrpc] could not write message, context is done")
+	case c.write <- msg:
+		return nil
+	}
 }
 
 // Close closes the websocket connection and cleans up pump goroutines.

internal/transport/websocket_server.go

@@ -1,6 +1,8 @@
 package transport
 
 import (
+	"context"
+	"fmt"
 	"log"
 	"sync"
 	"time"
@@ -61,11 +63,17 @@ func (s *WebsocketServer) Read() <-chan []byte {
 }
 
 // Write writes a message the websocket connection.
-func (s *WebsocketServer) Write(msg []byte) error {
-	// Send the message to the channel
-	s.write <- msg
-
-	return nil
+func (s *WebsocketServer) Write(ctx context.Context, msg []byte) error {
+	select {
+	case <-s.done:
+		return fmt.Errorf("[wsrpc] could not write message, websocket is closed")
+	case <-s.interrupt:
+		return fmt.Errorf("[wsrpc] could not write message, transport is closed")
+	case <-ctx.Done():
+		return fmt.Errorf("[wsrpc] could not write message, context is done")
+	case s.write <- msg:
+		return nil
+	}
 }
 
 // Close closes the websocket connection and cleans up pump goroutines. Notifies

server.go

@@ -177,7 +177,7 @@ func (s *Server) wshandler(w http.ResponseWriter, r *http.Request) {
 }
 
 // sendMsg writes the message to the connection which matches the public key.
-func (s *Server) sendMsg(pub [32]byte, msg []byte) error {
+func (s *Server) sendMsg(ctx context.Context, pub [32]byte, msg []byte) error {
 	// Find the transport matching the public key
 	s.mu.RLock()
 	tr, err := s.connMgr.getTransport(pub)
@@ -186,7 +186,7 @@ func (s *Server) sendMsg(pub [32]byte, msg []byte) error {
 		return err
 	}
 
-	return tr.Write(msg)
+	return tr.Write(ctx, msg)
 }
 
 // handleRead listens to the transport read channel and passes the message to the
@@ -248,7 +248,7 @@ func (s *Server) handleMessageRequest(pubKey credentials.StaticSizedPublicKey, r
 			return
 		}
 
-		if err := s.sendMsg(pubKey, replyMsg); err != nil {
+		if err := s.sendMsg(ctx, pubKey, replyMsg); err != nil {
 			log.Printf("error sending message: %s", err)
 		}
 	}
@@ -314,7 +314,7 @@ func (s *Server) Invoke(ctx context.Context, method string, args interface{}, re
 	}
 	pubKey := p.PublicKey
 
-	if err = s.sendMsg(pubKey, req); err != nil {
+	if err = s.sendMsg(ctx, pubKey, req); err != nil {
 		return err
 	}