job-search/internal/middleware/request_checks.go

package middleware

import (
	"mime"
	"net"
	"net/http"
	"strings"

	"git.avsdev.uk/Go/job-search/internal/config"
	"git.avsdev.uk/Go/job-search/internal/httpx"
)

// defaultMaxBodyBytes is the largest request body when max_body_bytes is 0:
// 1MB, as the configuration's default.
const defaultMaxBodyBytes = 1_000_000

// RequestChecks validates the Host header against headers.allowed_hosts (if
// any are set) and caps the request body at max_body_bytes: 0 means
// defaultMaxBodyBytes, and -1 (util.ByteSizeNoLimit) no limit.
func RequestChecks(
	conf config.JSCServer,
) func(http.Handler) http.Handler {
	allowedHosts := newHostAllowList(conf.Headers.AllowedHosts)

	maxBodyBytes := int64(conf.MaxBodyBytes)
	if maxBodyBytes == 0 {
		maxBodyBytes = defaultMaxBodyBytes
	}

	return func(next http.Handler) http.Handler {
		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			if !allowedHosts.allows(r.Host) {
				httpx.WriteProblem(
					w, r, http.StatusMisdirectedRequest, "host not allowed",
				)
				return
			}

			if maxBodyBytes > 0 && r.Body != nil && r.Body != http.NoBody {
				if r.ContentLength > maxBodyBytes {
					httpx.WriteProblem(
						w, r, http.StatusRequestEntityTooLarge, "request body too large",
					)
					return
				}
				r.Body = http.MaxBytesReader(w, r.Body, maxBodyBytes)
			}

			next.ServeHTTP(w, r)
		})
	}
}

// hostAllowList is the set of hosts a request's Host header may name: exact
// host names, and suffixes from "*.example.com" patterns. An empty list
// allows every host.
type hostAllowList struct {
	exact map[string]bool
	// suffixes are the domains of "*.domain" patterns, e.g. ".example.com".
	suffixes []string
}

// newHostAllowList builds a hostAllowList from host names and "*.domain"
// patterns, ignoring case and blank entries.
func newHostAllowList(hosts []string) hostAllowList {
	list := hostAllowList{
		exact: map[string]bool{},
	}
	for _, host := range hosts {
		host = strings.ToLower(strings.TrimSpace(host))
		if strings.HasPrefix(host, "*.") {
			list.suffixes = append(list.suffixes, host[1:])
		} else if host != "" {
			list.exact[host] = true
		}
	}
	return list
}

// allows reports whether the Host header host (with or without a port, and
// a trailing dot) is allowed. A suffix matches subdomains only, not the bare
// domain.
func (l hostAllowList) allows(host string) bool {
	if len(l.exact) == 0 && len(l.suffixes) == 0 {
		return true
	}

	if hostname, _, err := net.SplitHostPort(host); err == nil {
		host = hostname
	}

	host = strings.ToLower(strings.TrimSuffix(host, "."))
	if l.exact[host] {
		return true
	}

	for _, suffix := range l.suffixes {
		if strings.HasSuffix(host, suffix) && len(host) > len(suffix) {
			return true
		}
	}

	return false
}

// RequireJSON rejects requests with a body that is not JSON (415 Unsupported
// Media Type).
func RequireJSON(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		hasBody := r.ContentLength > 0 || len(r.TransferEncoding) > 0

		if hasBody {
			mediaType, _, err := mime.ParseMediaType(r.Header.Get("Content-Type"))
			isJSON := mediaType == "application/json" ||
				strings.HasSuffix(mediaType, "+json")
			if err != nil || !isJSON {
				httpx.WriteProblem(
					w,
					r,
					http.StatusUnsupportedMediaType,
					"Content-Type must be application/json",
				)
				return
			}
		}

		next.ServeHTTP(w, r)
	})
}

// NoStore marks responses as non-cacheable.
func NoStore(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Cache-Control", "no-store")
		next.ServeHTTP(w, r)
	})
}