Skip to content

Commit

Permalink
Chore: refine code
Browse files Browse the repository at this point in the history
  • Loading branch information
Loyalsoldier committed Jul 8, 2024
1 parent ad3e4c3 commit ef95676
Show file tree
Hide file tree
Showing 3 changed files with 8 additions and 50 deletions.
2 changes: 1 addition & 1 deletion lib/common.go
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
2 changes: 1 addition & 1 deletion lib/instance.go
Original file line number Diff line number Diff line change
Expand Up @@ -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)
}
Expand Down
54 changes: 6 additions & 48 deletions plugin/maxmind/mmdb_in.go
Original file line number Diff line number Diff line change
Expand Up @@ -3,8 +3,6 @@ package maxmind
import (
"encoding/json"
"fmt"
"io"
"net/http"
"os"
"path/filepath"
"strings"
Expand All @@ -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() {
Expand Down Expand Up @@ -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
}
Expand Down Expand Up @@ -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
}
Expand Down

0 comments on commit ef95676

Please sign in to comment.