From 5c47a87af7b026b391a29a9dab9afcc5794a3078 Mon Sep 17 00:00:00 2001
From: Simone Basso <bassosimone@gmail.com>
Date: Wed, 31 Mar 2021 14:30:30 +0200
Subject: [PATCH] feat(ooniprobe): discard lists not in selected categories
 (#278)

* feat(ooniprobe): discard lists not in selected categories

One day we may make an integration mistake and for any reason
we may end up with URLs that do not belong to the categories
originally selected by the user. If that happens, it's nice to
have a safety net where we remove URLs that do not belong to
the right category before proceeding with testing.

This diff was conceived while discussing the robustness of
https://github.com/ooni/probe/issues/1299 with @hellais.

* fix behavior and add unit test

* more robust
---
 .../internal/nettests/web_connectivity.go     | 29 ++++++-
 .../nettests/web_connectivity_test.go         | 77 +++++++++++++++++++
 2 files changed, 104 insertions(+), 2 deletions(-)
 create mode 100644 cmd/ooniprobe/internal/nettests/web_connectivity_test.go

diff --git a/cmd/ooniprobe/internal/nettests/web_connectivity.go b/cmd/ooniprobe/internal/nettests/web_connectivity.go
index fc752b2..c304701 100644
--- a/cmd/ooniprobe/internal/nettests/web_connectivity.go
+++ b/cmd/ooniprobe/internal/nettests/web_connectivity.go
@@ -9,6 +9,30 @@ import (
 	"github.com/ooni/probe-cli/v3/internal/engine/model"
 )
 
+// preventMistakes makes the code more robust with respect to any possible
+// integration issue where the backend returns to us URLs that don't
+// belong to the category codes we requested.
+func preventMistakes(input []model.URLInfo, categories []string) (output []model.URLInfo) {
+	if len(categories) <= 0 {
+		return input
+	}
+	for _, entry := range input {
+		var found bool
+		for _, cat := range categories {
+			if entry.CategoryCode == cat {
+				found = true
+				break
+			}
+		}
+		if !found {
+			log.Warnf("URL %+v not in %+v; skipping", entry, categories)
+			continue
+		}
+		output = append(output, entry)
+	}
+	return
+}
+
 func lookupURLs(ctl *Controller, categories []string) ([]string, map[int64]int64, error) {
 	inputloader := &engine.InputLoader{
 		CheckInConfig: &model.CheckInConfig{
@@ -29,11 +53,12 @@ func lookupURLs(ctl *Controller, categories []string) ([]string, map[int64]int64
 	}
 	log.Infof("Calling CheckIn API with %s runType", ctl.RunType)
 	testlist, err := inputloader.Load(context.Background())
-	var urls []string
-	urlIDMap := make(map[int64]int64)
 	if err != nil {
 		return nil, nil, err
 	}
+	testlist = preventMistakes(testlist, categories)
+	var urls []string
+	urlIDMap := make(map[int64]int64)
 	for idx, url := range testlist {
 		log.Debugf("Going over URL %d", idx)
 		urlID, err := database.CreateOrUpdateURL(
diff --git a/cmd/ooniprobe/internal/nettests/web_connectivity_test.go b/cmd/ooniprobe/internal/nettests/web_connectivity_test.go
new file mode 100644
index 0000000..5c17c64
--- /dev/null
+++ b/cmd/ooniprobe/internal/nettests/web_connectivity_test.go
@@ -0,0 +1,77 @@
+package nettests
+
+import (
+	"testing"
+
+	"github.com/google/go-cmp/cmp"
+	"github.com/ooni/probe-cli/v3/internal/engine/model"
+)
+
+func TestPreventMistakesWithCategories(t *testing.T) {
+	input := []model.URLInfo{{
+		CategoryCode: "NEWS",
+		URL:          "https://repubblica.it/",
+		CountryCode:  "IT",
+	}, {
+		CategoryCode: "HACK",
+		URL:          "https://2600.com",
+		CountryCode:  "XX",
+	}, {
+		CategoryCode: "FILE",
+		URL:          "https://addons.mozilla.org/",
+		CountryCode:  "XX",
+	}}
+	desired := []model.URLInfo{{
+		CategoryCode: "NEWS",
+		URL:          "https://repubblica.it/",
+		CountryCode:  "IT",
+	}, {
+		CategoryCode: "FILE",
+		URL:          "https://addons.mozilla.org/",
+		CountryCode:  "XX",
+	}}
+	output := preventMistakes(input, []string{"NEWS", "FILE"})
+	if diff := cmp.Diff(desired, output); diff != "" {
+		t.Fatal(diff)
+	}
+}
+
+func TestPreventMistakesWithoutCategoriesAndNil(t *testing.T) {
+	input := []model.URLInfo{{
+		CategoryCode: "NEWS",
+		URL:          "https://repubblica.it/",
+		CountryCode:  "IT",
+	}, {
+		CategoryCode: "HACK",
+		URL:          "https://2600.com",
+		CountryCode:  "XX",
+	}, {
+		CategoryCode: "FILE",
+		URL:          "https://addons.mozilla.org/",
+		CountryCode:  "XX",
+	}}
+	output := preventMistakes(input, nil)
+	if diff := cmp.Diff(input, output); diff != "" {
+		t.Fatal(diff)
+	}
+}
+
+func TestPreventMistakesWithoutCategoriesAndEmpty(t *testing.T) {
+	input := []model.URLInfo{{
+		CategoryCode: "NEWS",
+		URL:          "https://repubblica.it/",
+		CountryCode:  "IT",
+	}, {
+		CategoryCode: "HACK",
+		URL:          "https://2600.com",
+		CountryCode:  "XX",
+	}, {
+		CategoryCode: "FILE",
+		URL:          "https://addons.mozilla.org/",
+		CountryCode:  "XX",
+	}}
+	output := preventMistakes(input, []string{})
+	if diff := cmp.Diff(input, output); diff != "" {
+		t.Fatal(diff)
+	}
+}