feat: add before_script before making http request

ravisuhag · ravisuhag · commit 0d02ea11d4b5 · 2024-10-31T14:42:05.000+05:30
diff --git a/go.mod b/go.mod
@@ -229,7 +229,7 @@ require (
 	golang.org/x/crypto v0.11.0 // indirect
 	golang.org/x/mod v0.12.0 // indirect
 	golang.org/x/net v0.12.0 // indirect
-	golang.org/x/sync v0.3.0 // indirect
+	golang.org/x/sync v0.3.0
 	golang.org/x/sys v0.10.0 // indirect
 	golang.org/x/term v0.10.0 // indirect
 	golang.org/x/text v0.11.0 // indirect
diff --git a/plugins/extractors/http/execute_script.go b/plugins/extractors/http/execute_script.go
@@ -2,6 +2,7 @@ package http
 
 import (
 	"context"
+	"encoding/json"
 	"errors"
 	"fmt"
 	"reflect"
@@ -19,8 +20,7 @@ import (
 	"google.golang.org/protobuf/proto"
 )
 
-func (e *Extractor) executeScript(ctx context.Context, res interface{}, emit plugins.Emit) error {
-	scriptCfg := e.config.Script
+func (e *Extractor) executeScript(ctx context.Context, res interface{}, scriptCfg Script, emit plugins.Emit) error {
 	s, err := tengoutil.NewSecureScript(
 		([]byte)(scriptCfg.Source), e.scriptGlobals(ctx, res, emit),
 	)
@@ -40,12 +40,23 @@ func (e *Extractor) executeScript(ctx context.Context, res interface{}, emit plu
 		return fmt.Errorf("run: %w", err)
 	}
 
+	err = e.convertTengoObjToRequest(c.Get("request").Value())
+	if err != nil {
+		return err
+	}
+
 	return nil
 }
 
 func (e *Extractor) scriptGlobals(ctx context.Context, res interface{}, emit plugins.Emit) map[string]interface{} {
+	req, err := e.convertRequestToTengoObj()
+	if err != nil {
+		e.logger.Error(err.Error())
+	}
+
 	return map[string]interface{}{
 		"recipe_scope": &tengo.String{Value: e.UrnScope},
+		"request":      req,
 		"response":     res,
 		"new_asset": &tengo.UserFunction{
 			Name:  "new_asset",
@@ -68,6 +79,30 @@ func (e *Extractor) scriptGlobals(ctx context.Context, res interface{}, emit plu
 	}
 }
 
+func (e *Extractor) convertTengoObjToRequest(obj interface{}) error {
+	r, err := json.Marshal(obj)
+	if err != nil {
+		return err
+	}
+	err = json.Unmarshal(r, &e.config.Request)
+	if err != nil {
+		return err
+	}
+	return nil
+}
+func (e *Extractor) convertRequestToTengoObj() (tengo.Object, error) {
+	var res map[string]interface{}
+	r, err := json.Marshal(e.config.Request)
+	if err != nil {
+		return nil, err
+	}
+	err = json.Unmarshal(r, &res)
+	if err != nil {
+		return nil, err
+	}
+	return tengo.FromInterface(res)
+}
+
 func newAssetWrapper() tengo.CallableFunc {
 	typeURLs := knownTypeURLs()
 	return func(args ...tengo.Object) (tengo.Object, error) {
diff --git a/plugins/extractors/http/http_extractor.go b/plugins/extractors/http/http_extractor.go
@@ -29,17 +29,20 @@ func init() {
 //go:embed README.md
 var summary string
 
+type Script struct {
+	Engine          string `mapstructure:"engine" validate:"required,oneof=tengo"`
+	Source          string `mapstructure:"source" validate:"required"`
+	MaxAllocs       int64  `mapstructure:"max_allocs" validate:"gt=100" default:"5000"`
+	MaxConstObjects int    `mapstructure:"max_const_objects" validate:"gt=10" default:"500"`
+}
+
 // Config holds the set of configuration for the HTTP extractor.
 type Config struct {
 	Request      RequestConfig `mapstructure:"request"`
 	SuccessCodes []int         `mapstructure:"success_codes" validate:"dive,gte=200,lt=300" default:"[200]"`
 	Concurrency  int           `mapstructure:"concurrency" validate:"gte=1,lte=100" default:"5"`
-	Script       struct {
-		Engine          string `mapstructure:"engine" validate:"required,oneof=tengo"`
-		Source          string `mapstructure:"source" validate:"required"`
-		MaxAllocs       int64  `mapstructure:"max_allocs" validate:"gt=100" default:"5000"`
-		MaxConstObjects int    `mapstructure:"max_const_objects" validate:"gt=10" default:"500"`
-	} `mapstructure:"script"`
+	Script       Script        `mapstructure:"script"`
+	BeforeScript *Script       `mapstructure:"before_script"`
 }
 
 type RequestConfig struct {
@@ -122,12 +125,17 @@ func (e *Extractor) Init(ctx context.Context, config plugins.Config) error {
 // executes the script. The script has access to the response and can use the
 // same to 'emit' assets from within the script.
 func (e *Extractor) Extract(ctx context.Context, emit plugins.Emit) error {
+	if e.config.BeforeScript != nil {
+		if err := e.executeScript(ctx, nil, *e.config.BeforeScript, emit); err != nil {
+			return fmt.Errorf("http extractor: execute script: %w", err)
+		}
+	}
 	res, err := e.executeRequest(ctx, e.config.Request)
 	if err != nil {
 		return fmt.Errorf("http extractor: execute request: %w", err)
 	}
 
-	if err := e.executeScript(ctx, res, emit); err != nil {
+	if err := e.executeScript(ctx, res, e.config.Script, emit); err != nil {
 		return fmt.Errorf("http extractor: execute script: %w", err)
 	}
 
diff --git a/plugins/extractors/http/http_extractor_test.go b/plugins/extractors/http/http_extractor_test.go
@@ -202,6 +202,66 @@ func TestExtract(t *testing.T) {
 				testutils.Respond(t, w, http.StatusOK, `[]`)
 			},
 		},
+		{
+			name: "MatchRequestBeforeScript",
+			rawCfg: map[string]interface{}{
+				"before_script": map[string]interface{}{
+					"engine": "tengo",
+					"source": heredoc.Doc(`
+						fmt := import("fmt")
+						reqs := []
+						reqs = append(reqs, {
+							url: "{{serverURL}}/token",
+							method: "GET",
+							content_type: "application/json",
+							accept: "application/json",
+							timeout: "5s"
+						})
+						  
+						responses := execute_request(reqs...)
+						for r in responses {
+						  if is_error(r) {
+							continue
+						  }
+						  request.Headers = {"Authorization": format("Bearer %s", r.body.data.token)}
+						}
+					`),
+					"max_allocs":        5000,
+					"max_const_objects": 500,
+				},
+				"request": map[string]interface{}{
+					"url":          "{{serverURL}}/api/v2/endpoint",
+					"content_type": "application/json",
+					"accept":       "application/json",
+				},
+				"script": map[string]interface{}{
+					"engine": "tengo",
+					"source": "// do nothing",
+				},
+			},
+			handler: func(t *testing.T, w http.ResponseWriter, r *http.Request) {
+				switch r.URL.Path {
+				case "/token":
+					testutils.Respond(t, w, http.StatusOK,
+						`{"header":{"process_time":0.130462645,"messages":[],"error_code":""},"data":{"token":"testToken123123","service_id":1,"expires_at":1711537963}}`,
+					)
+				case "/api/v2/endpoint":
+					assert.Equal(t, r.Method, http.MethodGet)
+					assert.Equal(t, r.URL.Path, "/api/v2/endpoint")
+					assert.Equal(t, r.URL.RawQuery, "")
+					h := r.Header
+					assert.Equal(t, "", h.Get("Content-Type"))
+					assert.Equal(t, "Bearer testToken123123", h.Get("Authorization"))
+					assert.Equal(t, "application/json", h.Get("Accept"))
+					data, err := io.ReadAll(r.Body)
+					assert.NoError(t, err)
+					assert.Empty(t, data)
+					testutils.Respond(t, w, http.StatusOK, `[]`)
+				default:
+					t.Error("Unexpected HTTP call on", r.URL.Path)
+				}
+			},
+		},
 		{
 			name: "MatchRequestAdvanced",
 			rawCfg: map[string]interface{}{
@@ -890,4 +950,8 @@ func replaceServerURL(cfg map[string]interface{}, serverURL string) {
 	reqCfg["url"] = strings.Replace(reqCfg["url"].(string), "{{serverURL}}", serverURL, 1)
 	scriptCfg := cfg["script"].(map[string]interface{})
 	scriptCfg["source"] = strings.Replace(scriptCfg["source"].(string), "{{serverURL}}", serverURL, -1)
+	beforeScriptCfg, ok := cfg["before_script"].(map[string]interface{})
+	if ok {
+		beforeScriptCfg["source"] = strings.Replace(beforeScriptCfg["source"].(string), "{{serverURL}}", serverURL, -1)
+	}
 }
