Caleb Gardner
105 lines
2.0 KiB
Go
105 lines
2.0 KiB
Go
package data
|
|
|
|
import (
|
|
"bytes"
|
|
"io"
|
|
|
|
"github.com/CalebQ42/squashfs/internal/decompress"
|
|
)
|
|
|
|
type Reader struct {
|
|
master io.Reader
|
|
cur io.Reader
|
|
fragRdr io.Reader
|
|
d decompress.Decompressor
|
|
comRdr io.Reader
|
|
blockSizes []uint32
|
|
blockSize uint32
|
|
resetable bool
|
|
fileSize uint64
|
|
}
|
|
|
|
func NewReader(r io.Reader, d decompress.Decompressor, blockSizes []uint32, blockSize uint32, fileSize uint64) *Reader {
|
|
return &Reader{
|
|
d: d,
|
|
master: r,
|
|
blockSizes: blockSizes,
|
|
blockSize: blockSize,
|
|
resetable: true,
|
|
fileSize: fileSize,
|
|
}
|
|
}
|
|
|
|
func (r *Reader) AddFragment(rdr io.Reader) {
|
|
r.fragRdr = rdr
|
|
r.blockSizes = append(r.blockSizes, 0)
|
|
}
|
|
|
|
func realSize(siz uint32) uint32 {
|
|
return siz &^ (1 << 24)
|
|
}
|
|
|
|
func (r *Reader) advance() (err error) {
|
|
if clr, ok := r.cur.(io.Closer); ok {
|
|
clr.Close()
|
|
}
|
|
if len(r.blockSizes) == 0 {
|
|
return io.EOF
|
|
}
|
|
if len(r.blockSizes) == 1 && r.fragRdr != nil {
|
|
r.cur = r.fragRdr
|
|
} else {
|
|
size := realSize(r.blockSizes[0])
|
|
if size == 0 {
|
|
outSize := r.blockSize
|
|
if r.fileSize < uint64(r.blockSize) {
|
|
outSize = uint32(r.fileSize)
|
|
}
|
|
r.cur = bytes.NewReader(make([]byte, outSize))
|
|
} else {
|
|
r.cur = io.LimitReader(r.master, int64(size))
|
|
if size == r.blockSizes[0] {
|
|
if rs, ok := r.d.(decompress.Resetable); ok {
|
|
if r.comRdr == nil {
|
|
r.cur, err = r.d.Reader(r.cur)
|
|
if err != nil {
|
|
return
|
|
}
|
|
} else {
|
|
err = rs.Reset(r.comRdr, r.cur)
|
|
r.cur = r.comRdr
|
|
}
|
|
} else {
|
|
r.cur, err = r.d.Reader(r.cur)
|
|
}
|
|
}
|
|
}
|
|
}
|
|
r.blockSizes = r.blockSizes[1:]
|
|
return
|
|
}
|
|
|
|
func (r *Reader) Read(p []byte) (n int, err error) {
|
|
if r.cur == nil {
|
|
err = r.advance()
|
|
if err != nil {
|
|
return
|
|
}
|
|
}
|
|
n, err = r.cur.Read(p)
|
|
if err == io.EOF {
|
|
err = r.advance()
|
|
if err != nil {
|
|
return
|
|
}
|
|
var tmpN int
|
|
tmp := make([]byte, len(p)-n)
|
|
tmpN, err = r.Read(tmp)
|
|
for i := range tmp {
|
|
p[n+i] = tmp[i]
|
|
}
|
|
n += tmpN
|
|
}
|
|
return
|
|
}
|