1package middleware
2
3import (
4	"net/http"
5	"strings"
6)
7
8// AllowContentEncoding enforces a whitelist of request Content-Encoding otherwise responds
9// with a 415 Unsupported Media Type status.
10func AllowContentEncoding(contentEncoding ...string) func(next http.Handler) http.Handler {
11	allowedEncodings := make(map[string]struct{}, len(contentEncoding))
12	for _, encoding := range contentEncoding {
13		allowedEncodings[strings.TrimSpace(strings.ToLower(encoding))] = struct{}{}
14	}
15	return func(next http.Handler) http.Handler {
16		fn := func(w http.ResponseWriter, r *http.Request) {
17			requestEncodings := r.Header["Content-Encoding"]
18			// skip check for empty content body or no Content-Encoding
19			if r.ContentLength == 0 {
20				next.ServeHTTP(w, r)
21				return
22			}
23			// All encodings in the request must be allowed
24			for _, encoding := range requestEncodings {
25				if _, ok := allowedEncodings[strings.TrimSpace(strings.ToLower(encoding))]; !ok {
26					w.WriteHeader(http.StatusUnsupportedMediaType)
27					return
28				}
29			}
30			next.ServeHTTP(w, r)
31		}
32		return http.HandlerFunc(fn)
33	}
34}
35