Added support for writing Etag headers and handling If-None-Match req…

…uest 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

@@ -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 {

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)
+			}
+		})
+	}
+}

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).