marshaler_registry.go

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

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