Add chunk addresses getter iterator (#932)

pkg/file/addresses/addresses_getter.go

+// Copyright 2020 The Swarm Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+package addresses
+
+import (
+	"context"
+	"errors"
+
+	"github.com/ethersphere/bee/pkg/storage"
+	"github.com/ethersphere/bee/pkg/swarm"
+)
+
+var (
+	// ErrStopIterator is returned iterator function marks that iteration should
+	// be stopped.
+	ErrStopIterator = errors.New("stop iterator")
+)
+
+type addressesGetterStore struct {
+	getter storage.Getter
+	fn     swarm.AddressIterFunc
+}
+
+// NewGetter creates a new proxy storage.Getter which calls provided function
+// for each chunk address processed.
+func NewGetter(getter storage.Getter, fn swarm.AddressIterFunc) storage.Getter {
+	return &addressesGetterStore{getter, fn}
+}
+
+func (s *addressesGetterStore) Get(ctx context.Context, mode storage.ModeGet, addr swarm.Address) (ch swarm.Chunk, err error) {
+	ch, err = s.getter.Get(ctx, mode, addr)
+	if err != nil {
+		return
+	}
+
+	stop := s.fn(ch.Address())
+	if stop {
+		return ch, ErrStopIterator
+	}
+
+	return
+}

pkg/file/addresses/addresses_getter_test.go

+// Copyright 2020 The Swarm Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+package addresses_test
+
+import (
+	"context"
+	"io/ioutil"
+	"sync"
+	"testing"
+	"time"
+
+	"github.com/ethersphere/bee/pkg/file"
+	"github.com/ethersphere/bee/pkg/file/addresses"
+	"github.com/ethersphere/bee/pkg/file/joiner"
+	filetest "github.com/ethersphere/bee/pkg/file/testing"
+	"github.com/ethersphere/bee/pkg/storage"
+	"github.com/ethersphere/bee/pkg/storage/mock"
+	"github.com/ethersphere/bee/pkg/swarm"
+)
+
+func TestAddressesGetterIterateChunkAddresses(t *testing.T) {
+	store := mock.NewStorer()
+
+	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
+	defer cancel()
+
+	// create root chunk with 2 references and the referenced data chunks
+	rootChunk := filetest.GenerateTestRandomFileChunk(swarm.ZeroAddress, swarm.ChunkSize*2, swarm.SectionSize*2)
+	_, err := store.Put(ctx, storage.ModePutUpload, rootChunk)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	firstAddress := swarm.NewAddress(rootChunk.Data()[8 : swarm.SectionSize+8])
+	firstChunk := filetest.GenerateTestRandomFileChunk(firstAddress, swarm.ChunkSize, swarm.ChunkSize)
+	_, err = store.Put(ctx, storage.ModePutUpload, firstChunk)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	secondAddress := swarm.NewAddress(rootChunk.Data()[swarm.SectionSize+8:])
+	secondChunk := filetest.GenerateTestRandomFileChunk(secondAddress, swarm.ChunkSize, swarm.ChunkSize)
+	_, err = store.Put(ctx, storage.ModePutUpload, secondChunk)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	createdAddresses := []swarm.Address{rootChunk.Address(), firstAddress, secondAddress}
+
+	foundAddresses := make(map[string]struct{})
+	var foundAddressesMu sync.Mutex
+
+	addressIterFunc := func(addr swarm.Address) (stop bool) {
+		foundAddressesMu.Lock()
+		foundAddresses[addr.String()] = struct{}{}
+		foundAddressesMu.Unlock()
+		return false
+	}
+
+	addressesGetter := addresses.NewGetter(store, addressIterFunc)
+
+	j, _, err := joiner.New(ctx, addressesGetter, rootChunk.Address())
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	_, err = file.JoinReadAll(ctx, j, ioutil.Discard)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	if len(createdAddresses) != len(foundAddresses) {
+		t.Fatalf("expected to find %d addresses, got %d", len(createdAddresses), len(foundAddresses))
+	}
+
+	checkAddressFound := func(t *testing.T, foundAddresses map[string]struct{}, address swarm.Address) {
+		t.Helper()
+
+		if _, ok := foundAddresses[address.String()]; !ok {
+			t.Fatalf("expected address %s not found", address.String())
+		}
+	}
+
+	for _, createdAddress := range createdAddresses {
+		checkAddressFound(t, foundAddresses, createdAddress)
+	}
+}

pkg/swarm/swarm.go

@@ -103,6 +103,10 @@ func (a Address) MarshalJSON() ([]byte, error) {
 // ZeroAddress is the address that has no value.
 var ZeroAddress = NewAddress(nil)
 
+// AddressIterFunc is a callback on every address that is found by the iterator.
+// By returning a true for stop variable, iteration should stop.
+type AddressIterFunc func(address Address) (stop bool)
+
 type Chunk interface {
 	Address() Address
 	Data() []byte