mirror of
https://github.com/zeromicro/go-zero.git
synced 2026-05-07 15:10:01 +08:00
fix:issue-5110 (#5113)
This commit is contained in:
@@ -1,6 +1,7 @@
|
|||||||
package internal
|
package internal
|
||||||
|
|
||||||
import (
|
import (
|
||||||
|
"fmt"
|
||||||
"io"
|
"io"
|
||||||
"net/http"
|
"net/http"
|
||||||
|
|
||||||
@@ -12,6 +13,22 @@ import (
|
|||||||
"google.golang.org/grpc/status"
|
"google.golang.org/grpc/status"
|
||||||
)
|
)
|
||||||
|
|
||||||
|
// MetadataHeaderPrefix is the http prefix that represents custom metadata
|
||||||
|
// parameters to or from a gRPC call.
|
||||||
|
const MetadataHeaderPrefix = "Grpc-Metadata-"
|
||||||
|
|
||||||
|
// MetadataTrailerPrefix is prepended to gRPC metadata as it is converted to
|
||||||
|
// HTTP headers in a response handled by go-zero gateway
|
||||||
|
const MetadataTrailerPrefix = "Grpc-Trailer-"
|
||||||
|
|
||||||
|
func defaultOutgoingHeaderMatcher(key string) (string, bool) {
|
||||||
|
return fmt.Sprintf("%s%s", MetadataHeaderPrefix, key), true
|
||||||
|
}
|
||||||
|
|
||||||
|
func defaultOutgoingTrailerMatcher(key string) (string, bool) {
|
||||||
|
return fmt.Sprintf("%s%s", MetadataTrailerPrefix, key), true
|
||||||
|
}
|
||||||
|
|
||||||
type EventHandler struct {
|
type EventHandler struct {
|
||||||
Status *status.Status
|
Status *status.Status
|
||||||
writer io.Writer
|
writer io.Writer
|
||||||
@@ -31,9 +48,11 @@ func NewEventHandler(writer io.Writer, resolver jsonpb.AnyResolver) *EventHandle
|
|||||||
func (h *EventHandler) OnReceiveHeaders(md metadata.MD) {
|
func (h *EventHandler) OnReceiveHeaders(md metadata.MD) {
|
||||||
w, ok := h.writer.(http.ResponseWriter)
|
w, ok := h.writer.(http.ResponseWriter)
|
||||||
if ok {
|
if ok {
|
||||||
for k, v := range md {
|
for k, vs := range md {
|
||||||
for _, val := range v {
|
if h, ok := defaultOutgoingHeaderMatcher(k); ok {
|
||||||
w.Header().Add(k, val)
|
for _, v := range vs {
|
||||||
|
w.Header().Add(h, v)
|
||||||
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -48,9 +67,11 @@ func (h *EventHandler) OnReceiveResponse(message proto.Message) {
|
|||||||
func (h *EventHandler) OnReceiveTrailers(status *status.Status, md metadata.MD) {
|
func (h *EventHandler) OnReceiveTrailers(status *status.Status, md metadata.MD) {
|
||||||
w, ok := h.writer.(http.ResponseWriter)
|
w, ok := h.writer.(http.ResponseWriter)
|
||||||
if ok {
|
if ok {
|
||||||
for k, v := range md {
|
for k, vs := range md {
|
||||||
for _, val := range v {
|
if h, ok := defaultOutgoingTrailerMatcher(k); ok {
|
||||||
w.Header().Add(k, val)
|
for _, v := range vs {
|
||||||
|
w.Header().Add(h, v)
|
||||||
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -40,8 +40,8 @@ func TestEventHandler_OnReceiveTrailers(t *testing.T) {
|
|||||||
},
|
},
|
||||||
expectedStatus: codes.OK,
|
expectedStatus: codes.OK,
|
||||||
expectedHeader: map[string][]string{
|
expectedHeader: map[string][]string{
|
||||||
"X-Custom-Header": {"value1", "value2"},
|
"Grpc-Trailer-X-Custom-Header": {"value1", "value2"},
|
||||||
"X-Another-Header": {"single-value"},
|
"Grpc-Trailer-X-Another-Header": {"single-value"},
|
||||||
},
|
},
|
||||||
},
|
},
|
||||||
{
|
{
|
||||||
@@ -100,9 +100,9 @@ func TestEventHandler_OnReceiveHeaders(t *testing.T) {
|
|||||||
"x-another-header": []string{"single-value"},
|
"x-another-header": []string{"single-value"},
|
||||||
},
|
},
|
||||||
expectedHeader: map[string][]string{
|
expectedHeader: map[string][]string{
|
||||||
"Content-Type": {"application/json"},
|
"Grpc-Metadata-Content-Type": {"application/json"},
|
||||||
"X-Custom-Header": {"value1", "value2"},
|
"Grpc-Metadata-X-Custom-Header": {"value1", "value2"},
|
||||||
"X-Another-Header": {"single-value"},
|
"Grpc-Metadata-X-Another-Header": {"single-value"},
|
||||||
},
|
},
|
||||||
},
|
},
|
||||||
{
|
{
|
||||||
@@ -158,7 +158,81 @@ func TestEventHandler_OnReceiveHeaders_MultipleValues(t *testing.T) {
|
|||||||
"x-header-2": []string{"value3"},
|
"x-header-2": []string{"value3"},
|
||||||
})
|
})
|
||||||
|
|
||||||
// Check that headers are accumulated (not overwritten)
|
// Check that headers are accumulated (not overwritten) with proper prefix
|
||||||
assert.Equal(t, []string{"value1", "value2"}, recorder.Header()["X-Header-1"])
|
assert.Equal(t, []string{"value1", "value2"}, recorder.Header()["Grpc-Metadata-X-Header-1"])
|
||||||
assert.Equal(t, []string{"value3"}, recorder.Header()["X-Header-2"])
|
assert.Equal(t, []string{"value3"}, recorder.Header()["Grpc-Metadata-X-Header-2"])
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestEventHandler_OnReceiveHeaders_MetadataPrefix(t *testing.T) {
|
||||||
|
tests := []struct {
|
||||||
|
name string
|
||||||
|
metadata metadata.MD
|
||||||
|
expectedHeader map[string][]string
|
||||||
|
}{
|
||||||
|
{
|
||||||
|
name: "all metadata headers should be prefixed with Grpc-Metadata-",
|
||||||
|
metadata: metadata.MD{
|
||||||
|
"content-type": []string{"application/grpc"},
|
||||||
|
"x-custom-header": []string{"value1"},
|
||||||
|
"authorization": []string{"Bearer token"},
|
||||||
|
},
|
||||||
|
expectedHeader: map[string][]string{
|
||||||
|
"Grpc-Metadata-Content-Type": {"application/grpc"},
|
||||||
|
"Grpc-Metadata-X-Custom-Header": {"value1"},
|
||||||
|
"Grpc-Metadata-Authorization": {"Bearer token"},
|
||||||
|
},
|
||||||
|
},
|
||||||
|
{
|
||||||
|
name: "mixed case headers should be prefixed",
|
||||||
|
metadata: metadata.MD{
|
||||||
|
"Content-Type": []string{"APPLICATION/JSON"},
|
||||||
|
"X-Custom-Header": []string{"value1"},
|
||||||
|
},
|
||||||
|
expectedHeader: map[string][]string{
|
||||||
|
"Grpc-Metadata-Content-Type": {"APPLICATION/JSON"},
|
||||||
|
"Grpc-Metadata-X-Custom-Header": {"value1"},
|
||||||
|
},
|
||||||
|
},
|
||||||
|
{
|
||||||
|
name: "multiple values for same header",
|
||||||
|
metadata: metadata.MD{
|
||||||
|
"x-multi-header": []string{"value1", "value2", "value3"},
|
||||||
|
},
|
||||||
|
expectedHeader: map[string][]string{
|
||||||
|
"Grpc-Metadata-X-Multi-Header": {"value1", "value2", "value3"},
|
||||||
|
},
|
||||||
|
},
|
||||||
|
{
|
||||||
|
name: "empty metadata",
|
||||||
|
metadata: metadata.MD{},
|
||||||
|
expectedHeader: map[string][]string{},
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
|
for _, tt := range tests {
|
||||||
|
t.Run(tt.name, func(t *testing.T) {
|
||||||
|
recorder := httptest.NewRecorder()
|
||||||
|
h := NewEventHandler(recorder, nil)
|
||||||
|
|
||||||
|
h.OnReceiveHeaders(tt.metadata)
|
||||||
|
|
||||||
|
// Check that headers are set correctly
|
||||||
|
for key, expectedValues := range tt.expectedHeader {
|
||||||
|
actualValues := recorder.Header()[key]
|
||||||
|
assert.Equal(t, expectedValues, actualValues, "Header %s should match", key)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Ensure no unexpected headers are set
|
||||||
|
for actualKey := range recorder.Header() {
|
||||||
|
found := false
|
||||||
|
for expectedKey := range tt.expectedHeader {
|
||||||
|
if actualKey == expectedKey {
|
||||||
|
found = true
|
||||||
|
break
|
||||||
|
}
|
||||||
|
}
|
||||||
|
assert.True(t, found, "Unexpected header found: %s", actualKey)
|
||||||
|
}
|
||||||
|
})
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
Reference in New Issue
Block a user