From a63a69dfd94f1d4ed0419be0ed5784f90f726f7b Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?M=C3=A1rton=20Elek?= <elek@apache.org>
Date: Thu, 5 Oct 2023 14:40:11 +0200
Subject: [PATCH] satellite/nodeselection: support OR in placement definition

Change-Id: Icc7fd465b28c0c6f09f50c4ab8bffbcc77631dbd
---
 satellite/nodeselection/filter.go   | 27 +++++++++++++++-
 satellite/overlay/placement.go      |  8 +++++
 satellite/overlay/placement_test.go | 48 +++++++++++++++++++++++++++++
 3 files changed, 82 insertions(+), 1 deletion(-)

diff --git a/satellite/nodeselection/filter.go b/satellite/nodeselection/filter.go
index 7a5706da5..3bdd2c6d1 100644
--- a/satellite/nodeselection/filter.go
+++ b/satellite/nodeselection/filter.go
@@ -132,6 +132,27 @@ func (n NodeFilters) Match(node *SelectedNode) bool {
 	return true
 }
 
+// OrFilter will include the node, if at lest one of the filters are matched.
+type OrFilter []NodeFilter
+
+// Match implements NodeFilter interface.
+func (n OrFilter) Match(node *SelectedNode) bool {
+	for _, filter := range n {
+		if filter.Match(node) {
+			return true
+		}
+	}
+	return false
+}
+
+func (n OrFilter) String() string {
+	var parts []string
+	for _, filter := range n {
+		parts = append(parts, fmt.Sprintf("%s", filter))
+	}
+	return "(" + strings.Join(parts, " || ") + ")"
+}
+
 // WithCountryFilter is a helper to create a new filter with additional CountryFilter.
 func (n NodeFilters) WithCountryFilter(permit location.Set) NodeFilters {
 	return append(n, NewCountryFilter(permit))
@@ -143,12 +164,16 @@ func (n NodeFilters) WithExcludedIDs(ds []storj.NodeID) NodeFilters {
 }
 
 func (n NodeFilters) String() string {
+	if len(n) == 1 {
+		return fmt.Sprintf("%s", n[0])
+	}
+
 	var res []string
 	for _, filter := range n {
 		res = append(res, fmt.Sprintf("%s", filter))
 	}
 	sort.Strings(res)
-	return strings.Join(res, " && ")
+	return "(" + strings.Join(res, " && ") + ")"
 }
 
 // GetAnnotation implements NodeFilterWithAnnotation.
diff --git a/satellite/overlay/placement.go b/satellite/overlay/placement.go
index d7e5d840a..9106a0249 100644
--- a/satellite/overlay/placement.go
+++ b/satellite/overlay/placement.go
@@ -121,6 +121,14 @@ func (d *PlacementDefinitions) AddPlacementFromString(definitions string) error
 			res := nodeselection.NodeFilters{filter1, filter2}
 			return res, nil
 		},
+		mito.OpOr: func(env map[any]any, a, b any) (any, error) {
+			filter1, ok1 := a.(nodeselection.NodeFilter)
+			filter2, ok2 := b.(nodeselection.NodeFilter)
+			if !ok1 || !ok2 {
+				return nil, errs.New("OR is supported only between NodeFilter instances")
+			}
+			return nodeselection.OrFilter{filter1, filter2}, nil
+		},
 		"tag": func(nodeIDstr string, key string, value any) (nodeselection.NodeFilters, error) {
 			nodeID, err := storj.NodeIDFromString(nodeIDstr)
 			if err != nil {
diff --git a/satellite/overlay/placement_test.go b/satellite/overlay/placement_test.go
index 90dfa3194..9807db8f7 100644
--- a/satellite/overlay/placement_test.go
+++ b/satellite/overlay/placement_test.go
@@ -248,6 +248,54 @@ func TestPlacementFromString(t *testing.T) {
 
 	})
 
+	t.Run("OR", func(t *testing.T) {
+		p := NewPlacementDefinitions()
+		err := p.AddPlacementFromString(`11:country("GB") || country("DE")`)
+		require.NoError(t, err)
+
+		filters := p.placements[storj.PlacementConstraint(11)]
+		require.NotNil(t, filters)
+		require.True(t, filters.Match(&nodeselection.SelectedNode{
+			CountryCode: location.UnitedKingdom,
+		}))
+		require.True(t, filters.Match(&nodeselection.SelectedNode{
+			CountryCode: location.Germany,
+		}))
+		require.Equal(t, `(country("GB") || country("DE"))`, fmt.Sprintf("%s", filters))
+	})
+
+	t.Run("OR combined with AND", func(t *testing.T) {
+		p := NewPlacementDefinitions()
+		err := p.AddPlacementFromString(`11:((country("GB") || country("DE")) && tag("12whfK1EDvHJtajBiAUeajQLYcWqxcQmdYQU5zX5cCf6bAxfgu4","foo","bar"))`)
+		require.NoError(t, err)
+
+		filters := p.placements[storj.PlacementConstraint(11)]
+		require.NotNil(t, filters)
+		require.False(t, filters.Match(&nodeselection.SelectedNode{
+			CountryCode: location.UnitedKingdom,
+		}))
+		require.False(t, filters.Match(&nodeselection.SelectedNode{
+			Tags: nodeselection.NodeTags{
+				{
+					Signer: signer,
+					Name:   "foo",
+					Value:  []byte("bar"),
+				},
+			},
+		}))
+		require.True(t, filters.Match(&nodeselection.SelectedNode{
+			CountryCode: location.Germany,
+			Tags: nodeselection.NodeTags{
+				{
+					Signer: signer,
+					Name:   "foo",
+					Value:  []byte("bar"),
+				},
+			},
+		}))
+		require.Equal(t, `((country("GB") || country("DE")) && tag("12whfK1EDvHJtajBiAUeajQLYcWqxcQmdYQU5zX5cCf6bAxfgu4","foo","bar"))`, fmt.Sprintf("%s", filters))
+	})
+
 	t.Run("annotation usage", func(t *testing.T) {
 		t.Run("normal", func(t *testing.T) {
 			t.Parallel()