Chore: refine code

Loyalsoldier · Jul 8, 2024 · ef95676 · ef95676
1 parent ad3e4c3
commit ef95676
Show file tree

Hide file tree

Showing 3 changed files with 8 additions and 50 deletions.
diff --git a/lib/common.go b/lib/common.go
@@ -6,7 +6,7 @@ import (
 	"net/http"
 )
 
-func getRemoteURLContent(url string) ([]byte, error) {
+func GetRemoteURLContent(url string) ([]byte, error) {
 	resp, err := http.Get(url)
 	if err != nil {
 		return nil, err

diff --git a/lib/instance.go b/lib/instance.go
@@ -26,7 +26,7 @@ func (i *Instance) Init(configFile string) error {
 	var err error
 	configFile = strings.TrimSpace(configFile)
 	if strings.HasPrefix(strings.ToLower(configFile), "http://") || strings.HasPrefix(strings.ToLower(configFile), "https://") {
-		content, err = getRemoteURLContent(configFile)
+		content, err = GetRemoteURLContent(configFile)
 	} else {
 		content, err = os.ReadFile(configFile)
 	}

diff --git a/plugin/maxmind/mmdb_in.go b/plugin/maxmind/mmdb_in.go
@@ -3,8 +3,6 @@ package maxmind
 import (
 	"encoding/json"
 	"fmt"
-	"io"
-	"net/http"
 	"os"
 	"path/filepath"
 	"strings"
@@ -20,8 +18,6 @@ const (
 
 var (
 	defaultMMDBFile = filepath.Join("./", "geolite2", "GeoLite2-Country.mmdb")
-	tempMMDBPath    = filepath.Join("./", "tmp")
-	tempMMDBFile    = filepath.Join(tempMMDBPath, "input.mmdb")
 )
 
 func init() {
@@ -82,26 +78,20 @@ func (g *maxmindMMDBIn) GetDescription() string {
 }
 
 func (g *maxmindMMDBIn) Input(container lib.Container) (lib.Container, error) {
-	var fd io.ReadCloser
+	var content []byte
 	var err error
 	switch {
 	case strings.HasPrefix(strings.ToLower(g.URI), "http://"), strings.HasPrefix(strings.ToLower(g.URI), "https://"):
-		fd, err = g.downloadFile(g.URI)
+		content, err = lib.GetRemoteURLContent(g.URI)
 	default:
-		fd, err = os.Open(g.URI)
+		content, err = os.ReadFile(g.URI)
 	}
-
-	if err != nil {
-		return nil, err
-	}
-
-	err = g.moveFile(fd)
 	if err != nil {
 		return nil, err
 	}
 
 	entries := make(map[string]*lib.Entry)
-	err = g.generateEntries(entries)
+	err = g.generateEntries(content, entries)
 	if err != nil {
 		return nil, err
 	}
@@ -149,40 +139,8 @@ func (g *maxmindMMDBIn) Input(container lib.Container) (lib.Container, error) {
 	return container, nil
 }
 
-func (g *maxmindMMDBIn) downloadFile(url string) (io.ReadCloser, error) {
-	resp, err := http.Get(url)
-	if err != nil {
-		return nil, err
-	}
-
-	if resp.StatusCode != 200 {
-		return nil, fmt.Errorf("failed to get remote file %s, http status code %d", url, resp.StatusCode)
-	}
-
-	return resp.Body, nil
-}
-
-func (g *maxmindMMDBIn) moveFile(src io.ReadCloser) error {
-	defer src.Close()
-
-	err := os.MkdirAll(tempMMDBPath, 0755)
-	if err != nil {
-		return err
-	}
-
-	out, err := os.Create(tempMMDBFile)
-	if err != nil {
-		return err
-	}
-	defer out.Close()
-
-	_, err = io.Copy(out, src)
-
-	return err
-}
-
-func (g *maxmindMMDBIn) generateEntries(entries map[string]*lib.Entry) error {
-	db, err := maxminddb.Open(tempMMDBFile)
+func (g *maxmindMMDBIn) generateEntries(content []byte, entries map[string]*lib.Entry) error {
+	db, err := maxminddb.FromBytes(content)
 	if err != nil {
 		return err
 	}