github.com/cockroachdb/cockroach@v20.2.0-alpha.1+incompatible/pkg/util/protoutil/marshaler.go (about)

     1  // Copyright 2016 The Cockroach Authors.
     2  //
     3  // Use of this software is governed by the Business Source License
     4  // included in the file licenses/BSL.txt.
     5  //
     6  // As of the Change Date specified in that file, in accordance with
     7  // the Business Source License, use of this software will be governed
     8  // by the Apache License, Version 2.0, included in the file
     9  // licenses/APL.txt.
    10  
    11  package protoutil
    12  
    13  import (
    14  	"io"
    15  	"io/ioutil"
    16  
    17  	"github.com/cockroachdb/errors"
    18  	"github.com/gogo/protobuf/proto"
    19  	gwruntime "github.com/grpc-ecosystem/grpc-gateway/runtime"
    20  )
    21  
    22  var _ gwruntime.Marshaler = (*ProtoPb)(nil)
    23  
    24  // ProtoPb is a gwruntime.Marshaler that uses github.com/gogo/protobuf/proto.
    25  type ProtoPb struct{}
    26  
    27  // ContentType implements gwruntime.Marshaler.
    28  func (*ProtoPb) ContentType() string {
    29  	// NB: This is the same as httputil.ProtoContentType which we can't use due
    30  	// to an import cycle.
    31  	const ProtoContentType = "application/x-protobuf"
    32  	return ProtoContentType
    33  }
    34  
    35  // Marshal implements gwruntime.Marshaler.
    36  func (*ProtoPb) Marshal(v interface{}) ([]byte, error) {
    37  	// NB: we use proto.Message here because grpc-gateway passes us protos that
    38  	// we don't control and thus don't implement protoutil.Message.
    39  	if p, ok := v.(proto.Message); ok {
    40  		return proto.Marshal(p)
    41  	}
    42  	return nil, errors.Errorf("unexpected type %T does not implement %s", v, typeProtoMessage)
    43  }
    44  
    45  // Unmarshal implements gwruntime.Marshaler.
    46  func (*ProtoPb) Unmarshal(data []byte, v interface{}) error {
    47  	// NB: we use proto.Message here because grpc-gateway passes us protos that
    48  	// we don't control and thus don't implement protoutil.Message.
    49  	if p, ok := v.(proto.Message); ok {
    50  		return proto.Unmarshal(data, p)
    51  	}
    52  	return errors.Errorf("unexpected type %T does not implement %s", v, typeProtoMessage)
    53  }
    54  
    55  // NewDecoder implements gwruntime.Marshaler.
    56  func (*ProtoPb) NewDecoder(r io.Reader) gwruntime.Decoder {
    57  	return gwruntime.DecoderFunc(func(v interface{}) error {
    58  		// NB: we use proto.Message here because grpc-gateway passes us protos that
    59  		// we don't control and thus don't implement protoutil.Message.
    60  		if p, ok := v.(proto.Message); ok {
    61  			bytes, err := ioutil.ReadAll(r)
    62  			if err == nil {
    63  				err = proto.Unmarshal(bytes, p)
    64  			}
    65  			return err
    66  		}
    67  		return errors.Errorf("unexpected type %T does not implement %s", v, typeProtoMessage)
    68  	})
    69  }
    70  
    71  // NewEncoder implements gwruntime.Marshaler.
    72  func (*ProtoPb) NewEncoder(w io.Writer) gwruntime.Encoder {
    73  	return gwruntime.EncoderFunc(func(v interface{}) error {
    74  		// NB: we use proto.Message here because grpc-gateway passes us protos that
    75  		// we don't control and thus don't implement protoutil.Message.
    76  		if p, ok := v.(proto.Message); ok {
    77  			bytes, err := proto.Marshal(p)
    78  			if err == nil {
    79  				_, err = w.Write(bytes)
    80  			}
    81  			return err
    82  		}
    83  		return errors.Errorf("unexpected type %T does not implement %s", v, typeProtoMessage)
    84  	})
    85  }
    86  
    87  var _ gwruntime.Delimited = (*ProtoPb)(nil)
    88  
    89  // Delimiter implements gwruntime.Delimited.
    90  func (*ProtoPb) Delimiter() []byte {
    91  	return nil
    92  }