Add context cancellation and metadata progress

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
 }

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)
+				}
 			}
 		}
 

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,
 		)
 	})
 }

extractor/extractor.go

@@ -1,29 +1,30 @@
 package extractor
 
 import (
+	"context"
 	"io"
 
 	"github.com/xypwn/filediver/exec"
 	"github.com/xypwn/filediver/stingray"
 )
 
 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
 		}

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
 	}

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
 			}