feat(rewriter): add redirect(), get_cookie(), add_cookie() methods to rule engine (#36)

internal/rule/http/context.go

@@ -22,10 +22,10 @@ func WithResponse(ctx context.Context, r *http.Response) context.Context {
 	return context.WithValue(ctx, contextKeyResponse, r)
 }
 
-func ctxRequest(ctx context.Context) (*http.Request, bool) {
+func CtxRequest(ctx context.Context) (*http.Request, bool) {
 	return rule.Context[*http.Request](ctx, contextKeyRequest)
 }
 
-func ctxResponse(ctx context.Context) (*http.Response, bool) {
+func CtxResponse(ctx context.Context) (*http.Response, bool) {
 	return rule.Context[*http.Response](ctx, contextKeyResponse)
 }

internal/rule/http/option.go

@@ -13,6 +13,8 @@ func WithRequestFuncs() rule.OptionFunc {
 			addRequestHeaderFunc(),
 			delRequestHeadersFunc(),
 			setRequestHostFunc(),
+			getRequestCookieFunc(),
+			addRequestCookieFunc(),
 		}
 
 		if len(opts.Expr) == 0 {
@@ -29,6 +31,8 @@ func WithResponseFuncs() rule.OptionFunc {
 			setResponseHeaderFunc(),
 			addResponseHeaderFunc(),
 			delResponseHeadersFunc(),
+			addResponseCookieFunc(),
+			getResponseCookieFunc(),
 		}
 
 		if len(opts.Expr) == 0 {

internal/rule/http/request.go

@@ -3,6 +3,7 @@ package http
 import (
 	"context"
 	"fmt"
+	"net/http"
 	"net/url"
 	"strconv"
 	"strings"
@@ -28,7 +29,7 @@ func setRequestHostFunc() expr.Option {
 				return nil, errors.WithStack(err)
 			}
 
-			r, ok := ctxRequest(ctx)
+			r, ok := CtxRequest(ctx)
 			if !ok {
 				return nil, errors.New("could not find http request in context")
 			}
@@ -60,7 +61,7 @@ func setRequestURLFunc() expr.Option {
 				return false, errors.WithStack(err)
 			}
 
-			r, ok := ctxRequest(ctx)
+			r, ok := CtxRequest(ctx)
 			if !ok {
 				return nil, errors.New("could not find http request in context")
 			}
@@ -89,7 +90,7 @@ func addRequestHeaderFunc() expr.Option {
 
 			value := formatValue(params[2])
 
-			r, ok := ctxRequest(ctx)
+			r, ok := CtxRequest(ctx)
 			if !ok {
 				return nil, errors.New("could not find http request in context")
 			}
@@ -118,7 +119,7 @@ func setRequestHeaderFunc() expr.Option {
 
 			value := formatValue(params[2])
 
-			r, ok := ctxRequest(ctx)
+			r, ok := CtxRequest(ctx)
 			if !ok {
 				return nil, errors.New("could not find http request in context")
 			}
@@ -145,7 +146,7 @@ func delRequestHeadersFunc() expr.Option {
 				return nil, errors.WithStack(err)
 			}
 
-			r, ok := ctxRequest(ctx)
+			r, ok := CtxRequest(ctx)
 			if !ok {
 				return nil, errors.New("could not find http request in context")
 			}
@@ -167,6 +168,145 @@ func delRequestHeadersFunc() expr.Option {
 	)
 }
 
+type CookieVar struct {
+	Name     string        `expr:"name"`
+	Value    string        `expr:"value"`
+	Path     string        `expr:"path"`
+	Domain   string        `expr:"domain"`
+	Expires  time.Time     `expr:"expires"`
+	MaxAge   int           `expr:"max_age"`
+	Secure   bool          `expr:"secure"`
+	HttpOnly bool          `expr:"http_only"`
+	SameSite http.SameSite `expr:"same_site"`
+}
+
+func getRequestCookieFunc() 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)
+			}
+
+			r, ok := CtxRequest(ctx)
+			if !ok {
+				return nil, errors.New("could not find http request in context")
+			}
+
+			cookie, err := r.Cookie(name)
+			if err != nil && !errors.Is(err, http.ErrNoCookie) {
+				return nil, errors.WithStack(err)
+			}
+
+			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),
+	)
+}
+
+func addRequestCookieFunc() 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 := CtxRequest(ctx)
+			if !ok {
+				return nil, errors.New("could not find http request in context")
+			}
+
+			r.AddCookie(cookie)
+
+			return true, nil
+		},
+		new(func(context.Context, map[string]any) bool),
+	)
+}
+
+func cookieFrom(values map[string]any) (*http.Cookie, error) {
+	cookie := &http.Cookie{}
+
+	if name, ok := values["name"].(string); ok {
+		cookie.Name = name
+	}
+
+	if value, ok := values["value"].(string); ok {
+		cookie.Value = value
+	}
+
+	if domain, ok := values["domain"].(string); ok {
+		cookie.Domain = domain
+	}
+
+	if path, ok := values["path"].(string); ok {
+		cookie.Path = path
+	}
+
+	if httpOnly, ok := values["http_only"].(bool); ok {
+		cookie.HttpOnly = httpOnly
+	}
+
+	if maxAge, ok := values["max_age"].(int); ok {
+		cookie.MaxAge = maxAge
+	}
+
+	if secure, ok := values["secure"].(bool); ok {
+		cookie.Secure = secure
+	}
+
+	if sameSite, ok := values["same_site"].(http.SameSite); ok {
+		cookie.SameSite = sameSite
+	} else if sameSite, ok := values["same_site"].(int); ok {
+		cookie.SameSite = http.SameSite(sameSite)
+	}
+
+	if expires, ok := values["expires"].(time.Time); ok {
+		cookie.Expires = expires
+	} else if rawExpires, ok := values["expires"].(string); ok {
+		expires, err := time.Parse(http.TimeFormat, rawExpires)
+		if err != nil {
+			return nil, errors.WithStack(err)
+		}
+
+		cookie.Expires = expires
+	}
+
+	return cookie, nil
+}
+
 func formatValue(v any) string {
 	var value string
 	switch v := v.(type) {

internal/rule/http/request_test.go

@@ -2,6 +2,7 @@ package http
 
 import (
 	"context"
+	"fmt"
 	"net/http"
 	"testing"
 
@@ -185,6 +186,134 @@ func TestDelRequestHeaders(t *testing.T) {
 	}
 }
 
+func TestAddRequestCookie(t *testing.T) {
+	type TestCase struct {
+		Cookie     map[string]any
+		Check      func(t *testing.T, tc TestCase, req *http.Request)
+		ShouldFail bool
+	}
+
+	testCases := []TestCase{
+		{
+			Cookie: map[string]any{
+				"name": "test",
+			},
+			Check: func(t *testing.T, tc TestCase, req *http.Request) {
+				cookie, err := req.Cookie(tc.Cookie["name"].(string))
+				if err != nil {
+					t.Errorf("%+v", errors.WithStack(err))
+					return
+				}
+
+				if e, g := tc.Cookie["name"], cookie.Name; e != g {
+					t.Errorf("cookie.Name: expected '%v', got '%v'", e, g)
+				}
+			},
+		},
+		{
+			Cookie: map[string]any{
+				"name":  "foo",
+				"value": "test",
+			},
+			Check: func(t *testing.T, tc TestCase, req *http.Request) {
+				cookie, err := req.Cookie(tc.Cookie["name"].(string))
+				if err != nil {
+					t.Errorf("%+v", errors.WithStack(err))
+					return
+				}
+
+				if e, g := tc.Cookie["name"], cookie.Name; e != g {
+					t.Errorf("cookie.Name: expected '%v', got '%v'", e, g)
+				}
+
+				if e, g := tc.Cookie["value"], cookie.Value; e != g {
+					t.Errorf("cookie.Value: expected '%v', got '%v'", e, g)
+				}
+			},
+		},
+	}
+
+	for idx, tc := range testCases {
+		t.Run(fmt.Sprintf("Case_%d", idx), func(t *testing.T) {
+			type Vars struct {
+				NewCookie map[string]any `expr:"new_cookie"`
+			}
+
+			engine := createRuleEngine[Vars](t,
+				rule.WithExpr(addRequestCookieFunc()),
+				rule.WithRules(
+					`add_cookie(ctx, vars.new_cookie)`,
+				),
+			)
+
+			req, err := http.NewRequest("GET", "http://example.net", nil)
+			if err != nil {
+				t.Fatalf("%+v", errors.WithStack(err))
+			}
+
+			vars := Vars{
+				NewCookie: tc.Cookie,
+			}
+
+			ctx := context.Background()
+			ctx = WithRequest(ctx, req)
+
+			if _, err := engine.Apply(ctx, vars); err != nil {
+				t.Fatalf("%+v", errors.WithStack(err))
+			}
+
+			if tc.ShouldFail {
+				t.Error("engine.Apply() should have failed")
+			}
+
+			if tc.Check != nil {
+				tc.Check(t, tc, req)
+			}
+		})
+	}
+}
+
+func TestGetRequestCookie(t *testing.T) {
+	type Vars struct {
+		CookieName string `expr:"cookieName"`
+	}
+
+	engine := createRuleEngine[Vars](t,
+		rule.WithExpr(getRequestCookieFunc()),
+		rule.WithRules(
+			"let cookie = get_cookie(ctx, vars.cookieName); cookie.value",
+		),
+	)
+
+	req, err := http.NewRequest("GET", "http://example.net", nil)
+	if err != nil {
+		t.Fatalf("%+v", errors.WithStack(err))
+	}
+
+	vars := Vars{
+		CookieName: "foo",
+	}
+
+	cookie := &http.Cookie{
+		Name:  vars.CookieName,
+		Value: "bar",
+	}
+
+	req.AddCookie(cookie)
+
+	ctx := context.Background()
+	ctx = WithRequest(ctx, req)
+
+	results, err := engine.Apply(ctx, vars)
+	if err != nil {
+		t.Fatalf("%+v", errors.WithStack(err))
+	}
+
+	if e, g := cookie.Value, results[0]; e != g {
+		t.Errorf("result[0]: expected '%v', got '%v'", e, g)
+	}
+}
+
 func createRuleEngine[V any](t *testing.T, funcs ...rule.OptionFunc) *rule.Engine[V] {
 	engine, err := rule.NewEngine[V](funcs...)
 	if err != nil {

internal/rule/http/response.go

@@ -3,6 +3,7 @@ package http
 import (
 	"context"
 	"fmt"
+	"net/http"
 	"strconv"
 	"strings"
 	"time"
@@ -41,7 +42,7 @@ func addResponseHeaderFunc() expr.Option {
 				value = fmt.Sprintf("%v", rawValue)
 			}
 
-			r, ok := ctxResponse(ctx)
+			r, ok := CtxResponse(ctx)
 			if !ok {
 				return nil, errors.New("could not find http response in context")
 			}
@@ -82,7 +83,7 @@ func setResponseHeaderFunc() expr.Option {
 				value = fmt.Sprintf("%v", rawValue)
 			}
 
-			r, ok := ctxResponse(ctx)
+			r, ok := CtxResponse(ctx)
 			if !ok {
 				return nil, errors.New("could not find http response in context")
 			}
@@ -109,7 +110,7 @@ func delResponseHeadersFunc() expr.Option {
 				return nil, errors.WithStack(err)
 			}
 
-			r, ok := ctxResponse(ctx)
+			r, ok := CtxResponse(ctx)
 			if !ok {
 				return nil, errors.New("could not find http response in context")
 			}
@@ -130,3 +131,84 @@ func delResponseHeadersFunc() expr.Option {
 		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),
+	)
+}

internal/rule/http/response_test.go

@@ -2,9 +2,11 @@ package http
 
 import (
 	"context"
+	"fmt"
 	"io"
 	"net/http"
 	"testing"
+	"time"
 
 	"forge.cadoles.com/cadoles/bouncer/internal/rule"
 	"github.com/pkg/errors"
@@ -124,6 +126,182 @@ func TestResponseDelHeaders(t *testing.T) {
 	}
 }
 
+func TestAddResponseCookie(t *testing.T) {
+	type TestCase struct {
+		Cookie     map[string]any
+		Check      func(t *testing.T, tc TestCase, res *http.Response)
+		ShouldFail bool
+	}
+
+	testCases := []TestCase{
+		{
+			Cookie: map[string]any{
+				"name":      "foo",
+				"value":     "test",
+				"domain":    "example.net",
+				"path":      "/custom",
+				"same_site": http.SameSiteStrictMode,
+				"http_only": true,
+				"secure":    false,
+				"expires":   time.Now().UTC().Truncate(time.Second),
+			},
+			Check: func(t *testing.T, tc TestCase, res *http.Response) {
+				var cookie *http.Cookie
+				for _, c := range res.Cookies() {
+					if c.Name == tc.Cookie["name"] {
+						cookie = c
+						break
+					}
+				}
+				if cookie == nil {
+					t.Errorf("could not find cookie '%s'", tc.Cookie["name"])
+					return
+				}
+
+				if e, g := tc.Cookie["name"], cookie.Name; e != g {
+					t.Errorf("cookie.Name: expected '%v', got '%v'", e, g)
+				}
+
+				if e, g := tc.Cookie["value"], cookie.Value; e != g {
+					t.Errorf("cookie.Value: expected '%v', got '%v'", e, g)
+				}
+
+				if e, g := tc.Cookie["domain"], cookie.Domain; e != g {
+					t.Errorf("cookie.Domain: expected '%v', got '%v'", e, g)
+				}
+
+				if e, g := tc.Cookie["path"], cookie.Path; e != g {
+					t.Errorf("cookie.Path: expected '%v', got '%v'", e, g)
+				}
+
+				if e, g := tc.Cookie["secure"], cookie.Secure; e != g {
+					t.Errorf("cookie.Secure: expected '%v', got '%v'", e, g)
+				}
+
+				if e, g := tc.Cookie["http_only"], cookie.HttpOnly; e != g {
+					t.Errorf("cookie.HttpOnly: expected '%v', got '%v'", e, g)
+				}
+
+				if e, g := tc.Cookie["same_site"], cookie.SameSite; e != g {
+					t.Errorf("cookie.SameSite: expected '%v', got '%v'", e, g)
+				}
+
+				if e, g := tc.Cookie["expires"], cookie.Expires; e != g {
+					t.Errorf("cookie.Expires: expected '%v', got '%v'", e, g)
+				}
+			},
+		},
+		{
+			Cookie: map[string]any{
+				"name":    "foo",
+				"expires": time.Now().UTC().Format(http.TimeFormat),
+			},
+			Check: func(t *testing.T, tc TestCase, res *http.Response) {
+				var cookie *http.Cookie
+				for _, c := range res.Cookies() {
+					if c.Name == tc.Cookie["name"] {
+						cookie = c
+						break
+					}
+				}
+				if cookie == nil {
+					t.Errorf("could not find cookie '%s'", tc.Cookie["name"])
+					return
+				}
+
+				if e, g := tc.Cookie["expires"], cookie.Expires.Format(http.TimeFormat); e != g {
+					t.Errorf("cookie.Expires: expected '%v', got '%v'", e, g)
+				}
+			},
+		},
+	}
+
+	for idx, tc := range testCases {
+		t.Run(fmt.Sprintf("Case_%d", idx), func(t *testing.T) {
+			type Vars struct {
+				NewCookie map[string]any `expr:"new_cookie"`
+			}
+
+			engine := createRuleEngine[Vars](t,
+				rule.WithExpr(addResponseCookieFunc()),
+				rule.WithRules(
+					`add_cookie(ctx, vars.new_cookie)`,
+				),
+			)
+
+			req, err := http.NewRequest("GET", "http://example.net", nil)
+			if err != nil {
+				t.Fatalf("%+v", errors.WithStack(err))
+			}
+
+			resp := createResponse(req, http.StatusOK, nil)
+
+			vars := Vars{
+				NewCookie: tc.Cookie,
+			}
+
+			ctx := context.Background()
+			ctx = WithRequest(ctx, req)
+			ctx = WithResponse(ctx, resp)
+
+			if _, err := engine.Apply(ctx, vars); err != nil {
+				t.Fatalf("%+v", errors.WithStack(err))
+			}
+
+			if tc.ShouldFail {
+				t.Error("engine.Apply() should have failed")
+			}
+
+			if tc.Check != nil {
+				tc.Check(t, tc, resp)
+			}
+		})
+	}
+}
+
+func TestGetResponseCookie(t *testing.T) {
+	type Vars struct {
+		CookieName string `expr:"cookieName"`
+	}
+
+	engine := createRuleEngine[Vars](t,
+		rule.WithExpr(getResponseCookieFunc()),
+		rule.WithRules(
+			"let cookie = get_cookie(ctx, vars.cookieName); cookie.value",
+		),
+	)
+
+	req, err := http.NewRequest("GET", "http://example.net", nil)
+	if err != nil {
+		t.Fatalf("%+v", errors.WithStack(err))
+	}
+
+	resp := createResponse(req, http.StatusOK, nil)
+
+	vars := Vars{
+		CookieName: "foo",
+	}
+
+	cookie := &http.Cookie{
+		Name:  vars.CookieName,
+		Value: "bar",
+	}
+
+	resp.Header.Add("Set-Cookie", cookie.String())
+
+	ctx := context.Background()
+	ctx = WithResponse(ctx, resp)
+
+	results, err := engine.Apply(ctx, vars)
+	if err != nil {
+		t.Fatalf("%+v", errors.WithStack(err))
+	}
+
+	if e, g := cookie.Value, results[0]; e != g {
+		t.Errorf("result[0]: expected '%v', got '%v'", e, g)
+	}
+}
+
 func createResponse(req *http.Request, statusCode int, body io.Reader) *http.Response {
 	return &http.Response{
 		Status:        http.StatusText(statusCode),