forked from TrueCloudLab/distribution
109 lines
3.1 KiB
Go
109 lines
3.1 KiB
Go
package runtime
|
|
|
|
import (
|
|
"errors"
|
|
"mime"
|
|
"net/http"
|
|
|
|
"google.golang.org/grpc/grpclog"
|
|
"google.golang.org/protobuf/encoding/protojson"
|
|
)
|
|
|
|
// MIMEWildcard is the fallback MIME type used for requests which do not match
|
|
// a registered MIME type.
|
|
const MIMEWildcard = "*"
|
|
|
|
var (
|
|
acceptHeader = http.CanonicalHeaderKey("Accept")
|
|
contentTypeHeader = http.CanonicalHeaderKey("Content-Type")
|
|
|
|
defaultMarshaler = &HTTPBodyMarshaler{
|
|
Marshaler: &JSONPb{
|
|
MarshalOptions: protojson.MarshalOptions{
|
|
EmitUnpopulated: true,
|
|
},
|
|
UnmarshalOptions: protojson.UnmarshalOptions{
|
|
DiscardUnknown: true,
|
|
},
|
|
},
|
|
}
|
|
)
|
|
|
|
// MarshalerForRequest returns the inbound/outbound marshalers for this request.
|
|
// It checks the registry on the ServeMux for the MIME type set by the Content-Type header.
|
|
// If it isn't set (or the request Content-Type is empty), checks for "*".
|
|
// If there are multiple Content-Type headers set, choose the first one that it can
|
|
// exactly match in the registry.
|
|
// Otherwise, it follows the above logic for "*"/InboundMarshaler/OutboundMarshaler.
|
|
func MarshalerForRequest(mux *ServeMux, r *http.Request) (inbound Marshaler, outbound Marshaler) {
|
|
for _, acceptVal := range r.Header[acceptHeader] {
|
|
if m, ok := mux.marshalers.mimeMap[acceptVal]; ok {
|
|
outbound = m
|
|
break
|
|
}
|
|
}
|
|
|
|
for _, contentTypeVal := range r.Header[contentTypeHeader] {
|
|
contentType, _, err := mime.ParseMediaType(contentTypeVal)
|
|
if err != nil {
|
|
grpclog.Infof("Failed to parse Content-Type %s: %v", contentTypeVal, err)
|
|
continue
|
|
}
|
|
if m, ok := mux.marshalers.mimeMap[contentType]; ok {
|
|
inbound = m
|
|
break
|
|
}
|
|
}
|
|
|
|
if inbound == nil {
|
|
inbound = mux.marshalers.mimeMap[MIMEWildcard]
|
|
}
|
|
if outbound == nil {
|
|
outbound = inbound
|
|
}
|
|
|
|
return inbound, outbound
|
|
}
|
|
|
|
// marshalerRegistry is a mapping from MIME types to Marshalers.
|
|
type marshalerRegistry struct {
|
|
mimeMap map[string]Marshaler
|
|
}
|
|
|
|
// add adds a marshaler for a case-sensitive MIME type string ("*" to match any
|
|
// MIME type).
|
|
func (m marshalerRegistry) add(mime string, marshaler Marshaler) error {
|
|
if len(mime) == 0 {
|
|
return errors.New("empty MIME type")
|
|
}
|
|
|
|
m.mimeMap[mime] = marshaler
|
|
|
|
return nil
|
|
}
|
|
|
|
// makeMarshalerMIMERegistry returns a new registry of marshalers.
|
|
// It allows for a mapping of case-sensitive Content-Type MIME type string to runtime.Marshaler interfaces.
|
|
//
|
|
// For example, you could allow the client to specify the use of the runtime.JSONPb marshaler
|
|
// with a "application/jsonpb" Content-Type and the use of the runtime.JSONBuiltin marshaler
|
|
// with a "application/json" Content-Type.
|
|
// "*" can be used to match any Content-Type.
|
|
// This can be attached to a ServerMux with the marshaler option.
|
|
func makeMarshalerMIMERegistry() marshalerRegistry {
|
|
return marshalerRegistry{
|
|
mimeMap: map[string]Marshaler{
|
|
MIMEWildcard: defaultMarshaler,
|
|
},
|
|
}
|
|
}
|
|
|
|
// WithMarshalerOption returns a ServeMuxOption which associates inbound and outbound
|
|
// Marshalers to a MIME type in mux.
|
|
func WithMarshalerOption(mime string, marshaler Marshaler) ServeMuxOption {
|
|
return func(mux *ServeMux) {
|
|
if err := mux.marshalers.add(mime, marshaler); err != nil {
|
|
panic(err)
|
|
}
|
|
}
|
|
}
|