go-omni/cors/cors.go

package cors

import (
	"net/http"
	"strconv"
	"strings"
	"time"
)

// Policy is immutable after New and safe to share between concurrent requests.
type Policy struct {
	enabled, wildcard, credentials                    bool
	origins, methods, headers                         map[string]struct{}
	allowMethods, allowHeaders, exposeHeaders, maxAge string
}

// New validates and snapshots policy. It performs no network or DNS operations.
func New(c Config) (*Policy, error) {
	if err := c.Validate(); err != nil {
		return nil, err
	}
	p := &Policy{enabled: c.Enabled, credentials: c.AllowCredentials}
	if !c.Enabled {
		return p, nil
	}
	p.origins = make(map[string]struct{}, len(c.AllowedOrigins))
	for _, origin := range c.AllowedOrigins {
		if origin == "*" {
			p.wildcard = true
			continue
		}
		normalized, _ := normalizeOrigin(origin)
		p.origins[normalized] = struct{}{}
	}
	methods := c.AllowedMethods
	if len(methods) == 0 {
		methods = []string{http.MethodGet, http.MethodHead, http.MethodPost}
	}
	p.methods, p.allowMethods = compileList(methods, false)
	p.headers, p.allowHeaders = compileList(c.AllowedHeaders, true)
	_, p.exposeHeaders = compileList(c.ExposedHeaders, true)
	if c.MaxAge > 0 {
		p.maxAge = strconv.FormatInt(int64(c.MaxAge/time.Second), 10)
	}
	return p, nil
}

func compileList(values []string, header bool) (map[string]struct{}, string) {
	set := make(map[string]struct{}, len(values))
	list := make([]string, 0, len(values))
	for _, value := range values {
		key := value
		if header {
			key = strings.ToLower(value)
			value = http.CanonicalHeaderKey(value)
		}
		if _, exists := set[key]; exists {
			continue
		}
		set[key] = struct{}{}
		list = append(list, value)
	}
	return set, strings.Join(list, ", ")
}

// Wrap installs CORS outside routing/authentication. Only actual CORS preflights
// are short-circuited. Actual requests always continue through normal handlers;
// an unlisted Origin receives no permission to read the response in a browser.
// The original ResponseWriter, Request and body pass through without buffering.
func (p *Policy) Wrap(next http.Handler) http.Handler {
	if next == nil {
		panic("cors: nil handler")
	}
	if p == nil {
		panic("cors: nil policy")
	}
	if !p.enabled {
		return next
	}
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		h := w.Header()
		vary(h, "Origin")
		origins := r.Header.Values("Origin")
		// An OPTIONS request without Origin or Access-Control-Request-Method is
		// ordinary application traffic, including user-defined OPTIONS handlers.
		requestedMethods := r.Header.Values("Access-Control-Request-Method")
		preflight := r.Method == http.MethodOptions && len(origins) > 0 && len(requestedMethods) > 0
		if preflight {
			vary(h, "Access-Control-Request-Method")
			vary(h, "Access-Control-Request-Headers")
		}
		if len(origins) != 1 || !p.allowsOrigin(origins[0]) {
			if preflight {
				w.WriteHeader(http.StatusForbidden)
				return
			}
			next.ServeHTTP(w, r)
			return
		}
		if preflight {
			if len(requestedMethods) != 1 || !p.allowsMethod(requestedMethods[0]) || !p.allowsHeaders(r.Header.Values("Access-Control-Request-Headers")) {
				w.WriteHeader(http.StatusForbidden)
				return
			}
			p.allowOrigin(h, origins[0])
			h.Set("Access-Control-Allow-Methods", p.allowMethods)
			if p.allowHeaders != "" {
				h.Set("Access-Control-Allow-Headers", p.allowHeaders)
			}
			if p.maxAge != "" {
				h.Set("Access-Control-Max-Age", p.maxAge)
			}
			w.WriteHeader(http.StatusNoContent)
			return
		}
		// Match the actual method too: this grants response access, not permission
		// to execute a side effect. Authentication/CSRF checks remain in the app.
		if p.allowsMethod(r.Method) {
			p.allowOrigin(h, origins[0])
			if p.exposeHeaders != "" {
				h.Set("Access-Control-Expose-Headers", p.exposeHeaders)
			}
		}
		next.ServeHTTP(w, r)
	})
}

func (p *Policy) allowsOrigin(origin string) bool {
	if !p.wildcard {
		_, ok := p.origins[origin]
		return ok
	}
	normalized, ok := normalizeOrigin(origin)
	return ok && normalized == origin
}
func (p *Policy) allowsMethod(method string) bool { _, ok := p.methods[method]; return ok }
func (p *Policy) allowOrigin(h http.Header, origin string) {
	if p.wildcard {
		origin = "*"
	}
	h.Set("Access-Control-Allow-Origin", origin)
	if p.credentials {
		h.Set("Access-Control-Allow-Credentials", "true")
	}
}
func (p *Policy) allowsHeaders(lines []string) bool {
	count, size := 0, 0
	for _, line := range lines {
		size += len(line)
		if size > 8192 {
			return false
		}
		// Parse a header list without allocating an attacker-sized split slice.
		for {
			name, rest, more := strings.Cut(line, ",")
			name = strings.Trim(name, " \t")
			count++
			if count > 128 || !token(name) {
				return false
			}
			if _, ok := p.headers[strings.ToLower(name)]; !ok {
				return false
			}
			if !more {
				break
			}
			line = rest
		}
	}
	return true
}
func vary(h http.Header, name string) {
	for _, line := range h.Values("Vary") {
		for {
			v, rest, more := strings.Cut(line, ",")
			v = strings.TrimSpace(v)
			if v == "*" || strings.EqualFold(v, name) {
				return
			}
			if !more {
				break
			}
			line = rest
		}
	}
	h.Add("Vary", name)
}