Tweaks to make FromReader work

This commit is contained in:
Caleb Gardner
2021-09-27 01:27:58 -05:00
parent 80ff4466ae
commit e8a8c531a9
3 changed files with 55 additions and 55 deletions
+36 -34
View File
@@ -1,23 +1,28 @@
package rawreader package rawreader
import ( import (
"bytes"
"errors" "errors"
"io" "io"
) )
func ConvertReader(r io.Reader) RawReader { func ConvertReader(r io.Reader) (RawReader, error) {
if rr, ok := r.(RawReader); ok { if rr, ok := r.(RawReader); ok {
return rr return rr, nil
} }
if rs, is := r.(io.ReadSeeker); is { if rs, is := r.(io.ReadSeeker); is {
return &fromReadSeeker{ return &fromReadSeeker{
ReadSeeker: rs, ReadSeeker: rs,
}, nil
} }
buf := new(bytes.Buffer)
_, err := io.Copy(buf, r)
if err != nil {
return nil, err
} }
return &fromReader{ return &fromReader{
rdr: r, data: buf.Bytes(),
cache: make([]byte, 0), }, nil
}
} }
func ConvertReaderAt(r io.ReaderAt) RawReader { func ConvertReaderAt(r io.ReaderAt) RawReader {
@@ -37,27 +42,21 @@ type RawReader interface {
} }
type fromReader struct { type fromReader struct {
rdr io.Reader data []byte
cache []byte
off int off int
} }
func (r *fromReader) increaseCache(len int) error {
newCache := make([]byte, len)
_, err := r.rdr.Read(newCache)
if err != nil {
return err
}
r.cache = append(r.cache, newCache...)
return nil
}
func (r *fromReader) ReadAt(p []byte, off int64) (n int, err error) { func (r *fromReader) ReadAt(p []byte, off int64) (n int, err error) {
if int(off)+len(p) > len(r.cache) { n = len(p)
r.increaseCache((int(off) + len(p)) - len(r.cache)) if int(off)+len(p) > len(r.data) {
n = len(r.data) - int(off)
err = io.EOF
} }
for i := int64(0); i < int64(len(p)); i++ { if n < 0 {
p[i] = r.cache[off+i] n = 0
}
for i := 0; i < n; i++ {
p[i] = r.data[int(off)+i]
} }
return return
} }
@@ -65,32 +64,35 @@ func (r *fromReader) ReadAt(p []byte, off int64) (n int, err error) {
func (r *fromReader) Seek(off int64, whence int) (n int64, err error) { func (r *fromReader) Seek(off int64, whence int) (n int64, err error) {
switch whence { switch whence {
case io.SeekEnd: case io.SeekEnd:
return 0, errors.New("cannot SeekEnd RawReader") r.off = len(r.data) - int(off)
if r.off < 0 {
r.off = 0
err = io.EOF
}
case io.SeekCurrent: case io.SeekCurrent:
r.off += int(off) r.off += int(off)
case io.SeekStart: case io.SeekStart:
r.off = int(off) r.off = int(off)
} }
if r.off > len(r.cache) { if r.off > len(r.data) {
err = r.increaseCache(len(r.cache) - r.off) r.off = len(r.data)
if err != nil { return int64(r.off), io.EOF
r.off = len(r.cache)
}
} }
return int64(r.off), err return int64(r.off), err
} }
func (r *fromReader) Read(p []byte) (n int, err error) { func (r *fromReader) Read(p []byte) (n int, err error) {
if len(p)+r.off > len(r.cache) { n = len(p)
err = r.increaseCache((len(p) + r.off) - len(r.cache)) if r.off+len(p) > len(r.data) {
if err != nil { n = len(r.data) - r.off
return err = io.EOF
} }
if n < 0 {
n = 0
} }
for i := 0; i < len(p); i++ { for i := 0; i < n; i++ {
p[i] = r.cache[r.off+i] p[i] = r.data[r.off+i]
} }
r.off += len(p)
return return
} }
+14 -5
View File
@@ -46,12 +46,18 @@ func NewSquashfsReader(r io.ReaderAt) (*Reader, error) {
} }
//NewSquashfsReaderFromReader returns a new squashfs.Reader from an io.Reader. //NewSquashfsReaderFromReader returns a new squashfs.Reader from an io.Reader.
//If possible, try to use a io.ReaderAt. //If the io.Reader implements io.Seeker, the seek functions are used.
//With an io.Reader, much of the data has to be cached to memory. //It is NOT recommended to use a pure io.Reader as due to how squashfs
//archives are formatted, the ENTIRETY of the io.Reader's data is loaded into
//memory first before it can be used.
func NewSquashfsReaderFromReader(r io.Reader) (*Reader, error) { func NewSquashfsReaderFromReader(r io.Reader) (*Reader, error) {
var rdr Reader var rdr Reader
rdr.r = rawreader.ConvertReader(r) var err error
err := rdr.Init() rdr.r, err = rawreader.ConvertReader(r)
if err != nil {
return nil, err
}
err = rdr.Init()
if err != nil { if err != nil {
return nil, err return nil, err
} }
@@ -145,7 +151,10 @@ func (r *Reader) Init() error {
} }
unread := r.super.IDCount unread := r.super.IDCount
blockOffsets := make([]uint64, int(math.Ceil(float64(r.super.IDCount)/2048))) blockOffsets := make([]uint64, int(math.Ceil(float64(r.super.IDCount)/2048)))
r.r.Seek(int64(r.super.IDTableStart), io.SeekStart) _, err = r.r.Seek(int64(r.super.IDTableStart), io.SeekStart)
if err != nil {
return err
}
for i := range blockOffsets { for i := range blockOffsets {
err = binary.Read(r.r, binary.LittleEndian, &blockOffsets[i]) err = binary.Read(r.r, binary.LittleEndian, &blockOffsets[i])
if err != nil { if err != nil {
+2 -13
View File
@@ -40,25 +40,13 @@ func TestSquashfs(t *testing.T) {
if err != nil { if err != nil {
t.Fatal(err) t.Fatal(err)
} }
os.RemoveAll(wd + "/testing/" + squashfsName + ".d")
op := DefaultOptions() op := DefaultOptions()
op.Verbose = true op.Verbose = true
err = rdr.ExtractWithOptions(wd+"/testing/"+squashfsName+".d", op) err = rdr.ExtractWithOptions(wd+"/testing/"+squashfsName+".d", op)
if err != nil { if err != nil {
t.Fatal(err) t.Fatal(err)
} }
fmt.Println("stuff", rdr.super.CompressionType)
// fil := rdr.GetFileAtPath("*.desktop")
// if fil == nil {
// t.Fatal("Can't find desktop fil")
// }
// errs := fil.ExtractTo(wd + "/testing")
// if len(errs) > 0 {
// t.Fatal(errs)
// }
// errs = rdr.ExtractTo(wd + "/testing/" + squashfsName + ".d")
// if len(errs) > 0 {
// t.Fatal(errs)
// }
t.Fatal("No Problems") t.Fatal("No Problems")
} }
@@ -72,6 +60,7 @@ func TestSquashfsFromReader(t *testing.T) {
if err != nil { if err != nil {
t.Fatal(err) t.Fatal(err)
} }
os.RemoveAll("testing/" + squashfsName + ".d")
op := DefaultOptions() op := DefaultOptions()
op.Verbose = true op.Verbose = true
err = rdr.ExtractWithOptions("testing/"+squashfsName+".d", op) err = rdr.ExtractWithOptions("testing/"+squashfsName+".d", op)