Skip to content

Commit

Permalink
Update content_type.go
Browse files Browse the repository at this point in the history
  • Loading branch information
GocaMaric authored Dec 5, 2023
1 parent 58ca6d6 commit 920d5e5
Showing 1 changed file with 9 additions and 13 deletions.
22 changes: 9 additions & 13 deletions middleware/content_type.go
Original file line number Diff line number Diff line change
Expand Up @@ -6,44 +6,40 @@ import (
)

// SetHeader is a convenience handler to set a response header key/value
func SetHeader(key, value string) func(next http.Handler) http.Handler {
func SetHeader(key, value string) func(http.Handler) http.Handler {
return func(next http.Handler) http.Handler {
fn := func(w http.ResponseWriter, r *http.Request) {
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.Header().Set(key, value)
next.ServeHTTP(w, r)
}
return http.HandlerFunc(fn)
})
}
}

// AllowContentType enforces a whitelist of request Content-Types otherwise responds
// with a 415 Unsupported Media Type status.
func AllowContentType(contentTypes ...string) func(next http.Handler) http.Handler {
func AllowContentType(contentTypes ...string) func(http.Handler) http.Handler {
allowedContentTypes := make(map[string]struct{}, len(contentTypes))
for _, ctype := range contentTypes {
allowedContentTypes[strings.TrimSpace(strings.ToLower(ctype))] = struct{}{}
}

return func(next http.Handler) http.Handler {
fn := func(w http.ResponseWriter, r *http.Request) {
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if r.ContentLength == 0 {
// skip check for empty content body
// Skip check for empty content body
next.ServeHTTP(w, r)
return
}

s := strings.ToLower(strings.TrimSpace(r.Header.Get("Content-Type")))
if i := strings.Index(s, ";"); i > -1 {
s = s[0:i]
}
s := strings.ToLower(strings.TrimSpace(strings.Split(r.Header.Get("Content-Type"), ";")[0]))

if _, ok := allowedContentTypes[s]; ok {
next.ServeHTTP(w, r)
return
}

w.WriteHeader(http.StatusUnsupportedMediaType)
}
return http.HandlerFunc(fn)
})
}
}

0 comments on commit 920d5e5

Please sign in to comment.