frostfs-s3-gw/api/host_bucket_router.go

package api

import (
	"net/http"
	"strings"

	"github.com/go-chi/chi/v5"
)

type HostBucketRouter struct {
	routes        map[string]chi.Router
	bktParam      string
	defaultRouter chi.Router
}

func NewHostBucketRouter(bktParam string) HostBucketRouter {
	return HostBucketRouter{
		routes:   make(map[string]chi.Router),
		bktParam: bktParam,
	}
}

func (hr *HostBucketRouter) Default(router chi.Router) {
	hr.defaultRouter = router
}

func (hr HostBucketRouter) Map(host string, h chi.Router) {
	hr.routes[strings.ToLower(host)] = h
}

func (hr HostBucketRouter) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	bucket, domain := getBucketDomain(getHost(r))
	router, ok := hr.routes[strings.ToLower(domain)]
	if !ok {
		router = hr.defaultRouter
		if router == nil {
			http.Error(w, http.StatusText(404), 404)
			return
		}
	}

	if rctx := chi.RouteContext(r.Context()); rctx != nil && bucket != "" {
		rctx.URLParams.Add(hr.bktParam, bucket)
	}

	router.ServeHTTP(w, r)
}

func getBucketDomain(host string) (bucket string, domain string) {
	parts := strings.Split(host, ".")
	if len(parts) > 1 {
		return parts[0], strings.Join(parts[1:], ".")
	}
	return "", host
}

// getHost tries its best to return the request host.
// According to section 14.23 of RFC 2616 the Host header
// can include the port number if the default value of 80 is not used.
func getHost(r *http.Request) string {
	host := r.Host
	if r.URL.IsAbs() {
		host = r.URL.Host
	}

	if i := strings.Index(host, ":"); i != -1 {
		host = host[:i]
	}

	return host
}