From 16248701db5a72140da048f04194e8a2f2236f3d Mon Sep 17 00:00:00 2001 From: xypwn Date: Tue, 30 Apr 2024 16:13:37 +0200 Subject: [PATCH] Add context cancellation and metadata progress --- app/app.go | 69 +++++++++++++++++++++++--------- cmd/filediver-cli/main.go | 37 ++++++++++++++--- extractor/bik/extractor.go | 6 +-- extractor/extractor.go | 9 +++-- extractor/texture/extractor.go | 4 +- extractor/unit/extractor.go | 8 ++-- extractor/wwise/extractor.go | 38 ++++++++---------- stingray/dir.go | 28 +++++++++---- stingray/unit/texture/texture.go | 9 +++-- util/io.go | 28 +++++++++++++ 10 files changed, 164 insertions(+), 72 deletions(-) diff --git a/app/app.go b/app/app.go index 106f2c1..be7bdbc 100644 --- a/app/app.go +++ b/app/app.go @@ -2,6 +2,7 @@ package app import ( "bufio" + "context" "encoding/binary" "errors" "fmt" @@ -79,8 +80,8 @@ var ConfigFormat = ConfigTemplate{ Fallback: "raw", } -func parseWwiseDep(f *stingray.File) (string, error) { - r, err := f.Open(stingray.DataMain) +func parseWwiseDep(ctx context.Context, f *stingray.File) (string, error) { + r, err := f.Open(ctx, stingray.DataMain) if err != nil { return "", err } @@ -139,8 +140,8 @@ type App struct { } // Open game dir and read metadata. -func New(gameDir string, hashes []string) (*App, error) { - dataDir, err := stingray.OpenDataDir(filepath.Join(gameDir, "data")) +func OpenGameDir(ctx context.Context, gameDir string, hashes []string, onProgress func(curr, total int)) (*App, error) { + dataDir, err := stingray.OpenDataDir(ctx, filepath.Join(gameDir, "data"), onProgress) if err != nil { return nil, err } @@ -152,7 +153,7 @@ func New(gameDir string, hashes []string) (*App, error) { // wwise_dep files let us know the string of many of the wwise_banks for id, file := range dataDir.Files { if id.Type == stingray.Sum64([]byte("wwise_dep")) { - h, err := parseWwiseDep(file) + h, err := parseWwiseDep(ctx, file) if err != nil { return nil, fmt.Errorf("wwise_dep: %w", err) } @@ -274,15 +275,18 @@ func (a *App) MatchingFiles(includeGlob, excludeGlob string, cfgTemplate ConfigT } type extractContext struct { + ctx context.Context app *App file *stingray.File runner *exec.Runner config map[string]string outPath string + files []string } -func newExtractContext(app *App, file *stingray.File, runner *exec.Runner, config map[string]string, outPath string) *extractContext { +func newExtractContext(ctx context.Context, app *App, file *stingray.File, runner *exec.Runner, config map[string]string, outPath string) *extractContext { return &extractContext{ + ctx: ctx, app: app, file: file, runner: runner, @@ -291,6 +295,7 @@ func newExtractContext(app *App, file *stingray.File, runner *exec.Runner, confi } } +func (c *extractContext) OutPath() (string, error) { return c.outPath, nil } func (c *extractContext) File() *stingray.File { return c.file } func (c *extractContext) Runner() *exec.Runner { return c.runner } func (c *extractContext) Config() map[string]string { return c.config } @@ -299,18 +304,27 @@ func (c *extractContext) GetResource(name, typ stingray.Hash) (file *stingray.Fi return } func (c *extractContext) CreateFile(suffix string) (io.WriteCloser, error) { - return os.Create(c.outPath + suffix) -} -func (c *extractContext) CreateFileDir(dirSuffix, filename string) (io.WriteCloser, error) { - dir := c.outPath + dirSuffix - if err := os.MkdirAll(dir, os.ModePerm); err != nil { + path, err := c.AllocateFile(suffix) + if err != nil { return nil, err } - return os.Create(filepath.Join(dir, filename)) + return os.Create(path) +} +func (c *extractContext) AllocateFile(suffix string) (string, error) { + path := c.outPath + suffix + if err := os.MkdirAll(filepath.Dir(path), os.ModePerm); err != nil { + return "", err + } + c.files = append(c.files, path) + return path, nil +} +func (c *extractContext) Ctx() context.Context { return c.ctx } +func (c *extractContext) Files() []string { + return c.files } -func (c *extractContext) OutPath() (string, error) { return c.outPath, nil } -func (a *App) ExtractFile(id stingray.FileID, outDir string, extrCfg map[string]map[string]string, runner *exec.Runner) error { +// Returns path to extracted file/directory. +func (a *App) ExtractFile(ctx context.Context, id stingray.FileID, outDir string, extrCfg map[string]map[string]string, runner *exec.Runner) ([]string, error) { name, ok := a.Hashes[id.Name] if !ok { name = id.Name.String() @@ -322,7 +336,7 @@ func (a *App) ExtractFile(id stingray.FileID, outDir string, extrCfg map[string] file, ok := a.DataDir.Files[id] if !ok { - return fmt.Errorf("extract %v.%v: file does not found", name, typ) + return nil, fmt.Errorf("extract %v.%v: file does not exist", name, typ) } cfg := extrCfg[typ] @@ -376,17 +390,32 @@ func (a *App) ExtractFile(id stingray.FileID, outDir string, extrCfg map[string] outPath := filepath.Join(outDir, name) if err := os.MkdirAll(filepath.Dir(outPath), os.ModePerm); err != nil { - return err + return nil, err } - if err := extr(newExtractContext( + extrCtx := newExtractContext( + ctx, a, file, runner, cfg, outPath, - )); err != nil { - return fmt.Errorf("extract %v.%v: %w", name, typ, err) + ) + if err := extr(extrCtx); err != nil { + { + var err error + var errPath string + for _, path := range extrCtx.Files() { + if e := os.Remove(path); e != nil && !errors.Is(e, os.ErrNotExist) && err == nil { + err = e + errPath = path + } + } + if err != nil { + return nil, fmt.Errorf("cleanup %v: %w", errPath, err) + } + } + return nil, fmt.Errorf("extract %v.%v: %w", name, typ, err) } - return nil + return extrCtx.Files(), nil } diff --git a/cmd/filediver-cli/main.go b/cmd/filediver-cli/main.go index 78bdb64..cd9cb56 100644 --- a/cmd/filediver-cli/main.go +++ b/cmd/filediver-cli/main.go @@ -1,11 +1,15 @@ package main import ( + "context" _ "embed" + "errors" "fmt" "os" + "os/signal" "runtime/pprof" "sort" + "syscall" //"github.com/davecgh/go-spew/spew" @@ -111,11 +115,28 @@ extractor config: prt.Infof("Output directory: \"%v\"", *outDir) } - prt.Infof("Reading metadata...") - a, err := app.New(*gameDir, knownHashes) + ctx, cancel := context.WithCancel(context.Background()) + + sigs := make(chan os.Signal, 1) + signal.Notify(sigs, syscall.SIGINT, syscall.SIGTERM) + go func() { + <-sigs + cancel() + }() + + a, err := app.OpenGameDir(ctx, *gameDir, knownHashes, func(curr, total int) { + prt.Statusf("Reading metadata %.0f%%", float64(curr)/float64(total)*100) + }) if err != nil { - prt.Fatalf("%v", err) + if errors.Is(err, context.Canceled) { + prt.NoStatus() + prt.Warnf("Metadata read canceled, exiting") + return + } else { + prt.Errorf("%v", err) + } } + prt.NoStatus() files, err := a.MatchingFiles(*extrInclGlob, *extrExclGlob, app.ConfigFormat, extrCfg) if err != nil { @@ -182,10 +203,16 @@ extractor config: truncName = "..." + truncName[len(truncName)-37:] } prt.Statusf("File %v/%v: %v", i+1, len(files), truncName) - if err := a.ExtractFile(id, *outDir, extrCfg, runner); err == nil { + if _, err := a.ExtractFile(ctx, id, *outDir, extrCfg, runner); err == nil { numExtrFiles++ } else { - prt.Errorf("%v", err) + if errors.Is(err, context.Canceled) { + prt.NoStatus() + prt.Warnf("Extraction canceled, exiting cleanly") + return + } else { + prt.Errorf("%v", err) + } } } diff --git a/extractor/bik/extractor.go b/extractor/bik/extractor.go index 6d9f99a..b4c36b7 100644 --- a/extractor/bik/extractor.go +++ b/extractor/bik/extractor.go @@ -16,7 +16,7 @@ func extract(ctx extractor.Context, save func(ctx extractor.Context, r io.Reader dataTypes = append(dataTypes, stingray.DataGPU) } - r, err := ctx.File().OpenMulti(dataTypes...) + r, err := ctx.File().OpenMulti(ctx.Ctx(), dataTypes...) if err != nil { return err } @@ -52,7 +52,7 @@ func ConvertToMP4(ctx extractor.Context) error { } return extract(ctx, func(ctx extractor.Context, r io.Reader) error { - outPath, err := ctx.OutPath() + outPath, err := ctx.AllocateFile(".mp4") if err != nil { return err } @@ -62,7 +62,7 @@ func ConvertToMP4(ctx extractor.Context) error { r, "-f", "bink", "-i", "pipe:", - outPath+".mp4", + outPath, ) }) } diff --git a/extractor/extractor.go b/extractor/extractor.go index 8380be0..b154dfd 100644 --- a/extractor/extractor.go +++ b/extractor/extractor.go @@ -1,6 +1,7 @@ package extractor import ( + "context" "io" "github.com/xypwn/filediver/exec" @@ -8,22 +9,22 @@ import ( ) type Context interface { + Ctx() context.Context File() *stingray.File Runner() *exec.Runner Config() map[string]string GetResource(name, typ stingray.Hash) (file *stingray.File, exists bool) // Call WriteCloser.Close() when done. CreateFile(suffix string) (io.WriteCloser, error) - // Call WriteCloser.Close() when done. - CreateFileDir(dirSuffix, filename string) (io.WriteCloser, error) - OutPath() (string, error) + // Returns path to file. + AllocateFile(suffix string) (string, error) } type ExtractFunc func(ctx Context) error func ExtractFuncRaw(suffix string, types ...stingray.DataType) ExtractFunc { return func(ctx Context) error { - r, err := ctx.File().OpenMulti(types...) + r, err := ctx.File().OpenMulti(ctx.Ctx(), types...) if err != nil { return err } diff --git a/extractor/texture/extractor.go b/extractor/texture/extractor.go index 9a99bf9..e579d5d 100644 --- a/extractor/texture/extractor.go +++ b/extractor/texture/extractor.go @@ -14,7 +14,7 @@ func ExtractDDS(ctx extractor.Context) error { if !ctx.File().Exists(stingray.DataMain) { return errors.New("no main data") } - r, err := ctx.File().OpenMulti(stingray.DataMain, stingray.DataStream, stingray.DataGPU) + r, err := ctx.File().OpenMulti(ctx.Ctx(), stingray.DataMain, stingray.DataStream, stingray.DataGPU) if err != nil { return err } @@ -34,7 +34,7 @@ func ExtractDDS(ctx extractor.Context) error { } func ConvertToPNG(ctx extractor.Context) error { - tex, err := texture.Decode(ctx.File(), false) + tex, err := texture.Decode(ctx.Ctx(), ctx.File(), false) if err != nil { return err } diff --git a/extractor/unit/extractor.go b/extractor/unit/extractor.go index 63fe995..8e293c0 100644 --- a/extractor/unit/extractor.go +++ b/extractor/unit/extractor.go @@ -114,7 +114,7 @@ func writeTexture(ctx extractor.Context, doc *gltf.Document, id stingray.Hash, p return 0, fmt.Errorf("texture resource %v doesn't exist", id) } - tex, err := texture.Decode(file, false) + tex, err := texture.Decode(ctx.Ctx(), file, false) if err != nil { return 0, err } @@ -149,14 +149,14 @@ func writeTexture(ctx extractor.Context, doc *gltf.Document, id stingray.Hash, p } func Convert(ctx extractor.Context) error { - fMain, err := ctx.File().Open(stingray.DataMain) + fMain, err := ctx.File().Open(ctx.Ctx(), stingray.DataMain) if err != nil { return err } defer fMain.Close() var fGPU io.ReadSeekCloser if ctx.File().Exists(stingray.DataGPU) { - fGPU, err = ctx.File().Open(stingray.DataGPU) + fGPU, err = ctx.File().Open(ctx.Ctx(), stingray.DataGPU) if err != nil { return err } @@ -185,7 +185,7 @@ func Convert(ctx extractor.Context) error { return fmt.Errorf("referenced material resource %v doesn't exist", resID) } mat, err := func() (*material.Material, error) { - f, err := matRes.Open(stingray.DataMain) + f, err := matRes.Open(ctx.Ctx(), stingray.DataMain) if err != nil { return nil, err } diff --git a/extractor/wwise/extractor.go b/extractor/wwise/extractor.go index 93d514a..3abfeec 100644 --- a/extractor/wwise/extractor.go +++ b/extractor/wwise/extractor.go @@ -7,7 +7,6 @@ import ( "io" "math" "os" - "path/filepath" "github.com/go-audio/audio" "github.com/go-audio/wav" @@ -84,7 +83,7 @@ func pcmFloat32ToIntS16(dst []int, src []float32) { } } -func convertWemStream(ctx extractor.Context, outPath string, in io.ReadSeeker, format format) error { +func convertWemStream(ctx extractor.Context, outName string, in io.ReadSeeker, format format) error { if !ctx.Runner().Has("ffmpeg") { format = formatWav } @@ -95,7 +94,11 @@ func convertWemStream(ctx extractor.Context, outPath string, in io.ReadSeeker, f } switch format { case formatWav: - out, err := os.Create(outPath + ".wav") + outPath, err := ctx.AllocateFile(outName + ".wav") + if err != nil { + return err + } + out, err := os.Create(outPath) if err != nil { return err } @@ -139,6 +142,10 @@ func convertWemStream(ctx extractor.Context, outPath string, in io.ReadSeeker, f case formatAac: fmtExt = ".aac" } + outPath, err := ctx.AllocateFile(outName + fmtExt) + if err != nil { + return err + } if err := ctx.Runner().Run( "ffmpeg", nil, @@ -148,7 +155,7 @@ func convertWemStream(ctx extractor.Context, outPath string, in io.ReadSeeker, f "-ac", fmt.Sprint(dec.Channels()), "-channel_layout", fmt.Sprintf("0x%x", uint32(dec.ChannelLayout())), "-i", "pipe:", - outPath+fmtExt, + outPath, ); err != nil { return err } @@ -180,16 +187,12 @@ func ConvertWem(ctx extractor.Context) error { if err != nil { return err } - outPath, err := ctx.OutPath() - if err != nil { - return err - } - r, err := ctx.File().Open(stingray.DataStream) + r, err := ctx.File().Open(ctx.Ctx(), stingray.DataStream) if err != nil { return err } defer r.Close() - if err := convertWemStream(ctx, outPath, r, format); err != nil { + if err := convertWemStream(ctx, "", r, format); err != nil { return err } return nil @@ -226,7 +229,7 @@ func extractBnk(in io.ReadSeeker) (io.ReadSeeker, error) { } func ExtractBnk(ctx extractor.Context) error { - f, err := ctx.File().Open(stingray.DataMain) + f, err := ctx.File().Open(ctx.Ctx(), stingray.DataMain) if err != nil { return err } @@ -253,7 +256,7 @@ func ConvertBnk(ctx extractor.Context) error { return err } - in, err := ctx.File().Open(stingray.DataMain) + in, err := ctx.File().Open(ctx.Ctx(), stingray.DataMain) if err != nil { return err } @@ -273,22 +276,13 @@ func ConvertBnk(ctx extractor.Context) error { return err } - outPath, err := ctx.OutPath() - if err != nil { - return err - } - dirPath := outPath + ".bnk" - if err := os.MkdirAll(dirPath, os.ModePerm); err != nil { - return err - } for i := 0; i < bnk.NumFiles(); i++ { wemR, err := bnk.OpenFile(i) if err != nil { return err } if err := func() error { - outPath := filepath.Join(dirPath, fmt.Sprintf("%03d", i)) - if err := convertWemStream(ctx, outPath, wemR, format); err != nil { + if err := convertWemStream(ctx, fmt.Sprintf(".bnk/%03d", i), wemR, format); err != nil { return err } return nil diff --git a/stingray/dir.go b/stingray/dir.go index cadff88..a2d2c1f 100644 --- a/stingray/dir.go +++ b/stingray/dir.go @@ -2,10 +2,13 @@ package stingray import ( "bytes" + "context" "fmt" "io" "os" "path/filepath" + + "github.com/xypwn/filediver/util" ) const errPfx = "stingray: " @@ -45,7 +48,7 @@ func (r *preallocReader) Close() error { } // Call Close() on returned reader when done. -func (f *File) Open(typ DataType) (io.ReadSeekCloser, error) { +func (f *File) Open(ctx context.Context, typ DataType) (io.ReadSeekCloser, error) { fileR, err := f.triad.OpenFile(f.index, typ) if err != nil { return nil, err @@ -54,7 +57,7 @@ func (f *File) Open(typ DataType) (io.ReadSeekCloser, error) { if err != nil { return nil, err } - return r, nil + return util.NewContextReadSeekCloser(ctx, r), nil } type multiReadCloser struct { @@ -77,14 +80,14 @@ func (r *multiReadCloser) Close() error { // Skips any specified types that don't exist. // If you need seeking functionality, use Open(). // Call Close() on returned reader when done. -func (f *File) OpenMulti(types ...DataType) (io.ReadCloser, error) { +func (f *File) OpenMulti(ctx context.Context, types ...DataType) (io.ReadCloser, error) { var rdcs []io.ReadCloser var rds []io.Reader for _, dataType := range types { if !f.Exists(dataType) { continue } - r, err := f.Open(dataType) + r, err := f.Open(ctx, dataType) if err != nil { for _, rdc := range rdcs { rdc.Close() @@ -102,7 +105,8 @@ func (f *File) OpenMulti(types ...DataType) (io.ReadCloser, error) { // For testing purposes, takes a HUGE amount of time to execute. func (a *File) contentEqual(b *File, dt DataType) (bool, error) { - fa, err := a.Open(dt) + ctx := context.Background() + fa, err := a.Open(ctx, dt) if err != nil { return false, err } @@ -111,7 +115,7 @@ func (a *File) contentEqual(b *File, dt DataType) (bool, error) { if err != nil { return false, err } - fb, err := b.Open(dt) + fb, err := b.Open(ctx, dt) if err != nil { return false, err } @@ -127,7 +131,9 @@ type DataDir struct { Files map[FileID]*File } -func OpenDataDir(dirPath string) (*DataDir, error) { +// Opens the "data" game directory, reading all file metadata. Ctx allows for granular cancellation (before each triad open). +// onProgress is optional. +func OpenDataDir(ctx context.Context, dirPath string, onProgress func(curr, total int)) (*DataDir, error) { const errPfx = errPfx + "OpenDataDir: " ents, err := os.ReadDir(dirPath) @@ -139,7 +145,13 @@ func OpenDataDir(dirPath string) (*DataDir, error) { Files: make(map[FileID]*File), } - for _, ent := range ents { + for i, ent := range ents { + if err := ctx.Err(); err != nil { + return nil, err + } + if onProgress != nil { + onProgress(i, len(ents)) + } if !ent.Type().IsRegular() { continue } diff --git a/stingray/unit/texture/texture.go b/stingray/unit/texture/texture.go index 25d0622..dcdf9d2 100644 --- a/stingray/unit/texture/texture.go +++ b/stingray/unit/texture/texture.go @@ -1,6 +1,7 @@ package texture import ( + "context" "encoding/binary" "errors" "fmt" @@ -44,11 +45,11 @@ func DecodeInfo(r io.Reader) (*Info, error) { }, nil } -func decode(f *stingray.File, readMipMaps bool) (*dds.DDS, error) { +func decode(ctx context.Context, f *stingray.File, readMipMaps bool) (*dds.DDS, error) { if !f.Exists(stingray.DataMain) { return nil, errors.New("no main data") } - r, err := f.OpenMulti(stingray.DataMain, stingray.DataStream, stingray.DataGPU) + r, err := f.OpenMulti(ctx, stingray.DataMain, stingray.DataStream, stingray.DataGPU) if err != nil { return nil, err } @@ -65,8 +66,8 @@ func decode(f *stingray.File, readMipMaps bool) (*dds.DDS, error) { } // Decode DDS texture with Stingray wrapper. -func Decode(f *stingray.File, readMipMaps bool) (*dds.DDS, error) { - tex, err := decode(f, readMipMaps) +func Decode(ctx context.Context, f *stingray.File, readMipMaps bool) (*dds.DDS, error) { + tex, err := decode(ctx, f, readMipMaps) if err != nil { return nil, fmt.Errorf("stingray texture: %w", err) } diff --git a/util/io.go b/util/io.go index aeda125..122be0a 100644 --- a/util/io.go +++ b/util/io.go @@ -1,6 +1,7 @@ package util import ( + "context" "io" ) @@ -53,3 +54,30 @@ func (r *SectionReadSeeker) Seek(offset int64, whence int) (int64, error) { pos, err := r.r.Seek(r.off, io.SeekStart) return pos - r.base, err } + +// Cancellable ReadSeekCloser type +type contextReadSeekCloser struct { + io.ReadSeekCloser + ctx context.Context +} + +func NewContextReadSeekCloser(ctx context.Context, r io.ReadSeekCloser) io.ReadSeekCloser { + return &contextReadSeekCloser{ + ReadSeekCloser: r, + ctx: ctx, + } +} + +func (r *contextReadSeekCloser) Read(p []byte) (n int, err error) { + if err := r.ctx.Err(); err != nil { + return 0, err + } + return r.ReadSeekCloser.Read(p) +} + +func (r *contextReadSeekCloser) Seek(offset int64, whence int) (int64, error) { + if err := r.ctx.Err(); err != nil { + return 0, err + } + return r.ReadSeekCloser.Seek(offset, whence) +}