Skip to content

Commit 1354711

Browse files
yuguiyugui
committed
Unify APIs of custom marshaler into SeveMuxOptions
1 parent 2c1b1b9 commit 1354711

File tree

5 files changed

+256
-270
lines changed

5 files changed

+256
-270
lines changed

‎runtime/marshaler.go‎

Lines changed: 0 additions & 141 deletions
Original file line numberDiff line numberDiff line change
@@ -2,149 +2,8 @@ package runtime
22

33
import (
44
"io"
5-
"net/http"
65
)
76

8-
const mimeWildcard = "*"
9-
10-
var (
11-
inboundDefaultMarshaler = new(JSONBuiltin)
12-
outboundDefaultMarshaler = new(JSONBuiltin)
13-
14-
contentTypeHeader = http.CanonicalHeaderKey("Content-Type")
15-
)
16-
17-
// MarshalerForRequest returns the inbound/outbound marshalers for this request.
18-
// It checks the registry on the ServeMux for the MIME type set by the Content-Type header.
19-
// If it isn't set (or the request Content-Type is empty), checks for "*".
20-
// If that isn't set, uses the ServerMux's InboundMarshaler/OutboundMarshaler.
21-
// If there are multiple Content-Type headers set, choose the first one that it can
22-
// exactly match in the registry.
23-
// Otherwise, it follows the above logic for "*"/InboundMarshaler/OutboundMarshaler.
24-
func MarshalerForRequest(mux *ServeMux, r *http.Request) (inbound Marshaler, outbound Marshaler) {
25-
headerVals := append(append([]string(nil), r.Header[contentTypeHeader]...), "*")
26-
27-
for _, val := range headerVals {
28-
m := mux.MIMERegistry.lookup(val)
29-
if m != nil {
30-
if inbound == nil {
31-
inbound = m.inbound
32-
}
33-
if outbound == nil {
34-
outbound = m.outbound
35-
}
36-
}
37-
if inbound != nil && outbound != nil {
38-
// Got them both, return
39-
return inbound, outbound
40-
}
41-
}
42-
43-
if inbound == nil {
44-
inbound = mux.InboundMarshaler
45-
}
46-
if inbound == nil {
47-
inbound = inboundDefaultMarshaler
48-
}
49-
50-
if outbound == nil {
51-
outbound = mux.OutboundMarshaler
52-
}
53-
if outbound == nil {
54-
outbound = outboundDefaultMarshaler
55-
}
56-
57-
return inbound, outbound
58-
59-
}
60-
61-
// MarshalerMIMERegistry keeps a mapping from MIME types to mimeMarshalers.
62-
type MarshalerMIMERegistry struct {
63-
mimeMap map[string]*mimeMarshaler
64-
}
65-
66-
type mimeMarshaler struct {
67-
inbound Marshaler
68-
outbound Marshaler
69-
}
70-
71-
// AddMarshaler adds an inbound and outbund marshaler for a case-sensitive MIME type string ("*" to match any MIME type).
72-
// Inbound is the marshaler that is used when marshaling inbound requests from the client.
73-
// Outbound is the marshaler that is used when marshaling outbound responses to the client.
74-
func (m *MarshalerMIMERegistry) AddMarshaler(mime string, inbound, outbound Marshaler) {
75-
76-
if len(mime) == 0 {
77-
panic("Mime can't be an empty string")
78-
}
79-
80-
m.mimeMap[mime] = &mimeMarshaler{
81-
inbound: inbound,
82-
outbound: outbound,
83-
}
84-
85-
}
86-
87-
// AddInboundMarshaler adds an inbound marshaler for a case-sensitive MIME type string ("*" to match any MIME type).
88-
// Inbound is the marshaler that is used when marshaling inbound requests from the client.
89-
func (m *MarshalerMIMERegistry) AddInboundMarshaler(mime string, inbound Marshaler) {
90-
91-
if len(mime) == 0 {
92-
panic("Mime can't be an empty string")
93-
}
94-
95-
if _, ok := m.mimeMap[mime]; ok {
96-
//Already have this mime, just change inbound
97-
m.mimeMap[mime].inbound = inbound
98-
} else {
99-
m.mimeMap[mime] = &mimeMarshaler{
100-
inbound: inbound,
101-
outbound: nil,
102-
}
103-
}
104-
105-
}
106-
107-
// AddOutboundMarshaler adds an outbund marshaler for a case-sensitive MIME type string ("*" to match any MIME type).
108-
// Outbound is the marshaler that is used when marshaling outbound responses to the client.
109-
func (m *MarshalerMIMERegistry) AddOutboundMarshaler(mime string, outbound Marshaler) {
110-
mime = http.CanonicalHeaderKey(mime)
111-
if len(mime) == 0 {
112-
panic("Mime can't be an empty string")
113-
}
114-
115-
if _, ok := m.mimeMap[mime]; ok {
116-
//Already have this mime, just change outbound
117-
m.mimeMap[mime].outbound = outbound
118-
} else {
119-
m.mimeMap[mime] = &mimeMarshaler{
120-
inbound: nil,
121-
outbound: outbound,
122-
}
123-
}
124-
125-
}
126-
127-
func (m *MarshalerMIMERegistry) lookup(mime string) *mimeMarshaler {
128-
if m == nil {
129-
return nil
130-
}
131-
return m.mimeMap[mime]
132-
}
133-
134-
// NewMarshalerMIMERegistry returns a new registry of marshalers.
135-
// It allows for a mapping of case-sensitive Content-Type MIME type string to runtime.Marshaler interfaces.
136-
//
137-
// For example, you could allow the client to specify the use of the runtime.JSONPb marshaler
138-
// with a "applicaton/jsonpb" Content-Type and the use of the runtime.JSONBuiltin marshaler
139-
// with a "application/json" Content-Type.
140-
// "*" can be used to match any Content-Type.
141-
// This can be attached to a ServerMux with the MIMERegistry option.
142-
func NewMarshalerMIMERegistry() *MarshalerMIMERegistry {
143-
return &MarshalerMIMERegistry{
144-
mimeMap: make(map[string]*mimeMarshaler),
145-
}
146-
}
147-
1487
// Marshaler defines a conversion between byte sequence and gRPC payloads / fields.
1498
type Marshaler interface {
1509
// Marshal marshals "v" into byte sequence.

‎runtime/marshaler_registry.go‎

Lines changed: 138 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,138 @@
1+
package runtime
2+
3+
import (
4+
"errors"
5+
"net/http"
6+
)
7+
8+
const mimeWildcard = "*"
9+
10+
var (
11+
inboundDefaultMarshaler = new(JSONBuiltin)
12+
outboundDefaultMarshaler = new(JSONBuiltin)
13+
14+
contentTypeHeader = http.CanonicalHeaderKey("Content-Type")
15+
)
16+
17+
// MarshalerForRequest returns the inbound/outbound marshalers for this request.
18+
// It checks the registry on the ServeMux for the MIME type set by the Content-Type header.
19+
// If it isn't set (or the request Content-Type is empty), checks for "*".
20+
// If there are multiple Content-Type headers set, choose the first one that it can
21+
// exactly match in the registry.
22+
// Otherwise, it follows the above logic for "*"/InboundMarshaler/OutboundMarshaler.
23+
func MarshalerForRequest(mux *ServeMux, r *http.Request) (inbound Marshaler, outbound Marshaler) {
24+
headerVals := append(append([]string(nil), r.Header[contentTypeHeader]...), "*")
25+
26+
for _, val := range headerVals {
27+
m := mux.marshalers.lookup(val)
28+
if m != nil {
29+
if inbound == nil {
30+
inbound = m.inbound
31+
}
32+
if outbound == nil {
33+
outbound = m.outbound
34+
}
35+
}
36+
if inbound != nil && outbound != nil {
37+
// Got them both, return
38+
return inbound, outbound
39+
}
40+
}
41+
if inbound == nil {
42+
inbound = inboundDefaultMarshaler
43+
}
44+
if outbound == nil {
45+
outbound = outboundDefaultMarshaler
46+
}
47+
return inbound, outbound
48+
49+
}
50+
51+
// marshalerRegistry keeps a mapping from MIME types to mimeMarshalers.
52+
type marshalerRegistry map[string]*mimeMarshaler
53+
54+
type mimeMarshaler struct {
55+
inbound Marshaler
56+
outbound Marshaler
57+
}
58+
59+
// addMarshaler adds an inbound and outbund marshaler for a case-sensitive MIME type string ("*" to match any MIME type).
60+
// Inbound is the marshaler that is used when marshaling inbound requests from the client.
61+
// Outbound is the marshaler that is used when marshaling outbound responses to the client.
62+
func (r *marshalerRegistry) add(mime string, inbound, outbound Marshaler) error {
63+
if mime == "" {
64+
return errors.New("empty MIME type")
65+
}
66+
(*r)[mime] = &mimeMarshaler{
67+
inbound: inbound,
68+
outbound: outbound,
69+
}
70+
return nil
71+
}
72+
73+
// addInboundMarshaler adds an inbound marshaler for a case-sensitive MIME type string ("*" to match any MIME type).
74+
// Inbound is the marshaler that is used when marshaling inbound requests from the client.
75+
func (r *marshalerRegistry) addInbound(mime string, inbound Marshaler) error {
76+
if mime == "" {
77+
return errors.New("empty MIME type")
78+
}
79+
if entry := (*r)[mime]; entry != nil {
80+
entry.inbound = inbound
81+
return nil
82+
}
83+
(*r)[mime] = &mimeMarshaler{inbound: inbound}
84+
return nil
85+
}
86+
87+
// addOutBound adds an outbund marshaler for a case-sensitive MIME type string ("*" to match any MIME type).
88+
// Outbound is the marshaler that is used when marshaling outbound responses to the client.
89+
func (r *marshalerRegistry) addOutbound(mime string, outbound Marshaler) error {
90+
mime = http.CanonicalHeaderKey(mime)
91+
if mime == "" {
92+
return errors.New("empty MIME type")
93+
}
94+
if entry := (*r)[mime]; entry != nil {
95+
entry.outbound = outbound
96+
return nil
97+
}
98+
(*r)[mime] = &mimeMarshaler{outbound: outbound}
99+
return nil
100+
101+
}
102+
103+
func (r *marshalerRegistry) lookup(mime string) *mimeMarshaler {
104+
if r == nil {
105+
return nil
106+
}
107+
return (*r)[mime]
108+
}
109+
110+
// WithMarshalerOption returns a ServeMuxOption which associates inbound and outbound
111+
// Marshalers to a MIME type in mux.
112+
func WithMarshalerOption(mime string, in, out Marshaler) ServeMuxOption {
113+
return func(mux *ServeMux) {
114+
if err := mux.marshalers.add(mime, in, out); err != nil {
115+
panic(err)
116+
}
117+
}
118+
}
119+
120+
// WithInboundMarshalerOption returns a ServeMuxOption which associates an inbound
121+
// Marshaler to a MIME type in mux.
122+
func WithInboundMarshalerOption(mime string, in Marshaler) ServeMuxOption {
123+
return func(mux *ServeMux) {
124+
if err := mux.marshalers.addInbound(mime, in); err != nil {
125+
panic(err)
126+
}
127+
}
128+
}
129+
130+
// WithOutboundMarshalerOption returns a ServeMuxOption which associates an outbound
131+
// Marshaler to a MIME type in mux.
132+
func WithOutboundMarshalerOption(mime string, out Marshaler) ServeMuxOption {
133+
return func(mux *ServeMux) {
134+
if err := mux.marshalers.addOutbound(mime, out); err != nil {
135+
panic(err)
136+
}
137+
}
138+
}

0 commit comments

Comments
 (0)