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