marshaler_registry.go 3.1 KB

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