Bugfix: Retry removing tempfile in scope destructor (Velocidex#3642)

On windows, we can not remove a file that is still opened. This PR
allows the removal routine to be re-scheduled in the destructor list for
the root scope in order to retry removing it after the file is closed by
another destructor.

accessors/zip/zip.go

@@ -39,6 +39,7 @@ import (
 	"www.velocidex.com/golang/velociraptor/accessors"
 	"www.velocidex.com/golang/velociraptor/constants"
 	"www.velocidex.com/golang/velociraptor/json"
+	"www.velocidex.com/golang/velociraptor/services/debug"
 	"www.velocidex.com/golang/velociraptor/third_party/zip"
 	"www.velocidex.com/golang/velociraptor/utils"
 	"www.velocidex.com/golang/vfilter"
@@ -99,6 +100,7 @@ func (self *Tracker) Debug() string {
 func (self *Tracker) Reset() {
 	self.mu.Lock()
 	defer self.mu.Unlock()
+
 	self.refs = make(map[string]int)
 }
 
@@ -114,11 +116,25 @@ func (self *Tracker) Dec(filename string) {
 		} else {
 			self.refs[filename] = prev
 		}
+
 	} else {
 		panic(filename)
 	}
 }
 
+func (self *Tracker) ProfileWriter(ctx context.Context,
+	scope vfilter.Scope, output_chan chan vfilter.Row) {
+
+	self.mu.Lock()
+	defer self.mu.Unlock()
+
+	for filename, ref := range self.refs {
+		output_chan <- ordereddict.NewDict().
+			Set("Filename", filename).
+			Set("ReferenceCount", ref)
+	}
+}
+
 // Wrapper around zip.File with reference counting. Note that each
 // instance is holding a reference to the zip.Reader it came from. We
 // also increase references to ZipFileCache to manage its references
@@ -683,4 +699,10 @@ Example:
 	}, `Open a zip file as if it was a directory. Although zip files are case sensitive, this accessor behaves case insensitive`)
 
 	json.RegisterCustomEncoder(&ZipFileInfo{}, accessors.MarshalGlobFileInfo)
+
+	debug.RegisterProfileWriter(debug.ProfileWriterInfo{
+		Name:          "ZipTracker",
+		Description:   "Reference counting for open Zip files",
+		ProfileWriter: tracker.ProfileWriter,
+	})
 }

go.mod

