Add support for Etag & If-None-Match headers about grpc-gateway HOT 6 OPEN

import (
	"context"
	"crypto/md5"
	"encoding/hex"
	"net/http"

	"google.golang.org/protobuf/proto"
)

type etagWriter struct {
	http.ResponseWriter
	wroteHeader bool
	ifNoneMatch string
}

func (w *etagWriter) Write(b []byte) (int, error) {
	etag := w.Header().Get("Etag")
	if !w.wroteHeader && w.ifNoneMatch == etag {
		w.ResponseWriter.WriteHeader(http.StatusNotModified)
		return 0, nil
	} else {
		return w.ResponseWriter.Write(b)
	}
}

func (w *etagWriter) WriteHeader(code int) {
	w.wroteHeader = true
	w.ResponseWriter.WriteHeader(code)
}

// Unwrap returns the original http.ResponseWriter. This is necessary
// to expose Flush() and Push() on the underlying response writer.
func (w *etagWriter) Unwrap() http.ResponseWriter {
	return w.ResponseWriter
}

// IfNoneMatchHandler wraps an http.Handler and will return a NotModified
// response if the If-None-Match header matches the Etag header.
func IfNoneMatchHandler(h http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		ifNoneMatch := r.Header.Get("If-None-Match")
		if ifNoneMatch != "" {
			w = &etagWriter{
				ResponseWriter: w,
				ifNoneMatch:    ifNoneMatch,
			}
		}

		h.ServeHTTP(w, r)
	})
}

func ForwardResponseWithEtag(_ context.Context, w http.ResponseWriter, m proto.Message) error {
	// NOTE: Unfortunately we have to serialize the protobuf
	data, err := proto.Marshal(m)
	if err != nil {
		return err
	}

	// NOTE: We don't have access to the request, so this can't be limited to just GET methods
	if len(data) > 100 {
		h := md5.New()
		h.Write(data)
		etag := hex.EncodeToString(h.Sum(nil))
		w.Header().Set("Etag", "\""+etag+"\"")
	}

	return nil
}

mux := runtime.NewServeMux(runtime.WithForwardResponseOption(ForwardResponseWithEtag))

// Register generated gateway handlers

s := &http.Server{
    Handler: IfNoneMatchHandler(mux),
}

import (
	"crypto/md5"
	"encoding/hex"
	"fmt"
	"net/http"
)

type etagWriter struct {
	http.ResponseWriter
	wroteHeader bool
	ifNoneMatch string
	writeEtag   bool
	minBytes    int
}

func (w etagWriter) Write(b []byte) (int, error) {
	if w.wroteHeader || !w.writeEtag || len(b) < w.minBytes {
		return w.ResponseWriter.Write(b)
	}

	// Generate the Etag
	h := md5.New()
	h.Write(b)
	etag := fmt.Sprintf("\"%s\"", hex.EncodeToString(h.Sum(nil)))

	w.Header().Set("Etag", etag)

	if w.ifNoneMatch != "" && w.ifNoneMatch == etag {
		w.ResponseWriter.WriteHeader(http.StatusNotModified)
		return 0, nil
	} else {
		return w.ResponseWriter.Write(b)
	}
}

func (w etagWriter) WriteHeader(code int) {
	// Track if the headers have already been written
	w.wroteHeader = true
	w.ResponseWriter.WriteHeader(code)
}

// Unwrap returns the original http.ResponseWriter. This is necessary
// to expose Flush() and Push() on the underlying response writer.
func (w etagWriter) Unwrap() http.ResponseWriter {
	return w.ResponseWriter
}

// EtagHandler wraps an http.Handler and will write an Etag header to the
// response if the request method is GET and the response size is greater
// than or equal to minBytes.  It will also return a NotModified response
// if the If-None-Match header matches the Etag header.
func EtagHandler(h http.Handler, minBytes int) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w = etagWriter{
			ResponseWriter: w,
			ifNoneMatch:    r.Header.Get("If-None-Match"),
			writeEtag:      r.Method == http.MethodGet,
			minBytes:       minBytes,
		}

		h.ServeHTTP(w, r)
	})
}

mux := runtime.NewServeMux()

// Register generated gateway handlers

s := &http.Server{
    Handler: EtagHandler(mux, 100),
}