Go MultiReadSeeker (for stringing together multuple ReadSeeker structs - for example for edge includes using ServeContent)

mrs.go

package mrs

import "io"

type multiReadSeeker struct {
	Readers    []io.ReadSeeker
	Pos        int64
	BytesTotal int64
	BytesRead  int64
	locations  []location
}

type location struct {
	at   int64
	size int64
}

func AbsoluteOffset(offset int64, whence int, current int64, total int64) int64 {
	// make the offset absolute
	switch whence {
	case io.SeekEnd:
		offset = total + offset
	case io.SeekCurrent:
		offset = current + offset
	}
	return Clamp(0, offset, total)
}

func Clamp(a int64, value int64, b int64) int64 {
	if value < a {
		return a
	}
	if value > b {
		return b
	}
	return value
}

func New(readers ...io.ReadSeeker) *multiReadSeeker {
	mrs := &multiReadSeeker{Readers: readers}
	var offt int64
	for _, segment := range readers {
		segmentSize, _ := segment.Seek(0, io.SeekEnd)
		segment.Seek(0, io.SeekStart)
		loc := location{at: offt, size: segmentSize}
		mrs.locations = append(mrs.locations, loc)
		offt += segmentSize
		mrs.BytesTotal += segmentSize
	}
	return mrs
}

func (mrs *multiReadSeeker) Seek(offset int64, whence int) (int64, error) {
	newPos := AbsoluteOffset(offset, whence, mrs.Pos, mrs.BytesTotal)
	mrs.Pos = newPos
	for si, location := range mrs.locations {
		seg := mrs.Readers[si]
		seg.Seek(newPos-location.at, io.SeekStart)
	}
	return newPos, nil
}

func (mrs *multiReadSeeker) Read(p []byte) (int, error) {
	if mrs.Pos >= mrs.BytesTotal {
		return 0, io.EOF
	}

	siuc := -1 // segment index under cursor
	for si, location := range mrs.locations {
		if mrs.Pos >= location.at && mrs.Pos < (location.at+location.size) {
			siuc = si
			break
		}
	}
	// If none of the segments are under cursor, we rolled off the EOF or the
	// readers are empty - we are done.
	if siuc < 0 {
		return 0, io.EOF
	}

	segmentUnderCursor := mrs.Readers[siuc]
	offsetInSegment := mrs.Pos - mrs.locations[siuc].at

	// Seek in the segment to the read position

	segmentUnderCursor.Seek(offsetInSegment, io.SeekStart)
	n, errFromSegment := segmentUnderCursor.Read(p)

	segmentsRemain := siuc < len(mrs.Readers)-1
	if errFromSegment == io.EOF && segmentsRemain {
		errFromSegment = nil
	}
	mrs.Pos += int64(n)
	mrs.BytesRead += int64(n)
	return n, errFromSegment
}

mrs_test.go

package mrs

import "testing"
import "bytes"
import "io"
import "io/ioutil"

func TestMRS(t *testing.T) {
	r1 := bytes.NewReader([]byte("Hello"))
	r2 := bytes.NewReader([]byte(" and goodbye"))
	r3 := bytes.NewReader([]byte(""))
	r4 := bytes.NewReader([]byte(" and hello again!"))
	// 34 bytes total

	mrs := New(r1, r2, r3, r4)

	if mrs.Pos != 0 {
		t.Errorf("New ReadSeeker must have its Pos at 0")
	}

	if mrs.BytesTotal != 34 {
		t.Errorf("Must have a total length of 34, has %d", mrs.BytesTotal)
	}

	if mrs.Pos != 0 {
		t.Errorf("Seeker must be parked at 0, is at %d", mrs.Pos)
	}
	r, _ := mrs.Seek(0, io.SeekStart)
	if r != 0 {
		t.Errorf("Returned offset from Seek() must be 0, was %d", r)
	}
	if mrs.Pos != 0 {
		t.Errorf("Position in the seeker must be 0, is %d", mrs.Pos)
	}

	r, _ = mrs.Seek(2, io.SeekStart)
	if r != 2 {
		t.Errorf("Returned offset from Seek() must be 2, was %d", r)
	}
	if mrs.Pos != 2 {
		t.Errorf("Position in the seeker must be 2, is %d", mrs.Pos)
	}

	r, _ = mrs.Seek(981928, io.SeekStart)
	if r != 34 {
		t.Errorf("Returned offset from Seek() must be 34, was %d", r)
	}
	if mrs.Pos != 34 {
		t.Errorf("Position in the seeker must be 34, is %d", mrs.Pos)
	}

	mrs.Seek(1, io.SeekStart)
	b, err := ioutil.ReadAll(mrs)
	if err != nil {
		t.Errorf("Doing a ReadAll on the mrs must return no errors, was %+v", err)
	}
	if string(b) != "ello and goodbye and hello again!" {
		t.Errorf("Expected to read all except the first byte, but read %s", string(b))
	}
	mrs.Seek(0, io.SeekEnd)
	bytez := make([]byte, 3, 3)
	n, err := mrs.Read(bytez)
	if n != 0 {
		t.Errorf("Should have read 0, read %d", n)
	}
	if err != io.EOF {
		t.Errorf("Should have ended up with EOF, ended up with %+v", err)
	}

	mrs.Seek(0, io.SeekStart)
	onebyte := make([]byte, 1, 1)
	readBytes := 0
	for {
		n, err = mrs.Read(onebyte)
		readBytes += n
		if err == io.EOF {
			break
		}
	}
	if err != io.EOF {
		t.Errorf("Should have ended up with EOF, ended up with %+v", err)
	}
	if readBytes != 34 {
		t.Errorf("Should have read 34 bytes exactly but read %d", readBytes)
	}
}