fix: test encrypted

Author: acha-bill

pkg/feeds/getter_test.go

@@ -28,22 +28,40 @@ func TestGetWrappedChunk(t *testing.T) {
 		t.Fatal("data mismatch")
 	}
 
-	// old format (ts + ref)
-	timestamp := make([]byte, 8)
-	binary.BigEndian.PutUint64(timestamp, 1)
-	ch = soctesting.GenerateMockSOC(t, append(timestamp, wch.Address().Bytes()...)).Chunk()
-
-	err = storer.Put(context.Background(), wch)
-	if err != nil {
-		t.Fatal(err)
+	// old format
+	tt := []struct {
+		name string
+		addr []byte
+	}{
+		{
+			name: "unencrypted",
+			addr: wch.Address().Bytes(),
+		},
+		{
+			name: "encrypted",
+			addr: append(wch.Address().Bytes(), wch.Address().Bytes()...),
+		},
 	}
 
-	wch, err = GetWrappedChunk(context.Background(), storer.ChunkStore(), ch)
-	if err != nil {
-		t.Fatal(err)
-	}
+	for _, tc := range tt {
+		t.Run(tc.name, func(t *testing.T) {
+			timestamp := make([]byte, 8)
+			binary.BigEndian.PutUint64(timestamp, 1)
+			ch = soctesting.GenerateMockSOC(t, append(timestamp, tc.addr...)).Chunk()
 
-	if !bytes.Equal(wch.Data()[8:], []byte("data")) {
-		t.Fatal("data mismatch")
+			err = storer.Put(context.Background(), wch)
+			if err != nil {
+				t.Fatal(err)
+			}
+
+			wch, err = GetWrappedChunk(context.Background(), storer.ChunkStore(), ch)
+			if err != nil {
+				t.Fatal(err)
+			}
+
+			if !bytes.Equal(wch.Data()[8:], []byte("data")) {
+				t.Fatal("data mismatch")
+			}
+		})
 	}
 }

pkg/storage/inmemchunkstore/inmemchunkstore.go

@@ -32,7 +32,7 @@ func (c *ChunkStore) Get(_ context.Context, addr swarm.Address) (swarm.Chunk, er
 	c.mu.Lock()
 	defer c.mu.Unlock()
 
-	chunk, ok := c.chunks[addr.ByteString()]
+	chunk, ok := c.chunks[c.key(addr)]
 	if !ok {
 		return nil, storage.ErrNotFound
 	}
@@ -43,12 +43,12 @@ func (c *ChunkStore) Put(_ context.Context, ch swarm.Chunk) error {
 	c.mu.Lock()
 	defer c.mu.Unlock()
 
-	chunkCount, ok := c.chunks[ch.Address().ByteString()]
+	chunkCount, ok := c.chunks[c.key(ch.Address())]
 	if !ok {
 		chunkCount.chunk = swarm.NewChunk(ch.Address(), ch.Data()).WithStamp(ch.Stamp())
 	}
 	chunkCount.count++
-	c.chunks[ch.Address().ByteString()] = chunkCount
+	c.chunks[c.key(ch.Address())] = chunkCount
 
 	return nil
 }
@@ -57,7 +57,7 @@ func (c *ChunkStore) Has(_ context.Context, addr swarm.Address) (bool, error) {
 	c.mu.Lock()
 	defer c.mu.Unlock()
 
-	_, exists := c.chunks[addr.ByteString()]
+	_, exists := c.chunks[c.key(addr)]
 
 	return exists, nil
 }
@@ -66,12 +66,12 @@ func (c *ChunkStore) Delete(_ context.Context, addr swarm.Address) error {
 	c.mu.Lock()
 	defer c.mu.Unlock()
 
-	chunkCount := c.chunks[addr.ByteString()]
+	chunkCount := c.chunks[c.key(addr)]
 	chunkCount.count--
 	if chunkCount.count <= 0 {
 		delete(c.chunks, addr.ByteString())
 	} else {
-		c.chunks[addr.ByteString()] = chunkCount
+		c.chunks[c.key(addr)] = chunkCount
 	}
 
 	return nil
@@ -81,12 +81,12 @@ func (c *ChunkStore) Replace(_ context.Context, ch swarm.Chunk, emplace bool) er
 	c.mu.Lock()
 	defer c.mu.Unlock()
 
-	chunkCount := c.chunks[ch.Address().ByteString()]
+	chunkCount := c.chunks[c.key(ch.Address())]
 	chunkCount.chunk = ch
 	if emplace {
 		chunkCount.count++
 	}
-	c.chunks[ch.Address().ByteString()] = chunkCount
+	c.chunks[c.key(ch.Address())] = chunkCount
 
 	return nil
 }
@@ -111,3 +111,10 @@ func (c *ChunkStore) Iterate(_ context.Context, fn storage.IterateChunkFn) error
 func (c *ChunkStore) Close() error {
 	return nil
 }
+
+func (c *ChunkStore) key(addr swarm.Address) string {
+	if len(addr.Bytes()) < swarm.HashSize {
+		return addr.ByteString()
+	}
+	return string(addr.Bytes()[:swarm.HashSize])
+}