@@ -99,7 +99,7 @@ require (
 	www.velocidex.com/golang/go-prefetch v0.0.0-20220801101854-338dbe61982a
 	www.velocidex.com/golang/oleparse v0.0.0-20230217092320-383a0121aafe
 	www.velocidex.com/golang/regparser v0.0.0-20240404115756-2169ac0e3c09
-	www.velocidex.com/golang/vfilter v0.0.0-20240618023104-cd2ef63ee978
+	www.velocidex.com/golang/vfilter v0.0.0-20240725055846-2a1740af9bab
 )
 
 require (

go.sum

@@ -1371,7 +1371,7 @@ www.velocidex.com/golang/oleparse v0.0.0-20230217092320-383a0121aafe h1:o9jQWSwK
 www.velocidex.com/golang/oleparse v0.0.0-20230217092320-383a0121aafe/go.mod h1:R7IisRzDO7q5LVRJsCQf1xA50LrIavsPWzAjVE4THyY=
 www.velocidex.com/golang/regparser v0.0.0-20240404115756-2169ac0e3c09 h1:G1RWYBXP2lSzxKcrAU1YhiUlBetZ7hGIzIiWuuazvfo=
 www.velocidex.com/golang/regparser v0.0.0-20240404115756-2169ac0e3c09/go.mod h1:pxSECT5mWM3goJ4sxB4HCJNKnKqiAlpyT8XnvBwkLGU=
-www.velocidex.com/golang/vfilter v0.0.0-20240618023104-cd2ef63ee978 h1:4OsjvVgF4L4B9lzLjHFaxBLk4c6V7IK4uohb4z1qsGI=
-www.velocidex.com/golang/vfilter v0.0.0-20240618023104-cd2ef63ee978/go.mod h1:P50KPQr2LpWVAu7ilGH8CBLBASGtOJ2971yA9YhR8rY=
+www.velocidex.com/golang/vfilter v0.0.0-20240725055846-2a1740af9bab h1:vbxheOyXsbkW9lmG3hz4TTYFAUDTqO5lEtR1ZtI4iAs=
+www.velocidex.com/golang/vfilter v0.0.0-20240725055846-2a1740af9bab/go.mod h1:P50KPQr2LpWVAu7ilGH8CBLBASGtOJ2971yA9YhR8rY=
 www.velocidex.com/golang/vtypes v0.0.0-20240123105603-069d4a7f435c h1:rL/It+Ig+mvIhmy9vl5gg5b6CX2J12x0v2SXIT2RoWE=
 www.velocidex.com/golang/vtypes v0.0.0-20240123105603-069d4a7f435c/go.mod h1:tjaJNlBWbvH4cEMrEu678CFR2hrtcdyPINIpRxrOh4U=

vql/filesystem/tempfile.go

@@ -22,7 +22,6 @@ import (
 	"io/ioutil"
 	"os"
 	"runtime"
-	"time"
 
 	"github.com/Velocidex/ordereddict"
 	"www.velocidex.com/golang/velociraptor/acls"
@@ -97,18 +96,20 @@ func (self *TempfileFunction) Call(ctx context.Context,
 
 	if arg.RemoveLast {
 		scope.Log("Adding global destructor for %v", tmpfile.Name())
-		err := vql_subsystem.GetRootScope(scope).
-			AddDestructor(func() { RemoveFile(scope, tmpfile.Name()) })
+		root_scope := vql_subsystem.GetRootScope(scope)
+		err := root_scope.AddDestructor(func() {
+			RemoveFile(0, tmpfile.Name(), root_scope)
+		})
 		if err != nil {
-			RemoveFile(scope, tmpfile.Name())
+			RemoveFile(0, tmpfile.Name(), scope)
 			scope.Log("tempfile: %v", err)
 		}
 	} else {
 		err := scope.AddDestructor(func() {
-			RemoveFile(scope, tmpfile.Name())
+			RemoveFile(0, tmpfile.Name(), scope)
 		})
 		if err != nil {
-			RemoveFile(scope, tmpfile.Name())
+			RemoveFile(0, tmpfile.Name(), scope)
 			scope.Log("tempfile: %v", err)
 		}
 	}
@@ -156,37 +157,25 @@ func (self *TempdirFunction) Call(ctx context.Context,
 		return false
 	}
 
-	// Make sure the file is removed when the query is done.
-	removal := func() {
-		scope.Log("tempdir: removing tempdir %v", dir)
-
-		// On windows especially we can not remove files that
-		// are opened by something else, so we keep trying for
-		// a while.
-		for i := 0; i < 100; i++ {
-			err := os.RemoveAll(dir)
-			if err == nil {
-				break
-			}
-			time.Sleep(time.Second)
-		}
-	}
-
 	if arg.RemoveLast {
 		scope.Log("Adding global destructor for %v", dir)
-		err := vql_subsystem.GetRootScope(scope).AddDestructor(removal)
+		root_scope := vql_subsystem.GetRootScope(scope)
+		err := root_scope.AddDestructor(func() {
+			RemoveDirectory(0, dir, root_scope)
+		})
 		if err != nil {
-			removal()
+			RemoveDirectory(0, dir, scope)
 			scope.Log("tempdir: %v", err)
 		}
 
 	} else {
-		err := scope.AddDestructor(removal)
+		err := scope.AddDestructor(func() {
+			RemoveDirectory(0, dir, scope)
+		})
 		if err != nil {
-			removal()
+			RemoveDirectory(0, dir, scope)
 			scope.Log("tempdir: %v", err)
 		}
-
 	}
 	return dir
 }
@@ -201,18 +190,73 @@ func (self TempdirFunction) Info(scope vfilter.Scope,
 	}
 }
 
-func RemoveFile(scope vfilter.Scope, filename string) {
-	scope.Log("tempfile: removing tempfile %v", filename)
+// Make sure the file is removed when the query is done.
+func RemoveDirectory(retry int, tmpdir string, scope vfilter.Scope) {
+	if retry >= 10 {
+		scope.Log("RemoveDirectory: Retry count exceeded - giving up")
+		return
+	}
+
+	if retry > 0 {
+		scope.Log("RemoveDirectory: removing tempdir %v (Try %v)",
+			tmpdir, retry)
+	} else {
+		scope.Log("RemoveDirectory: removing tempdir %v", tmpdir)
+	}
 
 	// On windows especially we can not remove files that
 	// are opened by something else, so we keep trying for
 	// a while.
-	for i := 0; i < 100; i++ {
-		err := os.Remove(filename)
-		if err == nil {
-			break
+	err := os.RemoveAll(tmpdir)
+	if err != nil {
+		scope.Log("RemoveDirectory: Failed to remove %v: %v, reschedule", tmpdir, err)
+
+		// Add another detructor to try again a bit later.
+		scope.AddDestructor(func() {
+			RemoveFile(retry+1, tmpdir, scope)
+		})
+	} else {
+		if retry > 0 {
+			scope.Log("RemoveDirectory: removed tempdir %v (Try %v)",
+				tmpdir, retry)
+		} else {
+			scope.Log("RemoveDirectory: removed tempdir %v", tmpdir)
+		}
+	}
+}
+
+// Make sure the file is removed when the query is done.
+func RemoveFile(retry int, tmpfile string, scope vfilter.Scope) {
+	if retry >= 10 {
+		scope.Log("tempfile: Retry count exceeded - giving up")
+		return
+	}
+
+	if retry > 0 {
+		scope.Log("tempfile: removing tempfile %v (Try %v)",
+			tmpfile, retry)
+	} else {
+		scope.Log("tempfile: removing tempfile %v", tmpfile)
+	}
+
+	// On windows especially we can not remove files that
+	// are opened by something else, so we keep trying for
+	// a while.
+	err := os.Remove(tmpfile)
+	if err != nil {
+		scope.Log("tempfile: Failed to remove %v: %v, reschedule", tmpfile, err)
+
+		// Add another detructor to try again a bit later.
+		scope.AddDestructor(func() {
+			RemoveFile(retry+1, tmpfile, scope)
+		})
+	} else {
+		if retry > 0 {
+			scope.Log("tempfile: removed tempfile %v (Try %v)",
+				tmpfile, retry)
+		} else {
+			scope.Log("tempfile: removed tempfile %v", tmpfile)
 		}
-		time.Sleep(time.Second)
 	}
 }
 

vql/materializer/materialize.go

@@ -43,8 +43,10 @@ func NewTempFileMatrializer(
 	if err != nil {
 		return nil, err
 	}
-	vql_subsystem.GetRootScope(scope).AddDestructor(
-		func() { filesystem.RemoveFile(scope, tmpfile.Name()) })
+	root_scope := vql_subsystem.GetRootScope(scope)
+	root_scope.AddDestructor(func() {
+		filesystem.RemoveFile(0, tmpfile.Name(), root_scope)
+	})
 
 	result := &TempFileMatrializer{
 		filename: tmpfile.Name(),

vql/networking/http_client.go

@@ -40,6 +40,7 @@ import (
 	"www.velocidex.com/golang/velociraptor/utils"
 	"www.velocidex.com/golang/velociraptor/vql"
 	vql_subsystem "www.velocidex.com/golang/velociraptor/vql"
+	"www.velocidex.com/golang/velociraptor/vql/filesystem"
 	"www.velocidex.com/golang/velociraptor/vql/functions"
 	vfilter "www.velocidex.com/golang/vfilter"
 	"www.velocidex.com/golang/vfilter/arg_parser"
@@ -510,21 +511,23 @@ func (self *_HttpPlugin) Call(
 				return
 			}
 
-			remove := func() {
-				remove_tmpfile(tmpfile.Name(), scope)
-			}
 			if arg.RemoveLast {
 				scope.Log("Adding global destructor for %v", tmpfile.Name())
-				err := vql_subsystem.GetRootScope(scope).AddDestructor(remove)
+				root_scope := vql_subsystem.GetRootScope(scope)
+				err := root_scope.AddDestructor(func() {
+					filesystem.RemoveFile(0, tmpfile.Name(), root_scope)
+				})
 				if err != nil {
-					remove()
+					filesystem.RemoveFile(0, tmpfile.Name(), scope)
 					scope.Log("http_client: %v", err)
 					return
 				}
 			} else {
-				err := scope.AddDestructor(remove)
+				err := scope.AddDestructor(func() {
+					filesystem.RemoveFile(0, tmpfile.Name(), scope)
+				})
 				if err != nil {
-					remove()
+					filesystem.RemoveFile(0, tmpfile.Name(), scope)
 					scope.Log("http_client: %v", err)
 					return
 				}
@@ -597,23 +600,6 @@ func (self _HttpPlugin) Info(scope vfilter.Scope, type_map *vfilter.TypeMap) *vf
 	}
 }
 
-// Make sure the file is removed when the query is done.
-func remove_tmpfile(tmpfile string, scope vfilter.Scope) {
-	scope.Log("tempfile: removing tempfile %v", tmpfile)
-
-	// On windows especially we can not remove files that
-	// are opened by something else, so we keep trying for
-	// a while.
-	for i := 0; i < 100; i++ {
-		err := os.Remove(tmpfile)
-		if err == nil {
-			break
-		}
-		scope.Log("tempfile: Error %v - will retry", err)
-		time.Sleep(time.Second)
-	}
-}
-
 func SetProxy(handler func(*http.Request) (*url.URL, error)) {
 	mu.Lock()
 	defer mu.Unlock()

vql/parsers/pe_dump.go

@@ -109,10 +109,10 @@ func (self _PEDumpFunction) Call(
 			return false
 		}
 		defer tmpfile.Close()
-		_ = vql_subsystem.GetRootScope(scope).
-			AddDestructor(func() {
-				filesystem.RemoveFile(scope, tmpfile.Name())
-			})
+		root_scope := vql_subsystem.GetRootScope(scope)
+		_ = root_scope.AddDestructor(func() {
+			filesystem.RemoveFile(0, tmpfile.Name(), root_scope)
+		})
 
 		writer = tmpfile
 

vql/vql.go

@@ -195,7 +195,8 @@ func MakeNewScope() vfilter.Scope {
 	return result
 }
 
-// Used in tests to flush the global scope - needed **after** .
+// Used in tests to flush the global scope - needed **after**
+// overriding plugins with OverridePlugin, OverrideFunction etc.
 func ResetGlobalScopeCache() {
 	mu.Lock()
 	defer mu.Unlock()

vql/windows/winpmem.go

@@ -93,10 +93,11 @@ func (self WinpmemFunction) Call(
 		}
 
 		// Driver will only be uninstalled when then root scope is destroyed.
-		vql_subsystem.GetRootScope(scope).AddDestructor(func() {
+		root_scope := vql_subsystem.GetRootScope(scope)
+		root_scope.AddDestructor(func() {
 			err := winpmem.UninstallDriver(tmpfile.Name(), arg.ServiceName, logger)
 			if err == nil {
-				filesystem.RemoveFile(scope, tmpfile.Name())
+				filesystem.RemoveFile(0, tmpfile.Name(), root_scope)
 			}
 		})