From 4f76569bfe55747485eee564a77c64c4f70afee6 Mon Sep 17 00:00:00 2001
From: Wu Tingfeng <wutingfeng@outlook.com>
Date: Fri, 23 Dec 2022 01:52:16 +0800
Subject: [PATCH] Use temporary folder

---
 Makefile    |  2 +-
 fasttld.go  | 41 ++++++++++++++++++++++-------------------
 psl.go      | 19 ++++---------------
 psl_test.go |  6 ------
 4 files changed, 27 insertions(+), 41 deletions(-)

diff --git a/Makefile b/Makefile
index 490ce4c..bdd9fd9 100644
--- a/Makefile
+++ b/Makefile
@@ -20,7 +20,7 @@ mem_report:
 	go tool pprof mem.prof
 
 build_cli:
-	go build -o ./dist/fasttld -ldflags "-X 'github.com/elliotwutingfeng/go-fasttld/cmd/fasttld.version=v0.4.2'" ./cmd/main.go
+	go build -o ./dist/fasttld -ldflags "-X 'github.com/elliotwutingfeng/go-fasttld/cmd/fasttld.version=v0.4.3'" ./cmd/main.go
 
 demo:
 	go run ./examples/demo.go
diff --git a/fasttld.go b/fasttld.go
index 6b4c759..62f64e6 100644
--- a/fasttld.go
+++ b/fasttld.go
@@ -10,10 +10,10 @@ import (
 	"errors"
 	"log"
 	"net/url"
-	"os"
 	"strconv"
 	"strings"
 
+	"github.com/spf13/afero"
 	"github.com/tidwall/hashmap"
 	"golang.org/x/net/idna"
 )
@@ -431,31 +431,34 @@ func New(n SuffixListParams) (*FastTLD, error) {
 		tldTrie, err := trieConstruct(n.IncludePrivateSuffix, "")
 		return &FastTLD{cacheFilePath: "", tldTrie: tldTrie, includePrivateSuffix: n.IncludePrivateSuffix}, err
 	}
-	// If cacheFilePath is unreachable, use default Public Suffix List file.
-	if isValid, _ := checkCacheFile(n.CacheFilePath); !isValid {
-		defaultCacheFolderPath, defaultCacheFilePath, err := getDefaultCachePaths()
-		if err != nil || os.MkdirAll(defaultCacheFolderPath, 0644) != nil {
-			// default Public Suffix List file cannot be opened
+	extractor := &FastTLD{cacheFilePath: n.CacheFilePath, tldTrie: &trie{}, includePrivateSuffix: n.IncludePrivateSuffix}
+	// If cacheFilePath is unreachable, use temporary folder
+	if isValid, _ := checkCacheFile(extractor.cacheFilePath); !isValid {
+		filesystem := new(afero.OsFs)
+		defaultCacheFolderPath := afero.GetTempDir(filesystem, "")
+		defaultCacheFilePath := defaultCacheFolderPath + defaultPSLFileName
+		defaultCacheFolder, err := filesystem.Open(defaultCacheFolderPath)
+		if err != nil {
+			// temporary folder not accessible, fallback to inline Public Suffix list
 			return inlinePSL(err, n)
 		}
-		n.CacheFilePath = defaultCacheFilePath
-		if isValid, lastModifiedHours := checkCacheFile(n.CacheFilePath); !isValid || lastModifiedHours > pslMaxAgeHours {
-			if file, err := os.OpenFile(n.CacheFilePath, os.O_CREATE|os.O_WRONLY, 0644); err == nil {
-				if err := update(file, publicSuffixListSources); err != nil {
-					log.Println(err)
-				}
-				defer file.Close()
+		defer defaultCacheFolder.Close()
+		extractor.cacheFilePath = defaultCacheFilePath
+		isValid, lastModifiedHours := checkCacheFile(extractor.cacheFilePath)
+		if !isValid || lastModifiedHours > pslMaxAgeHours {
+			// update Public Suffix list cache if it is outdated
+			if updateErr := extractor.Update(); updateErr != nil {
+				// update failed, fallback to inline Public Suffix list
+				return inlinePSL(err, n)
 			}
-		}
-		if isValid, _ := checkCacheFile(n.CacheFilePath); !isValid {
-			return inlinePSL(err, n)
+			return extractor, err
 		}
 	}
 
-	tldTrie, err := trieConstruct(n.IncludePrivateSuffix, n.CacheFilePath)
+	tldTrie, err := trieConstruct(n.IncludePrivateSuffix, extractor.cacheFilePath)
 	if err != nil {
 		return inlinePSL(err, n)
 	}
-
-	return &FastTLD{cacheFilePath: n.CacheFilePath, tldTrie: tldTrie, includePrivateSuffix: n.IncludePrivateSuffix}, err
+	extractor.tldTrie = tldTrie
+	return extractor, err
 }
diff --git a/psl.go b/psl.go
index 51f7edd..b4504be 100644
--- a/psl.go
+++ b/psl.go
@@ -160,17 +160,6 @@ func update(file afero.File,
 	return nil
 }
 
-func getDefaultCachePaths() (string, string, error) {
-	currentFilePath, ok := getCurrentFilePath()
-	if !ok {
-		return "", "", errors.New("Cannot get path to current module file")
-	}
-	defaultCacheFolderPath := currentFilePath + string(os.PathSeparator) + defaultPSLFolder
-	defaultCacheFilePath := defaultCacheFolderPath + string(os.PathSeparator) + defaultPSLFileName
-
-	return defaultCacheFolderPath, defaultCacheFilePath, nil
-}
-
 func checkCacheFile(cacheFilePath string) (bool, float64) {
 	cacheFilePath, pathValidErr := filepath.Abs(strings.TrimSpace(cacheFilePath))
 	stat, fileinfoErr := os.Stat(cacheFilePath)
@@ -193,10 +182,10 @@ func checkCacheFile(cacheFilePath string) (bool, float64) {
 // Update updates the default Public Suffix list file and updates its suffix trie using the updated file.
 // If cache file path is not the same as the default cache file path, this will be a no-op.
 func (f *FastTLD) Update() error {
-	defaultCacheFolderPath, defaultCacheFilePath, err := getDefaultCachePaths()
-	if err := os.MkdirAll(defaultCacheFolderPath, 0644); err != nil {
-		return err
-	}
+	filesystem := new(afero.OsFs)
+	defaultCacheFolderPath := afero.GetTempDir(filesystem, "")
+	defaultCacheFilePath := defaultCacheFolderPath + defaultPSLFileName
+
 	if f.cacheFilePath != defaultCacheFilePath {
 		return errors.New("No-op. Only default Public Suffix list file can be updated")
 	}
diff --git a/psl_test.go b/psl_test.go
index 485d4c7..c1a7fc6 100644
--- a/psl_test.go
+++ b/psl_test.go
@@ -167,9 +167,3 @@ func TestFileLastModifiedHours(t *testing.T) {
 	}
 	defer file.Close()
 }
-
-func TestGetDefaultCachePaths(t *testing.T) {
-	if _, _, err := getDefaultCachePaths(); err != nil {
-		t.Errorf("Expected no error, got %v", err)
-	}
-}