bouncer/internal/rule/http/response.go

package http

import (
	"context"
	"fmt"
	"net/http"
	"strconv"
	"strings"
	"time"

	"forge.cadoles.com/Cadoles/go-proxy/wildcard"
	"forge.cadoles.com/cadoles/bouncer/internal/rule"
	"github.com/expr-lang/expr"
	"github.com/pkg/errors"
)

func addResponseHeaderFunc() expr.Option {
	return expr.Function(
		"add_header",
		func(params ...any) (any, error) {
			ctx, err := rule.Assert[context.Context](params[0])
			if err != nil {
				return nil, errors.WithStack(err)
			}

			name, err := rule.Assert[string](params[1])
			if err != nil {
				return nil, errors.WithStack(err)
			}

			rawValue := params[2]

			var value string
			switch v := rawValue.(type) {
			case []string:
				value = strings.Join(v, ",")
			case time.Time:
				value = strconv.FormatInt(v.UTC().Unix(), 10)
			case time.Duration:
				value = strconv.FormatInt(int64(v.Seconds()), 10)
			default:
				value = fmt.Sprintf("%v", rawValue)
			}

			r, ok := CtxResponse(ctx)
			if !ok {
				return nil, errors.New("could not find http response in context")
			}

			r.Header.Add(name, value)

			return true, nil
		},
		new(func(context.Context, string, string) bool),
	)
}

func setResponseHeaderFunc() expr.Option {
	return expr.Function(
		"set_header",
		func(params ...any) (any, error) {
			ctx, err := rule.Assert[context.Context](params[0])
			if err != nil {
				return nil, errors.WithStack(err)
			}

			name, err := rule.Assert[string](params[1])
			if err != nil {
				return nil, errors.WithStack(err)
			}

			rawValue := params[2]

			var value string
			switch v := rawValue.(type) {
			case []string:
				value = strings.Join(v, ",")
			case time.Time:
				value = strconv.FormatInt(v.UTC().Unix(), 10)
			case time.Duration:
				value = strconv.FormatInt(int64(v.Seconds()), 10)
			default:
				value = fmt.Sprintf("%v", rawValue)
			}

			r, ok := CtxResponse(ctx)
			if !ok {
				return nil, errors.New("could not find http response in context")
			}

			r.Header.Set(name, value)

			return true, nil
		},
		new(func(context.Context, string, string) bool),
	)
}

func delResponseHeadersFunc() expr.Option {
	return expr.Function(
		"del_headers",
		func(params ...any) (any, error) {
			ctx, err := rule.Assert[context.Context](params[0])
			if err != nil {
				return nil, errors.WithStack(err)
			}

			pattern, err := rule.Assert[string](params[1])
			if err != nil {
				return nil, errors.WithStack(err)
			}

			r, ok := CtxResponse(ctx)
			if !ok {
				return nil, errors.New("could not find http response in context")
			}

			deleted := false

			for key := range r.Header {
				if !wildcard.Match(key, pattern) {
					continue
				}

				r.Header.Del(key)
				deleted = true
			}

			return deleted, nil
		},
		new(func(context.Context, string) bool),
	)
}

func addResponseCookieFunc() expr.Option {
	return expr.Function(
		"add_cookie",
		func(params ...any) (any, error) {
			ctx, err := rule.Assert[context.Context](params[0])
			if err != nil {
				return nil, errors.WithStack(err)
			}

			values, err := rule.Assert[map[string]any](params[1])
			if err != nil {
				return nil, errors.WithStack(err)
			}

			cookie, err := cookieFrom(values)
			if err != nil {
				return nil, errors.WithStack(err)
			}

			r, ok := CtxResponse(ctx)
			if !ok {
				return nil, errors.New("could not find http request in context")
			}

			r.Header.Add("Set-Cookie", cookie.String())

			return true, nil
		},
		new(func(context.Context, map[string]any) bool),
	)
}

func getResponseCookieFunc() expr.Option {
	return expr.Function(
		"get_cookie",
		func(params ...any) (any, error) {
			ctx, err := rule.Assert[context.Context](params[0])
			if err != nil {
				return nil, errors.WithStack(err)
			}

			name, err := rule.Assert[string](params[1])
			if err != nil {
				return nil, errors.WithStack(err)
			}

			res, ok := CtxResponse(ctx)
			if !ok {
				return nil, errors.New("could not find http response in context")
			}

			var cookie *http.Cookie
			for _, c := range res.Cookies() {
				if c.Name != name {
					continue
				}

				cookie = c
				break
			}

			if cookie == nil {
				return nil, nil
			}

			return CookieVar{
				Name:     cookie.Name,
				Value:    cookie.Value,
				Path:     cookie.Path,
				Domain:   cookie.Domain,
				Expires:  cookie.Expires,
				MaxAge:   cookie.MaxAge,
				Secure:   cookie.Secure,
				HttpOnly: cookie.HttpOnly,
				SameSite: cookie.SameSite,
			}, nil
		},
		new(func(context.Context, string) CookieVar),
	)
}