marshaler_registry.go

Documentation: github.com/grpc-ecosystem/grpc-gateway/v2/runtime

     1  package runtime
     2  
     3  import (
     4  	"errors"
     5  	"mime"
     6  	"net/http"
     7  
     8  	"google.golang.org/grpc/grpclog"
     9  	"google.golang.org/protobuf/encoding/protojson"
    10  )
    11  
    12  // MIMEWildcard is the fallback MIME type used for requests which do not match
    13  // a registered MIME type.
    14  const MIMEWildcard = "*"
    15  
    16  var (
    17  	acceptHeader      = http.CanonicalHeaderKey("Accept")
    18  	contentTypeHeader = http.CanonicalHeaderKey("Content-Type")
    19  
    20  	defaultMarshaler = &HTTPBodyMarshaler{
    21  		Marshaler: &JSONPb{
    22  			MarshalOptions: protojson.MarshalOptions{
    23  				EmitUnpopulated: true,
    24  			},
    25  			UnmarshalOptions: protojson.UnmarshalOptions{
    26  				DiscardUnknown: true,
    27  			},
    28  		},
    29  	}
    30  )
    31  
    32  // MarshalerForRequest returns the inbound/outbound marshalers for this request.
    33  // It checks the registry on the ServeMux for the MIME type set by the Content-Type header.
    34  // If it isn't set (or the request Content-Type is empty), checks for "*".
    35  // If there are multiple Content-Type headers set, choose the first one that it can
    36  // exactly match in the registry.
    37  // Otherwise, it follows the above logic for "*"/InboundMarshaler/OutboundMarshaler.
    38  func MarshalerForRequest(mux *ServeMux, r *http.Request) (inbound Marshaler, outbound Marshaler) {
    39  	for _, acceptVal := range r.Header[acceptHeader] {
    40  		if m, ok := mux.marshalers.mimeMap[acceptVal]; ok {
    41  			outbound = m
    42  			break
    43  		}
    44  	}
    45  
    46  	for _, contentTypeVal := range r.Header[contentTypeHeader] {
    47  		contentType, _, err := mime.ParseMediaType(contentTypeVal)
    48  		if err != nil {
    49  			grpclog.Infof("Failed to parse Content-Type %s: %v", contentTypeVal, err)
    50  			continue
    51  		}
    52  		if m, ok := mux.marshalers.mimeMap[contentType]; ok {
    53  			inbound = m
    54  			break
    55  		}
    56  	}
    57  
    58  	if inbound == nil {
    59  		inbound = mux.marshalers.mimeMap[MIMEWildcard]
    60  	}
    61  	if outbound == nil {
    62  		outbound = inbound
    63  	}
    64  
    65  	return inbound, outbound
    66  }
    67  
    68  // marshalerRegistry is a mapping from MIME types to Marshalers.
    69  type marshalerRegistry struct {
    70  	mimeMap map[string]Marshaler
    71  }
    72  
    73  // add adds a marshaler for a case-sensitive MIME type string ("*" to match any
    74  // MIME type).
    75  func (m marshalerRegistry) add(mime string, marshaler Marshaler) error {
    76  	if len(mime) == 0 {
    77  		return errors.New("empty MIME type")
    78  	}
    79  
    80  	m.mimeMap[mime] = marshaler
    81  
    82  	return nil
    83  }
    84  
    85  // makeMarshalerMIMERegistry returns a new registry of marshalers.
    86  // It allows for a mapping of case-sensitive Content-Type MIME type string to runtime.Marshaler interfaces.
    87  //
    88  // For example, you could allow the client to specify the use of the runtime.JSONPb marshaler
    89  // with a "application/jsonpb" Content-Type and the use of the runtime.JSONBuiltin marshaler
    90  // with a "application/json" Content-Type.
    91  // "*" can be used to match any Content-Type.
    92  // This can be attached to a ServerMux with the marshaler option.
    93  func makeMarshalerMIMERegistry() marshalerRegistry {
    94  	return marshalerRegistry{
    95  		mimeMap: map[string]Marshaler{
    96  			MIMEWildcard: defaultMarshaler,
    97  		},
    98  	}
    99  }
   100  
   101  // WithMarshalerOption returns a ServeMuxOption which associates inbound and outbound
   102  // Marshalers to a MIME type in mux.
   103  func WithMarshalerOption(mime string, marshaler Marshaler) ServeMuxOption {
   104  	return func(mux *ServeMux) {
   105  		if err := mux.marshalers.add(mime, marshaler); err != nil {
   106  			panic(err)
   107  		}
   108  	}
   109  }
   110
View as plain text