From 00d153249ba60d2ad1d2fd1a331e350739782423 Mon Sep 17 00:00:00 2001
From: Dominik Richter <dominik.richter@gmail.com>
Date: Thu, 14 Sep 2023 07:59:22 -0700
Subject: [PATCH] =?UTF-8?q?=F0=9F=90=9B=20store+load=20recordings=20with?=
 =?UTF-8?q?=20errors=20(#1731)?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Up until this point, errors were marshalled into `error:{}` into JSON
files. This caused problems when trying to load the structure (because
the error type is not familiar with empty object).

This fixes the issue in the RawData to JSON marshal/unmarshal calls and
additionally handles type problems for int (which is by default read as
a float) and time (which is read as a string). It also fixes a problem
with reconnecting resources when the field in question is an error.

---------

Signed-off-by: Dominik Richter <dominik.richter@gmail.com>
---
 llx/rawdata.go         | 23 +++++++++++++++++++++--
 llx/rawdata_test.go    | 37 +++++++++++++++++++++++++++++++++++++
 providers/recording.go | 27 ++++++---------------------
 3 files changed, 64 insertions(+), 23 deletions(-)

diff --git a/llx/rawdata.go b/llx/rawdata.go
index d4495c9c51..2e334853a6 100644
--- a/llx/rawdata.go
+++ b/llx/rawdata.go
@@ -19,7 +19,7 @@ import (
 type RawData struct {
 	Type  types.Type  `json:"type"`
 	Value interface{} `json:"value"`
-	Error error       `json:"error,omitempty"`
+	Error error       `json:"-"`
 }
 
 // a helper structure exclusively used for json unmarshalling of errors
@@ -28,9 +28,28 @@ type errData struct {
 	Error string `json:"error"`
 }
 
+func (r *RawData) MarshalJSON() ([]byte, error) {
+	if r.Error != nil {
+		return json.Marshal(errData{Error: r.Error.Error()})
+	}
+
+	type rd2 RawData
+	return json.Marshal((*rd2)(r))
+}
+
 func (r *RawData) UnmarshalJSON(data []byte) error {
 	type tmp RawData
-	if err := json.Unmarshal(data, (*tmp)(r)); err == nil {
+	if err := json.Unmarshal(data, (*tmp)(r)); err == nil && r.Type != "" {
+		switch r.Type {
+		case types.Int:
+			r.Value = int64(r.Value.(float64))
+		case types.Time:
+			v, err := time.Parse(time.RFC3339, r.Value.(string))
+			if err != nil {
+				return errors.New("failed to parse time into raw data: " + err.Error())
+			}
+			r.Value = &v
+		}
 		return nil
 	}
 
diff --git a/llx/rawdata_test.go b/llx/rawdata_test.go
index 9371f83edb..5311b6158d 100644
--- a/llx/rawdata_test.go
+++ b/llx/rawdata_test.go
@@ -4,10 +4,13 @@
 package llx
 
 import (
+	"encoding/json"
+	"errors"
 	"testing"
 	"time"
 
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 	"go.mondoo.com/cnquery/types"
 )
 
@@ -172,3 +175,37 @@ func TestSuccess(t *testing.T) {
 		})
 	}
 }
+
+func TestRawData_JSON(t *testing.T) {
+	tests := []*RawData{
+		NilData,
+		BoolTrue,
+		BoolFalse,
+		IntData(0),
+		IntData(123),
+		FloatData(0),
+		FloatData(1.23),
+		StringData(""),
+		StringData("b"),
+		RegexData(""),
+		RegexData("r"),
+		TimeData(time.Time{}),
+		// TODO: the raw comparison here does not come out right, because of nano time
+		// TimeData(now),
+		ArrayData([]interface{}{"a", "b"}, types.String),
+		MapData(map[string]interface{}{"a": "b"}, types.String),
+		{Error: errors.New("test")},
+	}
+
+	for i := range tests {
+		o := tests[i]
+		t.Run(o.String(), func(t *testing.T) {
+			out, err := json.Marshal(o)
+			require.NoError(t, err)
+			var res RawData
+			err = json.Unmarshal(out, &res)
+			require.NoError(t, err)
+			assert.Equal(t, o, &res)
+		})
+	}
+}
diff --git a/providers/recording.go b/providers/recording.go
index a8d1ff3cd9..0f418bb205 100644
--- a/providers/recording.go
+++ b/providers/recording.go
@@ -174,7 +174,6 @@ func LoadRecordingFile(path string) (*recording, error) {
 
 	pres := &res
 	pres.refreshCache()
-	pres.fixTypes()
 
 	if err = pres.reconnectResources(); err != nil {
 		return nil, err
@@ -231,26 +230,6 @@ func (r *recording) refreshCache() {
 	}
 }
 
-// json during the unmarshal step will load some things in a way that we
-// can't process. For example: numbers are loaded as float64, but can also
-// be int64's in MQL. This fixes the loaded types.
-func (r *recording) fixTypes() {
-	for i := range r.Assets {
-		asset := r.Assets[i]
-		for j := range asset.Resources {
-			fixResourceTypes(&asset.Resources[j])
-		}
-	}
-}
-
-func fixResourceTypes(r *resourceRecording) {
-	for _, v := range r.Fields {
-		if v.Type == types.Int {
-			v.Value = int64(v.Value.(float64))
-		}
-	}
-}
-
 func (r *recording) reconnectResources() error {
 	var err error
 	for i := range r.Assets {
@@ -267,6 +246,12 @@ func (r *recording) reconnectResources() error {
 func (r *recording) reconnectResource(asset *assetRecording, resource *resourceRecording) error {
 	var err error
 	for k, v := range resource.Fields {
+		if v.Error != nil {
+			// in this case we have neither type information nor a value
+			resource.Fields[k].Error = v.Error
+			continue
+		}
+
 		typ := types.Type(v.Type)
 		resource.Fields[k].Value, err = tryReconnect(typ, v.Value, resource)
 		if err != nil {