chi-s3-vhs/s3.go

package main

import (
	"context"
	"fmt"
	"log"
	"net/http"
	"net/url"

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

type (
	Func func(h http.Handler) http.Handler

	ReqInfo struct {
		BucketName string
		ObjectName string
	}
)

const ctxRequestInfo = "FrostFS-S3-GW"

func GetReqInfo(ctx context.Context) *ReqInfo {
	if ctx == nil {
		return &ReqInfo{}
	} else if r, ok := ctx.Value(ctxRequestInfo).(*ReqInfo); ok {
		return r
	}
	return &ReqInfo{}
}

func SetReqInfo(ctx context.Context, req *ReqInfo) context.Context {
	if ctx == nil {
		return nil
	}
	return context.WithValue(ctx, ctxRequestInfo, req)
}

func AddBucketName() Func {
	return func(h http.Handler) http.Handler {
		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			ctx := r.Context()

			reqInfo := GetReqInfo(ctx)
			reqInfo.BucketName = chi.URLParam(r, "bucket")

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

// AddObjectName adds objects name to ReqInfo from context.
func AddObjectName() Func {
	return func(h http.Handler) http.Handler {
		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			ctx := r.Context()
			reqInfo := GetReqInfo(ctx)

			rctx := chi.RouteContext(ctx)
			// trim leading slash (always present)
			reqInfo.ObjectName = rctx.RoutePath[1:]

			if r.URL.RawPath != "" {
				// we have to do this because of
				// https://github.com/go-chi/chi/issues/641
				// https://github.com/go-chi/chi/issues/642
				if obj, err := url.PathUnescape(reqInfo.ObjectName); err != nil {
					// log error
				} else {
					reqInfo.ObjectName = obj
				}
			}

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

func Request() Func {
	return func(h http.Handler) http.Handler {
		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			reqInfo := &ReqInfo{}
			r = r.WithContext(SetReqInfo(r.Context(), reqInfo))
			h.ServeHTTP(w, r)
		})
	}
}

func BucketRouter(objRouter chi.Router) chi.Router {
	bktRouter := chi.NewRouter()
	bktRouter.Use(AddBucketName())

	bktRouter.Mount("/", objRouter)

	bktRouter.Group(func(r chi.Router) {
		r.MethodFunc(http.MethodGet, "/", BucketHandler())
	})

	return bktRouter
}

func ObjectRouter() chi.Router {
	objRouter := chi.NewRouter()
	objRouter.Use(AddObjectName())

	objRouter.Group(func(r chi.Router) {
		r.MethodFunc(http.MethodGet, "/*", ObjectHandler())
	})

	return objRouter
}

func ObjectHandler() http.HandlerFunc {
	return func(rw http.ResponseWriter, r *http.Request) {
		log.Println("object handler is triggered")
		ri := GetReqInfo(r.Context())
		if ri == nil {
			rw.WriteHeader(http.StatusInternalServerError)
			return
		}

		rw.WriteHeader(http.StatusOK)
		rw.Write([]byte(fmt.Sprintf("bucket:[%s] object:[%s]\n", ri.BucketName, ri.ObjectName)))
	}
}

func BucketHandler() http.HandlerFunc {
	return func(rw http.ResponseWriter, r *http.Request) {
		log.Printf("bucket handler is triggered")
		ri := GetReqInfo(r.Context())
		if ri == nil {
			log.Println("can't find request info")
			rw.WriteHeader(http.StatusInternalServerError)
			return
		}

		log.Printf("bucket:[%s]", ri.BucketName)
		rw.WriteHeader(http.StatusOK)
		rw.Write([]byte(fmt.Sprintf("bucket:[%s]\n", ri.BucketName)))
	}
}