diff --git a/satellite/overlay/service.go b/satellite/overlay/service.go index eec6171e3..2446c05ee 100644 --- a/satellite/overlay/service.go +++ b/satellite/overlay/service.go @@ -446,6 +446,24 @@ func (service *Service) GetNodesNetworkInOrder(ctx context.Context, nodeIDs []st return lastNets, nil } +// GetNodesOutOfPlacement checks if nodes from nodeIDs list are in allowed country according to specified geo placement +// and returns list of node ids which are not. +func (service *Service) GetNodesOutOfPlacement(ctx context.Context, nodeIDs []storj.NodeID, placement storj.PlacementConstraint) (offNodes []storj.NodeID, err error) { + defer mon.Task()(&ctx)(nil) + + nodes, err := service.DownloadSelectionCache.GetNodes(ctx, nodeIDs) + if err != nil { + return nil, err + } + offNodes = make([]storj.NodeID, 0, len(nodeIDs)) + for _, nodeID := range nodeIDs { + if selectedNode, ok := nodes[nodeID]; ok && !placement.AllowedCountry(selectedNode.CountryCode) { + offNodes = append(offNodes, selectedNode.ID) + } + } + return offNodes, nil +} + // FindStorageNodesForGracefulExit searches the overlay network for nodes that meet the provided requirements for graceful-exit requests. func (service *Service) FindStorageNodesForGracefulExit(ctx context.Context, req FindStorageNodesRequest) (_ []*SelectedNode, err error) { defer mon.Task()(&ctx)(&err) diff --git a/satellite/overlay/service_test.go b/satellite/overlay/service_test.go index 50b4c6587..67afeeb90 100644 --- a/satellite/overlay/service_test.go +++ b/satellite/overlay/service_test.go @@ -18,6 +18,7 @@ import ( "storj.io/common/memory" "storj.io/common/pb" "storj.io/common/storj" + "storj.io/common/storj/location" "storj.io/common/testcontext" "storj.io/common/testrand" "storj.io/storj/private/testplanet" @@ -1089,3 +1090,47 @@ func TestUpdateCheckInBelowMinVersionEvent(t *testing.T) { require.True(t, ne2.CreatedAt.After(ne1.CreatedAt)) }) } + +func TestService_GetNodesOutOfPlacement(t *testing.T) { + testplanet.Run(t, testplanet.Config{ + SatelliteCount: 1, StorageNodeCount: 4, UplinkCount: 1, + Reconfigure: testplanet.Reconfigure{ + Satellite: func(log *zap.Logger, index int, config *satellite.Config) { + config.Overlay.Node.AsOfSystemTime.Enabled = false + config.Overlay.Node.AsOfSystemTime.DefaultInterval = 0 + }, + }, + }, func(t *testing.T, ctx *testcontext.Context, planet *testplanet.Planet) { + service := planet.Satellites[0].Overlay.Service + + placement := storj.EU + + nodeIDs := []storj.NodeID{} + for _, node := range planet.StorageNodes { + nodeIDs = append(nodeIDs, node.ID()) + + err := service.TestNodeCountryCode(ctx, node.ID(), location.Poland.String()) + require.NoError(t, err) + } + + require.NoError(t, service.DownloadSelectionCache.Refresh(ctx)) + + offNodes, err := service.GetNodesOutOfPlacement(ctx, nodeIDs, placement) + require.NoError(t, err) + require.Empty(t, offNodes) + + expectedNodeIDs := []storj.NodeID{} + for _, node := range planet.StorageNodes { + expectedNodeIDs = append(expectedNodeIDs, node.ID()) + err := service.TestNodeCountryCode(ctx, node.ID(), location.Brazil.String()) + require.NoError(t, err) + + // we need to refresh cache because node country code was changed + require.NoError(t, service.DownloadSelectionCache.Refresh(ctx)) + + offNodes, err := service.GetNodesOutOfPlacement(ctx, nodeIDs, placement) + require.NoError(t, err) + require.ElementsMatch(t, expectedNodeIDs, offNodes) + } + }) +}