refactor(engine): scrub the whole measurement (#956)

Part of https://github.com/ooni/probe/issues/2297

internal/model/measurement.go

@@ -11,6 +11,8 @@ import (
 	"fmt"
 	"net"
 	"time"
+
+	"github.com/ooni/probe-cli/v3/internal/runtimex"
 )
 
 const (
@@ -168,39 +170,49 @@ func (m *Measurement) AddAnnotation(key, value string) {
 // is not the valid serialization of an IP address.
 var ErrInvalidProbeIP = errors.New("model: invalid probe IP")
 
-// Scrub scrubs the probeIP out of the measurement.
-func (m *Measurement) Scrub(probeIP string) (err error) {
-	// We now behave like we can share everything except the
-	// probe IP, which we instead cannot ever share
-	m.ProbeIP = DefaultProbeIP
-	return m.MaybeRewriteTestKeys(probeIP, json.Marshal)
-}
-
 // Scrubbed is the string that replaces IP addresses.
 const Scrubbed = `[scrubbed]`
 
-// MaybeRewriteTestKeys is the function called by Scrub that
-// ensures that m's serialization doesn't include the IP
-func (m *Measurement) MaybeRewriteTestKeys(
-	currentIP string, marshal func(interface{}) ([]byte, error)) error {
+// ScrubMeasurement removes [currentIP] from [m] by rewriting
+// it in place while preserving the underlying types
+func ScrubMeasurement(m *Measurement, currentIP string) error {
 	if net.ParseIP(currentIP) == nil {
 		return ErrInvalidProbeIP
 	}
-	data, err := marshal(m.TestKeys)
-	if err != nil {
+	m.ProbeIP = DefaultProbeIP
+	m.AddAnnotation("_probe_engine_sanitize_test_keys", "true")
+	if err := scrubTestKeys(m, currentIP); err != nil {
 		return err
 	}
-	// The check using Count is to save an unnecessary copy performed by
-	// ReplaceAll when there are no matches into the body. This is what
-	// we would like the common case to be, meaning that the code has done
-	// its job correctly and has not leaked the IP.
-	bpip := []byte(currentIP)
-	if bytes.Count(data, bpip) <= 0 {
-		return nil
+	testKeys := m.TestKeys
+	m.TestKeys = nil
+	if err := scrubTopLevelKeys(m, currentIP); err != nil {
+		return err
 	}
-	data = bytes.ReplaceAll(data, bpip, []byte(Scrubbed))
-	// We add an annotation such that hopefully later we can measure the
-	// number of cases where we failed to sanitize properly.
-	m.AddAnnotation("_probe_engine_sanitize_test_keys", "true")
-	return json.Unmarshal(data, &m.TestKeys)
+	m.TestKeys = testKeys
+	return nil
+}
+
+// scrubJSONUnmarshalTopLevelKeys allows to mock json.Unmarshal
+var scrubJSONUnmarshalTopLevelKeys = json.Unmarshal
+
+// scrubTopLevelKeys removes [currentIP] from the top-level keys
+// of [m] by rewriting these keys in place.
+func scrubTopLevelKeys(m *Measurement, currentIP string) error {
+	data, err := json.Marshal(m)
+	runtimex.PanicOnError(err, "json.Marshal(m) failed") // m must serialize
+	data = bytes.ReplaceAll(data, []byte(currentIP), []byte(Scrubbed))
+	return scrubJSONUnmarshalTopLevelKeys(data, &m)
+}
+
+// scrubJSONUnmarshalTestKeys allows to mock json.Unmarshal
+var scrubJSONUnmarshalTestKeys = json.Unmarshal
+
+// scrubTestKeys removes [currentIP] from the TestKeys by rewriting
+// them in place while preserving their original type
+func scrubTestKeys(m *Measurement, currentIP string) error {
+	data, err := json.Marshal(m.TestKeys)
+	runtimex.PanicOnError(err, "json.Marshal(m.TestKeys) failed") // m.TestKeys must serialize
+	data = bytes.ReplaceAll(data, []byte(currentIP), []byte(Scrubbed))
+	return scrubJSONUnmarshalTestKeys(data, &m.TestKeys)
 }