diff options
| author | Loyalsoldier <[email protected]> | 2024-07-08 14:10:42 +0800 |
|---|---|---|
| committer | Loyalsoldier <[email protected]> | 2024-07-08 14:23:41 +0800 |
| commit | 18e9f1aafac9bdb2783daa47f9a07c67e780f554 (patch) | |
| tree | 6bd2e5a610173c5bbd70f204d91687c99fa74f6a /plugin | |
| parent | 56ddd1bd048e6d42645f36cd473a211dca90b5dc (diff) | |
Chore: refine code
Diffstat (limited to 'plugin')
| -rw-r--r-- | plugin/maxmind/mmdb_in.go | 54 |
1 files changed, 6 insertions, 48 deletions
diff --git a/plugin/maxmind/mmdb_in.go b/plugin/maxmind/mmdb_in.go index 0c94a12c..3cf2f9a6 100644 --- 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 } |
