diff --git a/file.go b/file.go index 5c3cc3d..f595d12 100644 --- a/file.go +++ b/file.go @@ -13,13 +13,14 @@ import ( "github.com/CalebQ42/squashfs/internal/routinemanager" squashfslow "github.com/CalebQ42/squashfs/low" + "github.com/CalebQ42/squashfs/low/data" "github.com/CalebQ42/squashfs/low/inode" ) // File represents a file inside a squashfs archive. type File struct { - // full data.FullReader - // rdr data.Reader + full data.FullReader + rdr data.Reader rdrInit bool parent FS r *Reader diff --git a/low/README.md b/low/README.md index 33eb8eb..bb6cde4 100644 --- a/low/README.md +++ b/low/README.md @@ -1,3 +1,5 @@ # Lower-Level Squashfs This library is a lower level version of the main [squashfs](https://github.com/CalebQ42/squashfs) library that doesn't try to be easy to use and exposes a lot of information that is not necesary for must use cases. + +I will try to keep the API stable, but it is not guarenteed. diff --git a/low/data/fullreader.go b/low/data/fullreader.go index 59c9e9a..73c6544 100644 --- a/low/data/fullreader.go +++ b/low/data/fullreader.go @@ -3,18 +3,20 @@ package data import ( "errors" "io" + "runtime" "github.com/CalebQ42/squashfs/internal/decompress" ) type FullReader struct { - fileSize uint64 - blockSize uint32 - rdr io.ReaderAt - decomp decompress.Decompressor - sizes []uint32 - blockOffsets []uint64 - fragDat []byte + fileSize uint64 + blockSize uint32 + goroutineLimit uint16 + rdr io.ReaderAt + decomp decompress.Decompressor + sizes []uint32 + blockOffsets []uint64 + fragDat []byte } func NewFullReader(rdr io.ReaderAt, decomp decompress.Decompressor, size uint64, start uint64, blockSizes []uint32) FullReader { @@ -33,6 +35,13 @@ func NewFullReader(rdr io.ReaderAt, decomp decompress.Decompressor, size uint64, return out } +func (f *FullReader) Close() error { + f.fragDat = nil + f.sizes = nil + f.blockOffsets = nil + return nil +} + func (f *FullReader) AddFragData(blockStart uint64, offset uint32, blockSize uint32) error { realSize := blockSize &^ (1 << 24) dat := make([]byte, realSize) @@ -50,17 +59,30 @@ func (f *FullReader) AddFragData(blockStart uint64, offset uint32, blockSize uin return nil } +func (f *FullReader) SetGoroutineLimit(limit uint16) { + f.goroutineLimit = limit +} + +// The number of blocks, including the fragment block if present +func (f FullReader) BlockNum() uint32 { + out := len(f.sizes) + if f.fragDat != nil { + out++ + } + return uint32(out) +} + // Returns the data block at the given index -func (f FullReader) Block(i int) ([]byte, error) { - if i == len(f.sizes) && f.fragDat != nil { +func (f FullReader) Block(i uint32) ([]byte, error) { + if i == uint32(len(f.sizes)) && f.fragDat != nil { return f.fragDat, nil } - if i >= len(f.sizes) { + if i >= uint32(len(f.sizes)) { return nil, errors.New("invalid block index") } realSize := f.sizes[i] &^ (1 << 24) if realSize == 0 { - if i == len(f.sizes)-1 && f.fragDat == nil { + if i == uint32(len(f.sizes)-1) && f.fragDat == nil { return make([]byte, f.fileSize%uint64(f.blockSize)), nil } return make([]byte, f.blockSize), nil @@ -76,6 +98,92 @@ func (f FullReader) Block(i int) ([]byte, error) { return dat, nil } -func (f FullReader) WriteTo(w io.Writer) (int64, error) { +type blockResults struct { + idx uint32 + block []byte + err error +} +func (f FullReader) WriteTo(w io.Writer) (wrote int64, err error) { + routineLimit := f.goroutineLimit + if routineLimit == 0 { + routineLimit = uint16(runtime.NumCPU() / 2) + } + dispatchChan := make(chan struct{}, routineLimit) + for range int(routineLimit) { + dispatchChan <- struct{}{} + } + resChan := make(chan blockResults, routineLimit) + var results map[uint32]blockResults + if _, is := w.(io.WriterAt); !is { + results = make(map[uint32]blockResults) + } + for i := range f.BlockNum() { + go func(idx uint32) { + _, closed := <-dispatchChan + if closed { + resChan <- blockResults{} + return + } + block, err := f.Block(idx) + resChan <- blockResults{ + idx: idx, + block: block, + err: err, + } + dispatchChan <- struct{}{} + }(i) + } + out := int64(0) + errOut := make([]error, 0) + for i := uint32(0); i < f.BlockNum(); { + res := <-resChan + if res.err != nil { + close(dispatchChan) + errOut = append(errOut, res.err) + } + if len(errOut) > 0 { + continue + } + if wa, is := w.(io.WriterAt); is { + _, err := wa.WriteAt(res.block, int64(res.idx)*int64(f.blockSize)) + if err != nil { + errOut = append(errOut, err) + } else { + out = max(out, int64(res.idx)*int64(f.blockSize)+int64(len(res.block))) + } + continue + } + var err error + if res.idx == i { + _, err = w.Write(res.block) + if err != nil { + errOut = append(errOut, err) + } else { + out = max(out, int64(res.idx)*int64(f.blockSize)+int64(len(res.block))) + } + i++ + } else { + results[res.idx] = res + } + var has bool + for { + res, has = results[i] + if has { + _, err = w.Write(res.block) + if err != nil { + errOut = append(errOut, err) + } else { + out = max(out, int64(res.idx)*int64(f.blockSize)+int64(len(res.block))) + } + i++ + } else { + break + } + } + } + if len(errOut) > 0 { + return out, errors.Join(errOut...) + } + return out, nil } diff --git a/low/data/reader.go b/low/data/reader.go index 0ad59c2..dc7d9ac 100644 --- a/low/data/reader.go +++ b/low/data/reader.go @@ -1 +1,60 @@ package data + +import "io" + +type Reader struct { + f *FullReader + curBlock []byte + nextIdx uint32 + curOffset uint32 +} + +func NewReader(f *FullReader) (Reader, error) { + dat, err := f.Block(0) + if err != nil { + return Reader{}, err + } + return Reader{ + f: f, + curBlock: dat, + nextIdx: 1, + curOffset: 0, + }, nil +} + +func (d *Reader) Close() error { + d.curBlock = nil + return nil +} + +func (d *Reader) advanceBlock() error { + if d.nextIdx >= d.f.BlockNum() { + d.curBlock = nil + return io.EOF + } + var err error + d.curBlock, err = d.f.Block(d.nextIdx) + if err != nil { + return err + } + d.nextIdx++ + d.curOffset = 0 + return nil +} + +func (d *Reader) Read(buf []byte) (int, error) { + totRed := 0 + toRead := 0 + var err error + for totRed < len(buf) { + if int(d.curOffset) >= len(d.curBlock) { + err = d.advanceBlock() + if err != nil { + return totRed, err + } + } + toRead = min(len(d.curBlock)-int(d.curOffset), len(buf)-totRed) + copy(buf[totRed:], d.curBlock[d.curOffset:d.curOffset+uint32(toRead)]) + } + return totRed, nil +} diff --git a/low/file_base.go b/low/file_base.go index f4e57e3..635fac1 100644 --- a/low/file_base.go +++ b/low/file_base.go @@ -2,7 +2,6 @@ package squashfslow import ( "errors" - "io" "github.com/CalebQ42/squashfs/internal/metadata" "github.com/CalebQ42/squashfs/internal/toreader" @@ -84,6 +83,8 @@ func (b FileBase) IsRegular() bool { return b.Inode.Type == inode.Fil || b.Inode.Type == inode.EFil } +// Returns a regular file's readers. They are linked, so the data.Reader calls to the data.FullReader. +// Aka: closing the FullReader breaks the Reader func (b FileBase) GetRegFileReaders(r Reader) (data.Reader, data.FullReader, error) { if !b.IsRegular() { return data.Reader{}, data.FullReader{}, errors.New("not a regular file") @@ -91,41 +92,28 @@ func (b FileBase) GetRegFileReaders(r Reader) (data.Reader, data.FullReader, err var blockStart uint64 var fragIndex uint32 var fragOffset uint32 - var fragSize uint64 var sizes []uint32 + var fileSize uint64 if b.Inode.Type == inode.Fil { blockStart = uint64(b.Inode.Data.(inode.File).BlockStart) fragIndex = b.Inode.Data.(inode.File).FragInd fragOffset = b.Inode.Data.(inode.File).FragOffset sizes = b.Inode.Data.(inode.File).BlockSizes - fragSize = uint64(b.Inode.Data.(inode.File).Size % r.Superblock.BlockSize) + fileSize = uint64(b.Inode.Data.(inode.File).Size) } else { blockStart = b.Inode.Data.(inode.EFile).BlockStart fragIndex = b.Inode.Data.(inode.EFile).FragInd fragOffset = b.Inode.Data.(inode.EFile).FragOffset sizes = b.Inode.Data.(inode.EFile).BlockSizes - fragSize = b.Inode.Data.(inode.EFile).Size % uint64(r.Superblock.BlockSize) + fileSize = b.Inode.Data.(inode.EFile).Size } - frag := func() (io.Reader, error) { - ent, err := r.fragEntry(fragIndex) - if err != nil { - return nil, err - } - frag := data.NewReader(toreader.NewReader(r.r, int64(ent.Start)), r.d, []uint32{ent.Size}, uint64(r.Superblock.BlockSize), r.Superblock.BlockSize) - frag.Read(make([]byte, fragOffset)) - return io.LimitReader(&frag, int64(fragSize)), nil - } - outRdr := data.NewReader(toreader.NewReader(r.r, int64(blockStart)), r.d, sizes, fragSize, r.Superblock.BlockSize) + outFull := data.NewFullReader(r.r, r.d, fileSize, blockStart, sizes) if fragIndex != 0xffffffff { - f, err := frag() - if err != nil { - return data.Reader{}, data.FullReader{}, err - } - outRdr.AddFrag(f) + outFull.AddFragData(r.fragTable[fragIndex].Start, fragOffset, r.fragTable[fragIndex].Size) } - outFull := data.NewFullReader(r.r, int64(blockStart), r.d, sizes, fragSize, r.Superblock.BlockSize) - if fragIndex != 0xffffffff { - outFull.AddFrag(frag) + outRdr, err := data.NewReader(&outFull) + if err != nil { + return data.Reader{}, data.FullReader{}, err } return outRdr, outFull, nil } @@ -137,67 +125,24 @@ func (b FileBase) GetFullReader(r *Reader) (data.FullReader, error) { var blockStart uint64 var fragIndex uint32 var fragOffset uint32 - var fragSize uint64 var sizes []uint32 + var fileSize uint64 if b.Inode.Type == inode.Fil { blockStart = uint64(b.Inode.Data.(inode.File).BlockStart) fragIndex = b.Inode.Data.(inode.File).FragInd fragOffset = b.Inode.Data.(inode.File).FragOffset sizes = b.Inode.Data.(inode.File).BlockSizes - fragSize = uint64(b.Inode.Data.(inode.File).Size % r.Superblock.BlockSize) + fileSize = uint64(b.Inode.Data.(inode.File).Size) } else { blockStart = b.Inode.Data.(inode.EFile).BlockStart fragIndex = b.Inode.Data.(inode.EFile).FragInd fragOffset = b.Inode.Data.(inode.EFile).FragOffset sizes = b.Inode.Data.(inode.EFile).BlockSizes - fragSize = b.Inode.Data.(inode.EFile).Size % uint64(r.Superblock.BlockSize) + fileSize = b.Inode.Data.(inode.EFile).Size } - outFull := data.NewFullReader(r.r, int64(blockStart), r.d, sizes, fragSize, r.Superblock.BlockSize) + outFull := data.NewFullReader(r.r, r.d, fileSize, blockStart, sizes) if fragIndex != 0xffffffff { - outFull.AddFrag(func() (io.Reader, error) { - ent, err := r.fragEntry(fragIndex) - if err != nil { - return nil, err - } - frag := data.NewReader(toreader.NewReader(r.r, int64(ent.Start)), r.d, []uint32{ent.Size}, uint64(r.Superblock.BlockSize), r.Superblock.BlockSize) - frag.Read(make([]byte, fragOffset)) - return io.LimitReader(&frag, int64(fragSize)), nil - }) + outFull.AddFragData(r.fragTable[fragIndex].Start, fragOffset, r.fragTable[fragIndex].Size) } return outFull, nil } - -func (b FileBase) GetReader(r *Reader) (data.Reader, error) { - if !b.IsRegular() { - return data.Reader{}, errors.New("not a regular file") - } - var blockStart uint64 - var fragIndex uint32 - var fragOffset uint32 - var fragSize uint64 - var sizes []uint32 - if b.Inode.Type == inode.Fil { - blockStart = uint64(b.Inode.Data.(inode.File).BlockStart) - fragIndex = b.Inode.Data.(inode.File).FragInd - fragOffset = b.Inode.Data.(inode.File).FragOffset - sizes = b.Inode.Data.(inode.File).BlockSizes - fragSize = uint64(b.Inode.Data.(inode.File).Size % r.Superblock.BlockSize) - } else { - blockStart = b.Inode.Data.(inode.EFile).BlockStart - fragIndex = b.Inode.Data.(inode.EFile).FragInd - fragOffset = b.Inode.Data.(inode.EFile).FragOffset - sizes = b.Inode.Data.(inode.EFile).BlockSizes - fragSize = b.Inode.Data.(inode.EFile).Size % uint64(r.Superblock.BlockSize) - } - outRdr := data.NewReader(toreader.NewReader(r.r, int64(blockStart)), r.d, sizes, fragSize, r.Superblock.BlockSize) - if fragIndex != 0xffffffff { - ent, err := r.fragEntry(fragIndex) - if err != nil { - return data.Reader{}, err - } - frag := data.NewReader(toreader.NewReader(r.r, int64(ent.Start)), r.d, []uint32{ent.Size}, uint64(r.Superblock.BlockSize), r.Superblock.BlockSize) - frag.Read(make([]byte, fragOffset)) - outRdr.AddFrag(io.LimitReader(&frag, int64(fragSize))) - } - return outRdr, nil -}