bouncer/internal/proxy/director/layer/rewriter/layer.go

package rewriter

import (
	"net/http"

	proxy "forge.cadoles.com/Cadoles/go-proxy"
	"forge.cadoles.com/Cadoles/go-proxy/wildcard"
	"forge.cadoles.com/cadoles/bouncer/internal/proxy/director"
	"forge.cadoles.com/cadoles/bouncer/internal/proxy/director/layer/util"
	"forge.cadoles.com/cadoles/bouncer/internal/rule"
	ruleHTTP "forge.cadoles.com/cadoles/bouncer/internal/rule/http"
	"forge.cadoles.com/cadoles/bouncer/internal/store"
	"github.com/pkg/errors"
)

const LayerType store.LayerType = "rewriter"

type Layer struct {
	requestRuleEngineCache  *util.RuleEngineCache[*RequestVars, *LayerOptions]
	responseRuleEngineCache *util.RuleEngineCache[*ResponseVars, *LayerOptions]
}

func (l *Layer) LayerType() store.LayerType {
	return LayerType
}

func (l *Layer) Middleware(layer *store.Layer) proxy.Middleware {
	return func(next http.Handler) http.Handler {
		fn := func(w http.ResponseWriter, r *http.Request) {
			ctx := r.Context()

			options, err := fromStoreOptions(layer.Options)
			if err != nil {
				director.HandleError(ctx, w, r, http.StatusInternalServerError, errors.Wrap(err, "could not parse layer options"))
				return
			}

			matches := wildcard.MatchAny(r.URL.String(), options.MatchURLs...)
			if !matches {
				next.ServeHTTP(w, r)

				return
			}

			if err := l.applyRequestRules(ctx, r, layer, options); err != nil {
				var redirect *errRedirect
				if errors.As(err, &redirect) {
					http.Redirect(w, r, redirect.URL(), redirect.StatusCode())
					return
				}

				director.HandleError(ctx, w, r, http.StatusInternalServerError, errors.Wrap(err, "could not apply request rules"))

				return
			}

			next.ServeHTTP(w, r)
		}

		return http.HandlerFunc(fn)
	}
}

// ResponseTransformer implements director.ResponseTransformerLayer.
func (l *Layer) ResponseTransformer(layer *store.Layer) proxy.ResponseTransformer {
	return func(r *http.Response) error {
		options, err := fromStoreOptions(layer.Options)
		if err != nil {
			return errors.WithStack(err)
		}

		matches := wildcard.MatchAny(r.Request.URL.String(), options.MatchURLs...)
		if !matches {
			return nil
		}

		ctx := r.Request.Context()

		if err := l.applyResponseRules(ctx, r, layer, options); err != nil {
			return errors.WithStack(err)
		}

		return nil
	}
}

func New(funcs ...OptionFunc) *Layer {
	return &Layer{
		requestRuleEngineCache: util.NewInMemoryRuleEngineCache(func(options *LayerOptions) (*rule.Engine[*RequestVars], error) {
			engine, err := rule.NewEngine[*RequestVars](
				rule.WithRules(options.Rules.Request...),
				ruleHTTP.WithRequestFuncs(),
				WithRewriterFuncs(),
			)
			if err != nil {
				return nil, errors.WithStack(err)
			}

			return engine, nil
		}),
		responseRuleEngineCache: util.NewInMemoryRuleEngineCache(func(options *LayerOptions) (*rule.Engine[*ResponseVars], error) {
			engine, err := rule.NewEngine[*ResponseVars](
				rule.WithRules(options.Rules.Response...),
				ruleHTTP.WithResponseFuncs(),
			)
			if err != nil {
				return nil, errors.WithStack(err)
			}

			return engine, nil
		}),
	}
}

var (
	_ director.MiddlewareLayer          = &Layer{}
	_ director.ResponseTransformerLayer = &Layer{}
)