bouncer/internal/proxy/director/context.go

package director

import (
	"context"
	"net/url"

	"forge.cadoles.com/cadoles/bouncer/internal/store"
	"github.com/pkg/errors"
)

type contextKey string

const (
	contextKeyProxy       contextKey = "proxy"
	contextKeyLayers      contextKey = "layers"
	contextKeyOriginalURL contextKey = "originalURL"
)

var (
	errContextKeyNotFound     = errors.New("context key not found")
	errUnexpectedContextValue = errors.New("unexpected context value")
)

func withOriginalURL(ctx context.Context, url *url.URL) context.Context {
	return context.WithValue(ctx, contextKeyOriginalURL, url)
}

func OriginalURL(ctx context.Context) (*url.URL, error) {
	url, err := ctxValue[*url.URL](ctx, contextKeyOriginalURL)
	if err != nil {
		return nil, errors.WithStack(err)
	}

	return url, nil
}

func withProxy(ctx context.Context, proxy *store.Proxy) context.Context {
	return context.WithValue(ctx, contextKeyProxy, proxy)
}

func ctxProxy(ctx context.Context) (*store.Proxy, error) {
	proxy, err := ctxValue[*store.Proxy](ctx, contextKeyProxy)
	if err != nil {
		return nil, errors.WithStack(err)
	}

	return proxy, nil
}

func withLayers(ctx context.Context, layers []*store.Layer) context.Context {
	return context.WithValue(ctx, contextKeyLayers, layers)
}

func ctxLayers(ctx context.Context) ([]*store.Layer, error) {
	layers, err := ctxValue[[]*store.Layer](ctx, contextKeyLayers)
	if err != nil {
		return nil, errors.WithStack(err)
	}

	return layers, nil
}

func ctxValue[T any](ctx context.Context, key contextKey) (T, error) {
	raw := ctx.Value(key)
	if raw == nil {
		return *new(T), errors.WithStack(errContextKeyNotFound)
	}

	value, ok := raw.(T)
	if !ok {
		return *new(T), errors.WithStack(errUnexpectedContextValue)
	}

	return value, nil
}