Merge pull request #302 from balena-os/lmb/fix-concat-reader

Bugfix: concatReadSeekCloser.Read() with buffers of any size

integration/image/delta_test.go

@@ -205,8 +205,8 @@ var deltaTestCases = []struct {
 		images:      []string{"000", "001", "003"},
 		base:        "001",
 		target:      "003",
-		wantSizeMin: 8736,
-		wantSizeMax: 9248,
+		wantSizeMin: 4689,
+		wantSizeMax: 5201,
 	},
 
 	// Similar to the case above, but we diff between 000 and 003. (003 extends
@@ -234,16 +234,16 @@ var deltaTestCases = []struct {
 	{
 		base:        "006",
 		target:      "007",
-		wantSizeMin: 21888,
-		wantSizeMax: 21888,
+		wantSizeMin: 381,
+		wantSizeMax: 381,
 	},
 
 	// 008 has the same files as 006, but all on a single layer.
 	{
 		base:        "006",
 		target:      "008",
-		wantSizeMin: 21622,
-		wantSizeMax: 21622,
+		wantSizeMin: 115,
+		wantSizeMax: 115,
 	},
 	{
 		base:        "008",
@@ -263,8 +263,8 @@ var deltaTestCases = []struct {
 		images:      []string{"009", "010", "011"},
 		base:        "010",
 		target:      "011",
-		wantSizeMin: 49188,
-		wantSizeMax: 59940,
+		wantSizeMin: 1567,
+		wantSizeMax: 1567,
 	},
 
 	// This simulates the case in which the large base image itself was
@@ -282,7 +282,7 @@ var deltaTestCases = []struct {
 	},
 }
 
-// TestDeltaSizes checks if the sizes of generated deltas have not increased. In
+// TestDeltaSize checks if the sizes of generated deltas have not increased. In
 // other words, this test is designed to catch regressions in the delta sizes.
 //
 // The expected sizes (wantSize) were defined empirically so that they match
@@ -615,8 +615,8 @@ func ttrQueryDeltaSizeAsserting(ctx context.Context, t *testing.T, client apicli
 
 	tarRC, err := client.ImageSave(ctx, []string{ttrImageName(image)})
 	assert.Assert(t, err)
+	defer tarRC.Close()
 	size := deltaSizeFromTar(t, tarRC)
-	tarRC.Close()
 	return size
 }
 

pkg/ioutils/concat.go

@@ -71,20 +71,28 @@ func (self *concatReadSeekCloser) Read(p []byte) (n int, err error) {
 
 	// if the read ends within b
 	if self.off+int64(len(p)) >= self.aSize {
-		if _, err := self.b.Seek(max64(0, self.off-self.aSize), io.SeekStart); err != nil {
+		bOffset := max64(self.off-self.aSize, 0)
+
+		if _, err := self.b.Seek(bOffset, io.SeekStart); err != nil {
 			return 0, err
 		}
 
 		i := clampSliceIndex(self.aSize-self.off, 0, len(p))
-		nB, err := io.ReadFull(self.b, p[i:])
-
-		if err != nil {
-			return 0, err
+		j := clampSliceIndex(int64(i)+self.bSize-bOffset, i, len(p))
+
+		if i != j {
+			nB, err := io.ReadFull(self.b, p[i:j])
+			if err != nil {
+				return 0, err
+			}
+			n += nB
 		}
-		n += nB
 	}
 
 	self.off += int64(n)
+	if self.off == self.aSize+self.bSize {
+		err = io.EOF
+	}
 
 	return
 }

pkg/ioutils/concat_test.go

@@ -0,0 +1,86 @@
+package ioutils
+
+import (
+	"io"
+	"strconv"
+	"strings"
+	"testing"
+
+	"gotest.tools/v3/assert"
+)
+
+// nopCloser wraps an io.ReadSeeker, providing it with a no-op Close method.
+// Like io.NopCloser, but for io.ReadSeekers.
+type nopCloser struct {
+	io.ReadSeeker
+}
+
+func (nopCloser) Close() error { return nil }
+
+// Iterates, reading all data from a ConcatReadSeekCloser, using a buffer of a
+// given size. Checks if no unexpected errors are generated and if all expected
+// data is read.
+func TestConcatReadSeekCloserRead(t *testing.T) {
+	tests := []struct {
+		aContents  string
+		bContents  string
+		bufferSize int
+	}{
+		// Buffers fitting the concatenated readers perfectly
+		{"abcd", "1234", 4},
+		{"abcdef", "123456789", 3},
+		{"abc", "123", 6},
+
+		// Buffers _not_ fitting the concatenated readers perfectly
+		{"abc", "12", 2},
+		{"ab", "123", 2},
+		{"abcde", "123456", 3},
+
+		// Buffers that are larger than either or both concatenated readers
+		{"abc", "def", 10},
+		{"abcdefgh", "1", 2},
+		{"abcdefgh", "1", 3},
+		{"abcdefgh", "1", 8},
+		{"a", "12345678", 2},
+		{"a", "12345678", 3},
+		{"a", "12345678", 8},
+
+		// Either or both readers empty
+		{"", "1234", 4},
+		{"abcd", "", 4},
+		{"", "1234", 3},
+		{"abcd", "", 3},
+		{"", "", 4},
+	}
+
+	for _, tt := range tests {
+		testName := tt.aContents + "/" + tt.bContents + "/" + strconv.Itoa(tt.bufferSize)
+		t.Run(testName, func(t *testing.T) {
+			// Creates two Readers, a and b, with the given contents, and use
+			// them to create the ConcatReadSeekCloser c (the one to be tested).
+			a := nopCloser{strings.NewReader(tt.aContents)}
+			b := nopCloser{strings.NewReader(tt.bContents)}
+
+			c, err := ConcatReadSeekClosers(a, b)
+			assert.NilError(t, err, "error creating ConcatReadSeekCloser")
+
+			// Read until EOF.
+			buf := make([]byte, tt.bufferSize)
+			readData := ""
+
+			for {
+				n, err := c.Read(buf)
+
+				readData += string(buf[:n])
+				if err == io.EOF {
+					break
+				}
+				assert.NilError(t, err, "unexpected read error")
+			}
+
+			// Ensure we got all the data that was expected.
+			expectedReadData := tt.aContents + tt.bContents
+			assert.Equal(t, expectedReadData, readData, "got bad data")
+		})
+	}
+}