Cleanup tmp directory in the vulnerability ingestion service (#1033)

#1032

lunatrace/bsl/ingest-worker/cmd/ingestworker/main.go

@@ -11,7 +11,6 @@
 package main
 
 import (
-	"context"
 	"github.com/lunasec-io/lunasec/lunatrace/bsl/ingest-worker/cmd/ingestworker/cwe"
 	packageCommand "github.com/lunasec-io/lunasec/lunatrace/bsl/ingest-worker/cmd/ingestworker/package"
 	"github.com/lunasec-io/lunasec/lunatrace/bsl/ingest-worker/cmd/ingestworker/vulnerability"
@@ -23,7 +22,6 @@ import (
 	"github.com/lunasec-io/lunasec/lunatrace/bsl/ingest-worker/pkg/metadata/replicator"
 	"github.com/lunasec-io/lunasec/lunatrace/bsl/ingest-worker/pkg/scanner/licensecheck"
 	"github.com/lunasec-io/lunasec/lunatrace/bsl/ingest-worker/pkg/scanner/packagejson"
-	"github.com/lunasec-io/lunasec/lunatrace/cli/pkg/util"
 	"github.com/rs/zerolog"
 	"github.com/rs/zerolog/log"
 	"net/http"
@@ -49,12 +47,6 @@ func main() {
 		dbfx.Module,
 		registry.NPMModule,
 
-		fx.Invoke(func() {
-			util.RunOnProcessExit(func() {
-				util.RemoveCleanupDirs()
-			})
-		}),
-
 		fx.Provide(
 			cwe2.NewCWEIngester,
 		),
@@ -86,12 +78,5 @@ func main() {
 		fx.Provide(
 			packageCommand.NewCommand,
 		),
-
-		fx.Invoke(func(lc fx.Lifecycle) {
-			lc.Append(fx.Hook{OnStop: func(_ context.Context) error {
-				util.RemoveCleanupDirs()
-				return nil
-			}})
-		}),
 	)
 }

lunatrace/bsl/ingest-worker/pkg/vulnerability/ingester.go

@@ -171,7 +171,9 @@ func (f FileAdvisoryIngester) upsertVulnerabilities(
 }
 
 func (f FileAdvisoryIngester) IngestVulnerabilitiesFromSource(advisoryLocation, source, sourceRelativePath string) error {
-	advisoryLocation, err := ensureAdvisoriesExistFromSource(source, advisoryLocation)
+	advisoryLocation, cleanup, err := ensureAdvisoriesExistFromSource(source, advisoryLocation)
+
+	defer cleanup()
 	if err != nil {
 		return err
 	}

lunatrace/bsl/ingest-worker/pkg/vulnerability/source.go

@@ -1,20 +1,18 @@
 // Copyright by LunaSec (owned by Refinery Labs, Inc)
 //
-// Licensed under the Business Source License v1.1 
+// Licensed under the Business Source License v1.1
 // (the "License"); you may not use this file except in compliance with the
 // License. You may obtain a copy of the License at
 //
 // https://github.com/lunasec-io/lunasec/blob/master/licenses/BSL-LunaTrace.txt
 //
 // See the License for the specific language governing permissions and
 // limitations under the License.
-//
 package vulnerability
 
 import (
 	"fmt"
 	"github.com/go-git/go-git/v5"
-	"github.com/lunasec-io/lunasec/lunatrace/cli/pkg/util"
 	"github.com/rs/zerolog/log"
 	"io/ioutil"
 	"os"
@@ -38,19 +36,29 @@ func pullVulnerabilitiesFromSource(source, dst string) error {
 	return fmt.Errorf("cannot pull vulnerabilities for source: %s", source)
 }
 
-func ensureAdvisoriesExistFromSource(source, advisoryDir string) (string, error) {
+func ensureAdvisoriesExistFromSource(source, advisoryDir string) (string, func(), error) {
+	cleanup := func() {}
+
 	if advisoryDir != "" {
-		return advisoryDir, nil
+		return advisoryDir, cleanup, nil
 	}
 
 	advisoryDir, err := ioutil.TempDir("", source+"-advisories")
 	if err != nil {
 		log.Error().
 			Err(err).
 			Msg("unable to create temporary directory for advisories")
-		return "", err
+		return "", cleanup, err
+	}
+	cleanup = func() {
+		err = os.RemoveAll(advisoryDir)
+		if err != nil {
+			log.Error().
+				Err(err).
+				Str("advisory dir", advisoryDir).
+				Msg("failed to remove temporary advisory location")
+		}
 	}
-	util.EnsureDirIsCleanedUp(advisoryDir)
 
 	log.Info().
 		Str("source", source).
@@ -61,12 +69,12 @@ func ensureAdvisoriesExistFromSource(source, advisoryDir string) (string, error)
 		log.Error().
 			Err(err).
 			Msg("unable to pull vulnerabilities from source")
-		return "", err
+		return "", cleanup, err
 	}
 
 	log.Info().
 		Str("source", source).
 		Msg("collected advisories from source")
 
-	return advisoryDir, nil
+	return advisoryDir, cleanup, nil
 }