From 8e3570265ef19008eb196825428d2198d2ba06fd Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?M=2E=20Mert=20Y=C4=B1ld=C4=B1ran?= <mehmet@up9.com>
Date: Sun, 13 Mar 2022 18:29:28 -0700
Subject: [PATCH] Add `json()` helper support to `redact()` helper's arguments
 (#23)

---
 server/lib/eval.go      | 52 ++++++++++++++++++++++++++++++++++++-----
 server/lib/eval_test.go | 35 +++++++++++++++++++++++++++
 2 files changed, 81 insertions(+), 6 deletions(-)

diff --git a/server/lib/eval.go b/server/lib/eval.go
index 69299df..9053388 100644
--- a/server/lib/eval.go
+++ b/server/lib/eval.go
@@ -6,6 +6,7 @@ package basenine
 
 import (
 	"encoding/base64"
+	"errors"
 	"regexp"
 	"strconv"
 	"strings"
@@ -215,18 +216,57 @@ func _json(args ...interface{}) (interface{}, interface{}) {
 	return args[0], result[0]
 }
 
-func redact(args ...interface{}) (interface{}, interface{}) {
-	obj := args[0]
-	for _, param := range args[2:] {
-		jsonPath, err := jp.ParseString(stringOperand(param))
+func redactRecursively(obj interface{}, paths []string) (newObj interface{}, err error) {
+	newObj = obj
+	for i, path := range paths {
+		var jsonPath jp.Expr
+		jsonPath, err = jp.ParseString(path)
 		if err != nil {
-			continue
+			return
 		}
+
 		result := jsonPath.Get(obj)
 		if len(result) < 1 {
+			err = errors.New("No match")
+			return
+		}
+
+		if i < len(paths)-1 {
+			nextJSON, ok := result[0].(string)
+			if !ok {
+				err = errors.New("Not a string")
+			}
+
+			var nextObj interface{}
+			nextObj, err = oj.ParseString(nextJSON)
+			if err != nil {
+				return
+			}
+
+			var nextReturnObj interface{}
+			nextReturnObj, err = redactRecursively(nextObj, paths[i+1:])
+			if err != nil {
+				return
+			}
+
+			jsonPath.Set(newObj, oj.JSON(nextReturnObj))
+			return
+		}
+
+		jsonPath.Set(newObj, REDACTED)
+	}
+	return
+}
+
+func redact(args ...interface{}) (interface{}, interface{}) {
+	obj := args[0]
+	for _, param := range args[2:] {
+		paths := strings.Split(stringOperand(param), ".json().")
+		var err error
+		obj, err = redactRecursively(obj, paths)
+		if err != nil {
 			continue
 		}
-		jsonPath.Set(obj, REDACTED)
 	}
 	return obj, true
 }
diff --git a/server/lib/eval_test.go b/server/lib/eval_test.go
index 9f153b6..13aa41b 100644
--- a/server/lib/eval_test.go
+++ b/server/lib/eval_test.go
@@ -4,6 +4,8 @@ import (
 	"fmt"
 	"testing"
 
+	"github.com/ohler55/ojg/jp"
+	oj "github.com/ohler55/ojg/oj"
 	"github.com/stretchr/testify/assert"
 )
 
@@ -150,3 +152,36 @@ func TestEval(t *testing.T) {
 		assert.JSONEq(t, row.newJson, newJson)
 	}
 }
+
+func TestEvalRedactRecursive(t *testing.T) {
+	query := `redact("response.body.json().model")`
+	json := `{"response":{"body":"{\"id\":114905,\"model\":\"Camaro\",\"brand\":{\"name\":\"Chevrolet\"},\"year\":2021}"}}`
+	expected := fmt.Sprintf(`{"id":114905,"model":"%s","brand":{"name":"Chevrolet"},"year":2021}`, REDACTED)
+
+	expr, err := Parse(query)
+	if err != nil {
+		t.Fatal(err.Error())
+	}
+
+	_, err = Precompute(expr)
+	if err != nil {
+		t.Fatal(err.Error())
+	}
+
+	truth, newJson, err := Eval(expr, json)
+	if err != nil {
+		t.Fatal(err.Error())
+	}
+
+	assert.True(t, truth)
+
+	newObj, err := oj.ParseString(newJson)
+	assert.Nil(t, err)
+
+	jsonPath, err := jp.ParseString("response.body")
+	assert.Nil(t, err)
+
+	nestedJson := jsonPath.Get(newObj)[0].(string)
+
+	assert.JSONEq(t, expected, nestedJson)
+}