From d1499d3fab638d82a4e1a6aec68931fd2531c1cb Mon Sep 17 00:00:00 2001 From: Joshua Garnett Date: Sat, 27 Apr 2024 07:33:33 -0400 Subject: [PATCH] Added support for writing Etag headers and handling If-None-Match request headers * Taken from https://developer.mozilla.org/en-US/docs/Web/HTTP/Headers/ETag: "The ETag (or entity tag) HTTP response header is an identifier for a specific version of a resource. It lets caches be more efficient and save bandwidth, as a web server does not need to resend a full response if the content was not changed." * This feature is enabled by passing in the WithEnableEtagSupport option when creating the mux * The Etag is only written for GET requests with messages greater than 100 bytes and is set to the md5 hash of the message. * If a client sends the If-None-Match request header, the server will compare that value to the calculated etag. If they match, it will return http.StatusNotModified (304) and not write the message in the response. * Added a test to ensure the correct headers are being returned when the option is set and that If-None-Match works. --- runtime/handler.go | 26 ++++++++++++- runtime/handler_test.go | 82 +++++++++++++++++++++++++++++++++++++++++ runtime/mux.go | 8 ++++ 3 files changed, 114 insertions(+), 2 deletions(-) diff --git a/runtime/handler.go b/runtime/handler.go index 63dc812b781..e9067162293 100644 --- a/runtime/handler.go +++ b/runtime/handler.go @@ -2,6 +2,8 @@ package runtime import ( "context" + "crypto/md5" + "encoding/hex" "errors" "io" "net/http" @@ -181,8 +183,28 @@ func ForwardResponseMessage(ctx context.Context, mux *ServeMux, marshaler Marsha w.Header().Set("Content-Length", strconv.Itoa(len(buf))) } - if _, err = w.Write(buf); err != nil { - grpclog.Infof("Failed to write response: %v", err) + // If enabled, generate an Etag for any GET requests with messages larger than 100 bytes. + // Writing the Etag in the response takes 41 bytes, so it's not worth doing for smaller messages. + if mux.enableEtagSupport && req.Method == http.MethodGet && len(buf) > 100 { + h := md5.New() + h.Write(buf) + etag := hex.EncodeToString(h.Sum(nil)) + w.Header().Set("Etag", "\""+etag+"\"") + + // Check if the client has provided If-None-Match and if it matches the generated Etag. + // If it does, send a 304 Not Modified response. + ifNoneMatch := req.Header.Get("If-None-Match") + if ifNoneMatch != "" && ifNoneMatch == etag { + w.WriteHeader(http.StatusNotModified) + } else { + if _, err = w.Write(buf); err != nil { + grpclog.Infof("Failed to write response: %v", err) + } + } + } else { + if _, err = w.Write(buf); err != nil { + grpclog.Infof("Failed to write response: %v", err) + } } if doForwardTrailers { diff --git a/runtime/handler_test.go b/runtime/handler_test.go index b0ad6e2c90d..a464d31d0b7 100644 --- a/runtime/handler_test.go +++ b/runtime/handler_test.go @@ -493,3 +493,85 @@ func TestOutgoingTrailerMatcher(t *testing.T) { }) } } + +func TestOutgoingEtagIfNoneMatch(t *testing.T) { + t.Parallel() + for _, tc := range []struct { + msg *pb.SimpleMessage + requestHeaders http.Header + method string + name string + headers http.Header + expectedStatus int + }{ + { + msg: &pb.SimpleMessage{Id: "foo"}, + method: http.MethodGet, + name: "small message", + headers: http.Header{ + "Content-Length": []string{"12"}, + "Content-Type": []string{"application/json"}, + }, + expectedStatus: http.StatusOK, + }, + { + msg: &pb.SimpleMessage{Id: "Lorem ipsum dolor sit amet, consectetur adipiscing elit. Aliquam rhoncus magna ante, sed malesuada nibh vehicula in nec."}, + method: http.MethodGet, + name: "large message", + headers: http.Header{ + "Content-Length": []string{"129"}, + "Content-Type": []string{"application/json"}, + "Etag": []string{"\"41bf5d28a47f59b2a649e44f2607b0ea\""}, + }, + expectedStatus: http.StatusOK, + }, + { + msg: &pb.SimpleMessage{Id: "Lorem ipsum dolor sit amet, consectetur adipiscing elit. Aliquam rhoncus magna ante, sed malesuada nibh vehicula in nec."}, + method: http.MethodGet, + requestHeaders: http.Header{ + "If-None-Match": []string{"41bf5d28a47f59b2a649e44f2607b0ea"}, + }, + name: "large message with If-None-Match header", + headers: http.Header{ + "Content-Length": []string{"129"}, + "Content-Type": []string{"application/json"}, + "Etag": []string{"\"41bf5d28a47f59b2a649e44f2607b0ea\""}, + }, + expectedStatus: http.StatusNotModified, + }, + { + msg: &pb.SimpleMessage{Id: "Lorem ipsum dolor sit amet, consectetur adipiscing elit. Aliquam rhoncus magna ante, sed malesuada nibh vehicula in nec."}, + method: http.MethodPost, + requestHeaders: http.Header{ + "If-None-Match": []string{"41bf5d28a47f59b2a649e44f2607b0ea"}, + }, + name: "large message with If-None-Match header", + headers: http.Header{ + "Content-Length": []string{"129"}, + "Content-Type": []string{"application/json"}, + }, + expectedStatus: http.StatusOK, + }, + } { + tc := tc + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + + req := httptest.NewRequest(tc.method, "http://example.com/foo", nil) + req.Header = tc.requestHeaders + resp := httptest.NewRecorder() + + runtime.ForwardResponseMessage(context.Background(), runtime.NewServeMux(runtime.WithEnableEtagSupport()), &runtime.JSONPb{}, resp, req, tc.msg) + + w := resp.Result() + defer w.Body.Close() + if w.StatusCode != tc.expectedStatus { + t.Fatalf("StatusCode %d want %d", w.StatusCode, http.StatusOK) + } + + if !reflect.DeepEqual(w.Header, tc.headers) { + t.Fatalf("Header %v want %v", w.Header, tc.headers) + } + }) + } +} diff --git a/runtime/mux.go b/runtime/mux.go index ed9a7e4387d..5655be22bae 100644 --- a/runtime/mux.go +++ b/runtime/mux.go @@ -63,6 +63,7 @@ type ServeMux struct { streamErrorHandler StreamErrorHandlerFunc routingErrorHandler RoutingErrorHandlerFunc disablePathLengthFallback bool + enableEtagSupport bool unescapingMode UnescapingMode } @@ -224,6 +225,13 @@ func WithDisablePathLengthFallback() ServeMuxOption { } } +// WithEnableEtagSupport returns a ServeMuxOption that enables writing Etags and handling If-None-Match request headers +func WithEnableEtagSupport() ServeMuxOption { + return func(serveMux *ServeMux) { + serveMux.enableEtagSupport = true + } +} + // WithHealthEndpointAt returns a ServeMuxOption that will add an endpoint to the created ServeMux at the path specified by endpointPath. // When called the handler will forward the request to the upstream grpc service health check (defined in the // gRPC Health Checking Protocol).