bouncer/internal/proxy/director/director.go

package director

import (
	"context"
	"net/http"
	"net/url"
	"sort"

	"forge.cadoles.com/Cadoles/go-proxy"
	"forge.cadoles.com/Cadoles/go-proxy/wildcard"
	"forge.cadoles.com/cadoles/bouncer/internal/store"
	"github.com/pkg/errors"
	"gitlab.com/wpetit/goweb/logger"
)

type Director struct {
	proxyRepository store.ProxyRepository
	layerRepository store.LayerRepository
	layerRegistry   *LayerRegistry
}

func (d *Director) rewriteRequest(r *http.Request) (*http.Request, error) {
	ctx := r.Context()

	proxies, err := d.getProxies(ctx)
	if err != nil {
		return r, errors.WithStack(err)
	}

	var match *store.Proxy

MAIN:
	for _, p := range proxies {
		for _, from := range p.From {
			if matches := wildcard.Match(r.Host, from); !matches {
				continue
			}

			match = p
			break MAIN
		}
	}

	if match == nil {
		return r, nil
	}

	toURL, err := url.Parse(match.To)
	if err != nil {
		return r, errors.WithStack(err)
	}

	r.URL.Host = toURL.Host
	r.URL.Scheme = toURL.Scheme

	ctx = logger.With(ctx,
		logger.F("proxy", match.Name),
		logger.F("host", r.Host),
		logger.F("remoteAddr", r.RemoteAddr),
	)

	ctx = withProxy(ctx, match)

	layers, err := d.getLayers(ctx, match.Name)
	if err != nil {
		return r, errors.WithStack(err)
	}

	ctx = withLayers(ctx, layers)
	r = r.WithContext(ctx)

	return r, nil
}

func (d *Director) getProxies(ctx context.Context) ([]*store.Proxy, error) {
	headers, err := d.proxyRepository.QueryProxy(ctx, store.WithProxyQueryEnabled(true))
	if err != nil {
		return nil, errors.WithStack(err)
	}

	sort.Sort(store.ByProxyWeight(headers))

	proxies := make([]*store.Proxy, 0, len(headers))

	for _, h := range headers {
		if !h.Enabled {
			continue
		}

		proxy, err := d.proxyRepository.GetProxy(ctx, h.Name)
		if err != nil {
			return nil, errors.WithStack(err)
		}

		proxies = append(proxies, proxy)
	}

	return proxies, nil
}

func (d *Director) getLayers(ctx context.Context, proxyName store.ProxyName) ([]*store.Layer, error) {
	headers, err := d.layerRepository.QueryLayers(ctx, proxyName, store.WithLayerQueryEnabled(true))
	if err != nil {
		return nil, errors.WithStack(err)
	}

	sort.Sort(store.ByLayerWeight(headers))

	layers := make([]*store.Layer, 0, len(headers))

	for _, h := range headers {
		if !h.Enabled {
			continue
		}

		layer, err := d.layerRepository.GetLayer(ctx, proxyName, h.Name)
		if err != nil {
			return nil, errors.WithStack(err)
		}

		layers = append(layers, layer)
	}

	return layers, nil
}

func (d *Director) RequestTransformer() proxy.RequestTransformer {
	return func(r *http.Request) {
		// ctx := r.Context()
		// proxy, layers, err := proxyLayersFromContext(ctx)
		// if err != nil {
		// 	if errors.Is(err, errContextKeyNotFound) {
		// 		return
		// 	}

		// 	logger.Error(ctx, "could not retrieve proxy and layers from context", logger.E(errors.WithStack(err)))

		// 	return
		// }
	}
}

func (d *Director) ResponseTransformer() proxy.ResponseTransformer {
	return func(r *http.Response) error {
		// ctx := r.Request.Context()
		// proxy, layers, err := proxyLayersFromContext(ctx)
		// if err != nil {
		// 	if errors.Is(err, errContextKeyNotFound) {
		// 		return nil
		// 	}

		// 	return errors.WithStack(err)
		// }

		return nil
	}
}

func (d *Director) Middleware() proxy.Middleware {
	return func(next http.Handler) http.Handler {
		fn := func(w http.ResponseWriter, r *http.Request) {
			r, err := d.rewriteRequest(r)
			if err != nil {
				logger.Error(r.Context(), "could not rewrite request", logger.E(errors.WithStack(err)))
				http.Error(w, http.StatusText(http.StatusInternalServerError), http.StatusInternalServerError)

				return
			}

			ctx := r.Context()

			_, layers, err := proxyLayersFromContext(ctx)
			if err != nil {
				if errors.Is(err, errContextKeyNotFound) {
					return
				}

				logger.Error(ctx, "could not retrieve proxy and layers from context", logger.E(errors.WithStack(err)))
				http.Error(w, http.StatusText(http.StatusInternalServerError), http.StatusInternalServerError)

				return
			}

			httpMiddlewares := make([]proxy.Middleware, 0)
			for _, layer := range layers {
				middleware, ok := d.layerRegistry.GetMiddleware(layer.Type, layer.Name)
				if !ok {
					continue
				}

				httpMiddlewares = append(httpMiddlewares, middleware.Middleware(layer))
			}

			handler := createMiddlewareChain(next, httpMiddlewares)

			handler.ServeHTTP(w, r)
		}

		return http.HandlerFunc(fn)
	}
}

func New(proxyRepository store.ProxyRepository, layerRepository store.LayerRepository, layers ...Layer) *Director {
	registry := NewLayerRegistry(layers...)

	return &Director{proxyRepository, layerRepository, registry}
}

func proxyLayersFromContext(ctx context.Context) (*store.Proxy, []*store.Layer, error) {
	proxy, err := ctxProxy(ctx)
	if err != nil {
		if errors.Is(err, errContextKeyNotFound) {
			return nil, nil, nil
		}

		logger.Error(ctx, "could not retrieve proxy from context", logger.E(errors.WithStack(err)))

		return nil, nil, nil
	}

	layers, err := ctxLayers(ctx)
	if err != nil {
		if errors.Is(err, errContextKeyNotFound) {
			return nil, nil, nil
		}

		logger.Error(ctx, "could not retrieve layers from context", logger.E(errors.WithStack(err)))

		return nil, nil, nil
	}

	return proxy, layers, nil
}