Update to new goprotobuf

codec/codec.go

@@ -28,7 +28,7 @@ type ProtoReader interface {
 
 // ReadProto reads a uvarint size and then a protobuf from r.
 // If the size read is zero, nothing more is read.
-func ReadProto(r ProtoReader, pb interface{}) error {
+func ReadProto(r ProtoReader, pb proto.Message) error {
 	size, err := binary.ReadUvarint(r)
 	if err != nil {
 		return err
@@ -44,7 +44,7 @@ func ReadProto(r ProtoReader, pb interface{}) error {
 // WriteProto writes a uvarint size and then a protobuf to w.
 // If the data takes no space (like rpc.InvalidRequest),
 // only a zero size is written.
-func WriteProto(w io.Writer, pb interface{}) error {
+func WriteProto(w io.Writer, pb proto.Message) error {
 	// Allocate enough space for the biggest uvarint
 	var size [binary.MaxVarintLen64]byte
 
@@ -93,15 +93,25 @@ func (s *ServerCodec) ReadRequestHeader(req *rpc.Request) error {
 // ReadRequestBody reads a uvarint from the connection and decodes that many
 // subsequent bytes into the given protobuf (which should be a pointer to a
 // struct that is generated by the proto package).
-func (s *ServerCodec) ReadRequestBody(pb interface{}) error {
+func (s *ServerCodec) ReadRequestBody(obj interface{}) error {
+	pb, ok := obj.(proto.Message)
+	if !ok {
+		return fmt.Errorf("%T does not implement proto.Message", obj)
+	}
+
 	return ReadProto(s.r, pb)
 }
 
 // WriteResponse writes the appropriate header protobuf and the given protobuf
 // to the connection (each prefixed with a uvarint indicating its size).  If
 // the response was invalid, the size of the body of the resp is reported as
 // having size zero and is not sent.
-func (s *ServerCodec) WriteResponse(resp *rpc.Response, pb interface{}) error {
+func (s *ServerCodec) WriteResponse(resp *rpc.Response, obj interface{}) error {
+	pb, ok := obj.(proto.Message)
+	if !ok {
+		return fmt.Errorf("%T does not implement proto.Message", obj)
+	}
+
 	// Write the header
 	header := wire.Header{
 		Method: &resp.ServiceMethod,
@@ -140,7 +150,12 @@ func NewClientCodec(conn net.Conn) *ClientCodec {
 
 // WriteRequest writes the appropriate header protobuf and the given protobuf
 // to the connection (each prefixed with a uvarint indicating its size).
-func (c *ClientCodec) WriteRequest(req *rpc.Request, pb interface{}) error {
+func (c *ClientCodec) WriteRequest(req *rpc.Request, obj interface{}) error {
+	pb, ok := obj.(proto.Message)
+	if !ok {
+		return fmt.Errorf("%T does not implement proto.Message", obj)
+	}
+
 	// Write the header
 	header := wire.Header{
 		Method: &req.ServiceMethod,
@@ -180,7 +195,12 @@ func (c *ClientCodec) ReadResponseHeader(resp *rpc.Response) error {
 // struct that is generated by the proto package).  If the uvarint size read
 // is zero, nothing is done (this indicates an error condition, which was
 // encapsulated in the header)
-func (c *ClientCodec) ReadResponseBody(pb interface{}) error {
+func (c *ClientCodec) ReadResponseBody(obj interface{}) error {
+	pb, ok := obj.(proto.Message)
+	if !ok {
+		return fmt.Errorf("%T does not implement proto.Message", obj)
+	}
+
 	return ReadProto(c.r, pb)
 }
 

codec/codec_test.go

@@ -11,11 +11,15 @@ import (
 
 type InvalidRequest struct{}
 
+func (InvalidRequest) ProtoMessage()  {}
+func (InvalidRequest) Reset()         {}
+func (InvalidRequest) String() string { return "" }
+
 func TestWriteProto(t *testing.T) {
 	tests := []struct {
 		Desc string
-		Out  interface{}
-		In   interface{}
+		Out  proto.Message
+		In   proto.Message
 	}{
 		{
 			Desc: "Test the zero size proto",