github.com/cloudbase/juju-core@v0.0.0-20140504232958-a7271ac7912f/rpc/jsoncodec/codec.go (about)

     1  // The jsoncodec package provides a JSON codec for the rpc package.
     2  package jsoncodec
     3  
     4  import (
     5  	"encoding/json"
     6  	"fmt"
     7  	"io"
     8  	"sync"
     9  	"sync/atomic"
    10  
    11  	"github.com/juju/loggo"
    12  
    13  	"launchpad.net/juju-core/rpc"
    14  )
    15  
    16  var logger = loggo.GetLogger("juju.rpc.jsoncodec")
    17  
    18  // JSONConn sends and receives messages to an underlying connection
    19  // in JSON format.
    20  type JSONConn interface {
    21  	// Send sends a message.
    22  	Send(msg interface{}) error
    23  	// Receive receives a message into msg.
    24  	Receive(msg interface{}) error
    25  	Close() error
    26  }
    27  
    28  // Codec implements rpc.Codec for a connection.
    29  type Codec struct {
    30  	// msg holds the message that's just been read by ReadHeader, so
    31  	// that the body can be read by ReadBody.
    32  	msg         inMsg
    33  	conn        JSONConn
    34  	logMessages int32
    35  	mu          sync.Mutex
    36  	closing     bool
    37  }
    38  
    39  // New returns an rpc codec that uses conn to send and receive
    40  // messages.
    41  func New(conn JSONConn) *Codec {
    42  	return &Codec{
    43  		conn: conn,
    44  	}
    45  }
    46  
    47  // SetLogging sets whether messages will be logged
    48  // by the codec.
    49  func (c *Codec) SetLogging(on bool) {
    50  	val := int32(0)
    51  	if on {
    52  		val = 1
    53  	}
    54  	atomic.StoreInt32(&c.logMessages, val)
    55  }
    56  
    57  func (c *Codec) isLogging() bool {
    58  	return atomic.LoadInt32(&c.logMessages) != 0
    59  }
    60  
    61  // inMsg holds an incoming message.  We don't know the type of the
    62  // parameters or response yet, so we delay parsing by storing them
    63  // in a RawMessage.
    64  type inMsg struct {
    65  	RequestId uint64
    66  	Type      string
    67  	Id        string
    68  	Request   string
    69  	Params    json.RawMessage
    70  	Error     string
    71  	ErrorCode string
    72  	Response  json.RawMessage
    73  }
    74  
    75  // outMsg holds an outgoing message.
    76  type outMsg struct {
    77  	RequestId uint64
    78  	Type      string      `json:",omitempty"`
    79  	Id        string      `json:",omitempty"`
    80  	Request   string      `json:",omitempty"`
    81  	Params    interface{} `json:",omitempty"`
    82  	Error     string      `json:",omitempty"`
    83  	ErrorCode string      `json:",omitempty"`
    84  	Response  interface{} `json:",omitempty"`
    85  }
    86  
    87  func (c *Codec) Close() error {
    88  	c.mu.Lock()
    89  	c.closing = true
    90  	c.mu.Unlock()
    91  	return c.conn.Close()
    92  }
    93  
    94  func (c *Codec) isClosing() bool {
    95  	c.mu.Lock()
    96  	defer c.mu.Unlock()
    97  	return c.closing
    98  }
    99  
   100  func (c *Codec) ReadHeader(hdr *rpc.Header) error {
   101  	c.msg = inMsg{} // avoid any potential cross-message contamination.
   102  	var err error
   103  	if c.isLogging() {
   104  		var m json.RawMessage
   105  		err = c.conn.Receive(&m)
   106  		if err == nil {
   107  			logger.Tracef("<- %s", m)
   108  			err = json.Unmarshal(m, &c.msg)
   109  		} else {
   110  			logger.Tracef("<- error: %v (closing %v)", err, c.isClosing())
   111  		}
   112  	} else {
   113  		err = c.conn.Receive(&c.msg)
   114  	}
   115  	if err != nil {
   116  		// If we've closed the connection, we may get a spurious error,
   117  		// so ignore it.
   118  		if c.isClosing() || err == io.EOF {
   119  			return io.EOF
   120  		}
   121  		return fmt.Errorf("error receiving message: %v", err)
   122  	}
   123  	hdr.RequestId = c.msg.RequestId
   124  	hdr.Request = rpc.Request{
   125  		Type:   c.msg.Type,
   126  		Id:     c.msg.Id,
   127  		Action: c.msg.Request,
   128  	}
   129  	hdr.Error = c.msg.Error
   130  	hdr.ErrorCode = c.msg.ErrorCode
   131  	return nil
   132  }
   133  
   134  func (c *Codec) ReadBody(body interface{}, isRequest bool) error {
   135  	if body == nil {
   136  		return nil
   137  	}
   138  	var rawBody json.RawMessage
   139  	if isRequest {
   140  		rawBody = c.msg.Params
   141  	} else {
   142  		rawBody = c.msg.Response
   143  	}
   144  	if len(rawBody) == 0 {
   145  		// If the response or params are omitted, it's
   146  		// equivalent to an empty object.
   147  		return nil
   148  	}
   149  	return json.Unmarshal(rawBody, body)
   150  }
   151  
   152  // DumpRequest returns JSON-formatted data representing
   153  // the RPC message with the given header and body,
   154  // as it would be written by Codec.WriteMessage.
   155  // If the body cannot be marshalled as JSON, the data
   156  // will hold a JSON string describing the error.
   157  func DumpRequest(hdr *rpc.Header, body interface{}) []byte {
   158  	var m outMsg
   159  	m.init(hdr, body)
   160  	data, err := json.Marshal(&m)
   161  	if err != nil {
   162  		return []byte(fmt.Sprintf("%q", "marshal error: "+err.Error()))
   163  	}
   164  	return data
   165  }
   166  
   167  func (c *Codec) WriteMessage(hdr *rpc.Header, body interface{}) error {
   168  	var m outMsg
   169  	m.init(hdr, body)
   170  	if c.isLogging() {
   171  		data, err := json.Marshal(&m)
   172  		if err != nil {
   173  			logger.Tracef("-> marshal error: %v", err)
   174  			return err
   175  		}
   176  		logger.Tracef("-> %s", data)
   177  	}
   178  	return c.conn.Send(&m)
   179  }
   180  
   181  // init fills out the receiving outMsg with information from the given
   182  // header and body.
   183  func (m *outMsg) init(hdr *rpc.Header, body interface{}) {
   184  	m.RequestId = hdr.RequestId
   185  	m.Type = hdr.Request.Type
   186  	m.Id = hdr.Request.Id
   187  	m.Request = hdr.Request.Action
   188  	m.Error = hdr.Error
   189  	m.ErrorCode = hdr.ErrorCode
   190  	if hdr.IsRequest() {
   191  		m.Params = body
   192  	} else {
   193  		m.Response = body
   194  	}
   195  }