github.com/niedbalski/juju@v0.0.0-20190215020005-8ff100488e47/rpc/jsoncodec/conn.go (about)

     1  // Copyright 2013 Canonical Ltd.
     2  // Licensed under the AGPLv3, see LICENCE file for details.
     3  
     4  package jsoncodec
     5  
     6  import (
     7  	"encoding/json"
     8  	"io"
     9  	"sync"
    10  
    11  	"github.com/gorilla/websocket"
    12  	"github.com/juju/errors"
    13  )
    14  
    15  // NewWebsocket returns an rpc codec that uses the given websocket
    16  // connection to send and receive messages.
    17  func NewWebsocket(conn *websocket.Conn) *Codec {
    18  	return New(NewWebsocketConn(conn))
    19  }
    20  
    21  type wsJSONConn struct {
    22  	conn *websocket.Conn
    23  	// gorilla websockets can have at most one concurrent writer, and
    24  	// one concurrent reader.
    25  	writeMutex sync.Mutex
    26  	readMutex  sync.Mutex
    27  }
    28  
    29  // NewWebsocketConn returns a JSONConn implementation
    30  // that uses the given connection for transport.
    31  func NewWebsocketConn(conn *websocket.Conn) JSONConn {
    32  	return &wsJSONConn{conn: conn}
    33  }
    34  
    35  func (conn *wsJSONConn) Send(msg interface{}) error {
    36  	conn.writeMutex.Lock()
    37  	defer conn.writeMutex.Unlock()
    38  	return conn.conn.WriteJSON(msg)
    39  }
    40  
    41  func (conn *wsJSONConn) Receive(msg interface{}) error {
    42  	conn.readMutex.Lock()
    43  	defer conn.readMutex.Unlock()
    44  	// When receiving a message, if error has been closed from the other
    45  	// side, wrap with io.EOF as this is the expected error.
    46  	err := conn.conn.ReadJSON(msg)
    47  	if err != nil {
    48  		if websocket.IsCloseError(err,
    49  			websocket.CloseNormalClosure,
    50  			websocket.CloseGoingAway,
    51  			websocket.CloseNoStatusReceived,
    52  			websocket.CloseAbnormalClosure) {
    53  			err = errors.Wrap(err, io.EOF)
    54  		}
    55  	}
    56  	return err
    57  }
    58  
    59  func (conn *wsJSONConn) Close() error {
    60  	// Tell the other end we are closing.
    61  	conn.writeMutex.Lock()
    62  	conn.conn.WriteMessage(websocket.CloseMessage, []byte{})
    63  	conn.writeMutex.Unlock()
    64  	return conn.conn.Close()
    65  }
    66  
    67  // NewNet returns an rpc codec that uses the given connection
    68  // to send and receive messages.
    69  func NewNet(conn io.ReadWriteCloser) *Codec {
    70  	return New(NetJSONConn(conn))
    71  }
    72  
    73  func NetJSONConn(conn io.ReadWriteCloser) JSONConn {
    74  	return &netConn{
    75  		enc:  json.NewEncoder(conn),
    76  		dec:  json.NewDecoder(conn),
    77  		conn: conn,
    78  	}
    79  }
    80  
    81  type netConn struct {
    82  	enc  *json.Encoder
    83  	dec  *json.Decoder
    84  	conn io.ReadWriteCloser
    85  }
    86  
    87  func (conn *netConn) Send(msg interface{}) error {
    88  	return conn.enc.Encode(msg)
    89  }
    90  
    91  func (conn *netConn) Receive(msg interface{}) error {
    92  	return conn.dec.Decode(msg)
    93  }
    94  
    95  func (conn *netConn) Close() error {
    96  	return conn.conn.Close()
    97  }