chi-s3-vhs/main.go

package main

import (
	"context"
	"flag"
	"fmt"
	"log"
	"net"
	"net/http"
	"os"
	"os/signal"
	"strings"

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

type AddressTypeRouter struct {
	vhsGlobalEnabled    bool            // setting from gateway application
	vhsNamespaceEnabled map[string]bool // setting from gateway application

	pathStyleRouter chi.Router // look for 'defaultRouter' in S3 Gateway
	vhsRouter       chi.Router // look for 'bucketRouter' in S3 Gateway
}

type nsSettings struct {
	settings map[string]bool
}

func (s *nsSettings) Set(value string) error {
	v := strings.Split(value, "=")
	if len(v) != 2 {
		return fmt.Errorf("invalid argument: %s", value)
	}
	switch v[1] {
	case "enabled":
		s.settings[v[0]] = true
	case "disabled":
		s.settings[v[0]] = false
	}
	return nil
}
func (s *nsSettings) String() string {
	return fmt.Sprintf("%v", s.settings)
}

func main() {
	addr := flag.String("a", ":44412", "http listen address")
	globalVHSEnabled := flag.Bool("global", false, "enable global VHS")
	namespaceVHSEnabled := nsSettings{
		settings: make(map[string]bool),
	}
	flag.Var(&namespaceVHSEnabled, "ns", "namespace settings")
	flag.Parse()

	ctx, cancel := signal.NotifyContext(context.Background(), os.Interrupt)
	defer cancel()

	objRouter := ObjectRouter()          // router for object operations
	bktRouter := BucketRouter(objRouter) // router for bucket operations
	pathStyleRouter := chi.NewRouter()   // router for path-style requests
	pathStyleRouter.Mount("/{bucket}", bktRouter)
	atr := AddressTypeRouter{ // router for path-style and virtual-host-style switch
		vhsGlobalEnabled:    *globalVHSEnabled,
		vhsNamespaceEnabled: namespaceVHSEnabled.settings,
		pathStyleRouter:     pathStyleRouter,
		vhsRouter:           bktRouter,
	}

	api := chi.NewRouter() // global router
	api.Use(Request())
	api.Mount("/", atr)

	srv := &http.Server{}
	srv.Handler = api
	l, err := net.Listen("tcp", *addr)
	if err != nil {
		log.Fatal(fmt.Sprintf("failed to listen on %s: %v", *addr, err))
	}
	defer l.Close()

	go srv.Serve(l)
	log.Printf("listening on %s", *addr)
	log.Printf("global VHS enabled: %v", *globalVHSEnabled)
	log.Printf("namespace override: %s", &namespaceVHSEnabled)
	<-ctx.Done()
}

func (atr AddressTypeRouter) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	namespace := r.Header.Get("X-FrostFS-Namespace")
	router := atr.pathStyleRouter

	vhsEnabled := atr.isVHSAddress(r, namespace)
	if vhsEnabled {
		log.Println("VHS is enabled")
		router = atr.vhsRouter
		bucketName := atr.vhsBucketName(r)
		if rctx := chi.RouteContext(r.Context()); rctx != nil {
			rctx.URLParams.Add("bucket", bucketName)
		}
	}

	router.ServeHTTP(w, r)
}

func (atr AddressTypeRouter) isVHSAddress(r *http.Request, namespace string) bool {
	result := atr.vhsGlobalEnabled // default global value

	// check namespace override
	if v, ok := atr.vhsNamespaceEnabled[namespace]; ok {
		result = v
	}

	// check header override
	switch r.Header.Get("X-FrostFS-VHS") {
	case "enabled":
		result = true
	case "disabled":
		result = false
	default:
	}

	return result
}

func (atr AddressTypeRouter) vhsBucketName(r *http.Request) string {
	hostname := r.Host
	if r.URL.IsAbs() {
		hostname = r.URL.Hostname()
	}

	parts := strings.Split(hostname, ".")

	return parts[0]
}