From ae3305270d302052e1f131ed178c8e3a33d4ed50 Mon Sep 17 00:00:00 2001
From: Charles Duffy <charles@dyfis.net>
Date: Tue, 2 Feb 2021 16:08:17 -0600
Subject: [PATCH] Implement GetSectionReader (#10)

---
 reader.go      | 28 ++++++++++++++++++++++++++++
 reader_test.go | 41 +++++++++++++++++++++++++++++++++++++++++
 2 files changed, 69 insertions(+)

diff --git a/reader.go b/reader.go
index 4b82493..8b91f01 100644
--- a/reader.go
+++ b/reader.go
@@ -52,6 +52,11 @@ type Reader struct {
 	pad int64
 }
 
+var (
+	ErrNotSeekable = errors.New("ar: SectionReader only available when wrapping a Seeker")
+	ErrNotReaderAt = errors.New("ar: SectionReader only available when wrapping a ReaderAt")
+)
+
 // Copies read data to r. Strips the global ar header.
 func NewReader(r io.Reader) *Reader {
 	io.CopyN(ioutil.Discard, r, 8) // Discard global header
@@ -153,3 +158,26 @@ func (rd *Reader) Read(b []byte) (n int, err error) {
 
 	return
 }
+
+// Returns a SectionReader that can be used in the future to read this section
+// of the archive. Requires that we be wrapping a ReaderAt that is also a
+// ReadSeeker; will fail otherwise.
+func (rd *Reader) GetSectionReader() (*io.SectionReader, error) {
+	seeker, ok := rd.r.(io.Seeker)
+	if !ok {
+		return nil, ErrNotSeekable
+	}
+	readerAt, ok := rd.r.(io.ReaderAt)
+	if !ok {
+		return nil, ErrNotReaderAt
+	}
+	if rd.nb == 0 {
+		return nil, io.EOF
+	}
+	offset, err := seeker.Seek(0, os.SEEK_CUR)
+	if err != nil {
+		return nil, err
+	}
+	retval := io.NewSectionReader(readerAt, offset, rd.nb)
+	return retval, rd.skipUnread()
+}
diff --git a/reader_test.go b/reader_test.go
index 9fad1d0..d8da4ae 100644
--- a/reader_test.go
+++ b/reader_test.go
@@ -86,6 +86,47 @@ func TestReadBody(t *testing.T) {
 	}
 }
 
+func TestSectionReader(t *testing.T) {
+	f, err := os.Open("./fixtures/multi_archive.a")
+	defer f.Close()
+
+	if err != nil {
+		t.Errorf(err.Error())
+	}
+	reader := NewReader(f)
+	allSectionReaders := make([]*io.SectionReader, 0)
+	for {
+		_, err := reader.Next()
+		if err == io.EOF {
+			break
+		}
+		if err != nil {
+			t.Errorf(err.Error())
+		}
+		sectionReader, err := reader.GetSectionReader()
+		if err != nil {
+			t.Errorf(err.Error())
+		}
+		if sectionReader == nil {
+			t.Errorf("A nil SectionReader was returned by GetSectionReader")
+		}
+		allSectionReaders = append(allSectionReaders, sectionReader)
+	}
+	// Now, _after_ the loop completed, make sure we can go back and read the SectionReaders
+	var buf bytes.Buffer
+	for _, sr := range allSectionReaders {
+		_, err := io.Copy(&buf, sr)
+		if err != nil {
+			t.Errorf(err.Error())
+		}
+	}
+	expected := []byte("Hello world!\nI love lamp.\n")
+	actual := buf.Bytes()
+	if !bytes.Equal(expected, actual) {
+		t.Errorf("Concatted byte buffer should be %s but is %s", expected, actual)
+	}
+}
+
 func TestReadMulti(t *testing.T) {
 	f, err := os.Open("./fixtures/multi_archive.a")
 	defer f.Close()