1package middleware
2
3import (
4	"net/http"
5	"strings"
6)
7
8// SetHeader is a convenience handler to set a response header key/value
9func SetHeader(key, value string) func(next http.Handler) http.Handler {
10	return func(next http.Handler) http.Handler {
11		fn := func(w http.ResponseWriter, r *http.Request) {
12			w.Header().Set(key, value)
13			next.ServeHTTP(w, r)
14		}
15		return http.HandlerFunc(fn)
16	}
17}
18
19// AllowContentType enforces a whitelist of request Content-Types otherwise responds
20// with a 415 Unsupported Media Type status.
21func AllowContentType(contentTypes ...string) func(next http.Handler) http.Handler {
22	cT := []string{}
23	for _, t := range contentTypes {
24		cT = append(cT, strings.ToLower(t))
25	}
26
27	return func(next http.Handler) http.Handler {
28		fn := func(w http.ResponseWriter, r *http.Request) {
29			if r.ContentLength == 0 {
30				// skip check for empty content body
31				next.ServeHTTP(w, r)
32				return
33			}
34
35			s := strings.ToLower(strings.TrimSpace(r.Header.Get("Content-Type")))
36			if i := strings.Index(s, ";"); i > -1 {
37				s = s[0:i]
38			}
39
40			for _, t := range cT {
41				if t == s {
42					next.ServeHTTP(w, r)
43					return
44				}
45			}
46
47			w.WriteHeader(http.StatusUnsupportedMediaType)
48		}
49		return http.HandlerFunc(fn)
50	}
51}
52