Keep cookies in store

builtin/url.go

@@ -0,0 +1,13 @@
+package builtin
+
+import (
+	"net/url"
+)
+
+func Url(rawURL string) *url.URL {
+	u, err := url.Parse(rawURL)
+	if err != nil {
+		panic(err)
+	}
+	return u
+}

http.go

@@ -35,6 +35,7 @@ const (
 	httpStoreBodyKey     = "body"
 	httpStoreRawBodyKey  = "rawBody"
 	httpStoreHeaderKey   = "headers"
+	httpStoreCookieKey   = "cookies"
 	httpStoreResponseKey = "res"
 )
 
@@ -377,6 +378,21 @@ func (rnr *httpRunner) Run(ctx context.Context, r *httpRequest) error {
 	d[httpStoreRawBodyKey] = string(resBody)
 	d[httpStoreHeaderKey] = res.Header
 
+	cookies := res.Cookies()
+
+	if len(cookies) > 0 {
+		keyMap := make(map[string]*http.Cookie)
+
+		for _, c := range cookies {
+			// If the Domain attribute is not specified, the host is taken over
+			c.Domain = rnr.endpoint.Host
+			keyMap[c.Name] = c
+		}
+
+		d[httpStoreCookieKey] = keyMap
+		rnr.operator.recordToCookie(cookies)
+	}
+
 	rnr.operator.record(map[string]any{
 		string(httpStoreResponseKey): d,
 	})

integration_test.go

@@ -45,6 +45,7 @@ func TestRunUsingHTTPBin(t *testing.T) {
 		book string
 	}{
 		{"testdata/book/httpbin.yml"},
+		{"testdata/book/cookie.yml"},
 	}
 	for _, tt := range tests {
 		t.Run(tt.book, func(t *testing.T) {

operator.go

@@ -7,6 +7,7 @@ import (
 	"io"
 	"math"
 	"math/rand"
+	"net/http"
 	"os"
 	"path/filepath"
 	"sort"
@@ -380,6 +381,10 @@ func (o *operator) recordToLatest(key string, value any) error {
 	return o.store.recordToLatest(key, value)
 }
 
+func (o *operator) recordToCookie(cookies []*http.Cookie) {
+	o.store.recordToCookie(cookies)
+}
+
 func (o *operator) generateID() ID {
 	return ID{
 		Type:        IDTypeRunbook,

option.go

@@ -787,6 +787,7 @@ func setupBuiltinFunctions(opts ...Option) []Option {
 	// Built-in functions are added at the beginning of an option and are overridden by subsequent options
 	return append([]Option{
 		// NOTE: Please add here the built-in functions you want to enable.
+		Func("url", func(v string) *url.URL { return builtin.Url(v) }),
 		Func("urlencode", url.QueryEscape),
 		Func("base64encode", func(v any) string { return base64.StdEncoding.EncodeToString([]byte(cast.ToString(v))) }),
 		Func("base64decode", func(v any) string {

option_test.go

@@ -865,6 +865,7 @@ func TestSetupBuiltinFunctions(t *testing.T) {
 	tests := []struct {
 		fn string
 	}{
+		{"url"},
 		{"urlencode"},
 		{"base64encode"},
 		{"base64decode"},

store.go

@@ -2,8 +2,10 @@ package runn
 
 import (
 	"errors"
+	"net/http"
 	"os"
 	"strings"
+	"time"
 )
 
 const (
@@ -17,6 +19,7 @@ const (
 	storeFuncValue   = "[func]"
 	storeStepRunKey  = "run"
 	storeOutcomeKey  = "outcome"
+	storeCookieKey   = "cookies"
 )
 
 type store struct {
@@ -29,6 +32,7 @@ type store struct {
 	parentVars  map[string]any
 	useMap      bool // Use map syntax in `steps:`.
 	loopIndex   *int
+	cookies     map[string]map[string]*http.Cookie
 }
 
 func (s *store) recordAsMapped(k string, v map[string]any) {
@@ -106,6 +110,28 @@ func (s *store) recordToLatest(key string, value any) error {
 	return errors.New("failed to record")
 }
 
+func (s *store) recordToCookie(cookies []*http.Cookie) {
+	cookieMap := make(map[string]map[string]*http.Cookie)
+	for _, cookie := range cookies {
+		domain := cookie.Domain
+		if domain == "" {
+			domain = "localhost"
+		}
+		keyMap, ok := cookieMap[domain]
+		if !ok || keyMap == nil {
+			keyMap = make(map[string]*http.Cookie)
+		}
+		if !cookie.Expires.IsZero() && cookie.Expires.Before(time.Now()) {
+			// Remove expired cookie
+			delete(keyMap, cookie.Name)
+		} else {
+			keyMap[cookie.Name] = cookie
+		}
+		cookieMap[domain] = keyMap
+	}
+	s.cookies = cookieMap
+}
+
 func (s *store) toNormalizedMap() map[string]any {
 	store := map[string]any{}
 	store[storeEnvKey] = envMap()
@@ -124,6 +150,9 @@ func (s *store) toNormalizedMap() map[string]any {
 	if s.loopIndex != nil {
 		store[loopCountVarKey] = *s.loopIndex
 	}
+	if s.cookies != nil {
+		store[storeCookieKey] = s.cookies
+	}
 	return store
 }
 
@@ -148,14 +177,17 @@ func (s *store) toMap() map[string]any {
 	if s.loopIndex != nil {
 		store[loopCountVarKey] = *s.loopIndex
 	}
+	if s.cookies != nil {
+		store[storeCookieKey] = s.cookies
+	}
 	return store
 }
 
 func (s *store) clearSteps() {
 	s.steps = []map[string]any{}
 	s.stepMapKeys = []string{}
 	s.stepMap = map[string]map[string]any{}
-	// keep vars, bindVars
+	// keep vars, bindVars, cookies
 	s.parentVars = map[string]any{}
 	s.loopIndex = nil
 }