mirror of
https://github.com/ScuroNeko/mtg.git
synced 2026-08-31 13:04:02 +03:00
Refactor firehol
This commit is contained in:
@@ -22,7 +22,7 @@ func (h httpFile) Open(ctx context.Context) (io.ReadCloser, error) {
|
|||||||
response, err := h.http.Do(request)
|
response, err := h.http.Do(request)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
if response != nil {
|
if response != nil {
|
||||||
io.Copy(io.Discard, response.Body)
|
io.Copy(io.Discard, response.Body) // nolint: errcheck
|
||||||
response.Body.Close()
|
response.Body.Close()
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -36,6 +36,10 @@ func (h httpFile) Open(ctx context.Context) (io.ReadCloser, error) {
|
|||||||
return response.Body, nil
|
return response.Body, nil
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func (h httpFile) String() string {
|
||||||
|
return h.url
|
||||||
|
}
|
||||||
|
|
||||||
func NewHTTP(client *http.Client, endpoint string) (File, error) {
|
func NewHTTP(client *http.Client, endpoint string) (File, error) {
|
||||||
if client == nil {
|
if client == nil {
|
||||||
return nil, ErrBadHTTPClient
|
return nil, ErrBadHTTPClient
|
||||||
|
|||||||
@@ -22,7 +22,7 @@ type HTTPTestSuite struct {
|
|||||||
}
|
}
|
||||||
|
|
||||||
func (suite *HTTPTestSuite) makeFile(path string) (files.File, error) {
|
func (suite *HTTPTestSuite) makeFile(path string) (files.File, error) {
|
||||||
return files.NewHTTP(suite.httpClient, suite.httpServer.URL+"/"+path)
|
return files.NewHTTP(suite.httpClient, suite.httpServer.URL+"/"+path) // nolint: wrapcheck
|
||||||
}
|
}
|
||||||
|
|
||||||
func (suite *HTTPTestSuite) SetupSuite() {
|
func (suite *HTTPTestSuite) SetupSuite() {
|
||||||
|
|||||||
@@ -10,4 +10,5 @@ var ErrBadHTTPClient = errors.New("incorrect http client")
|
|||||||
|
|
||||||
type File interface {
|
type File interface {
|
||||||
Open(context.Context) (io.ReadCloser, error)
|
Open(context.Context) (io.ReadCloser, error)
|
||||||
|
String() string
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -4,18 +4,19 @@ import (
|
|||||||
"context"
|
"context"
|
||||||
"fmt"
|
"fmt"
|
||||||
"io"
|
"io"
|
||||||
"io/fs"
|
|
||||||
"os"
|
"os"
|
||||||
"path/filepath"
|
|
||||||
)
|
)
|
||||||
|
|
||||||
type localFile struct {
|
type localFile struct {
|
||||||
root fs.FS
|
path string
|
||||||
name string
|
|
||||||
}
|
}
|
||||||
|
|
||||||
func (l localFile) Open(ctx context.Context) (io.ReadCloser, error) {
|
func (l localFile) Open(ctx context.Context) (io.ReadCloser, error) {
|
||||||
return l.root.Open(l.name)
|
return os.Open(l.path) // nolint: wrapcheck
|
||||||
|
}
|
||||||
|
|
||||||
|
func (l localFile) String() string {
|
||||||
|
return l.path
|
||||||
}
|
}
|
||||||
|
|
||||||
func NewLocal(path string) (File, error) {
|
func NewLocal(path string) (File, error) {
|
||||||
@@ -24,7 +25,6 @@ func NewLocal(path string) (File, error) {
|
|||||||
}
|
}
|
||||||
|
|
||||||
return localFile{
|
return localFile{
|
||||||
root: os.DirFS(filepath.Dir(path)),
|
path: path,
|
||||||
name: filepath.Base(path),
|
|
||||||
}, nil
|
}, nil
|
||||||
}
|
}
|
||||||
|
|||||||
+55
-131
@@ -4,16 +4,13 @@ import (
|
|||||||
"bufio"
|
"bufio"
|
||||||
"context"
|
"context"
|
||||||
"fmt"
|
"fmt"
|
||||||
"io"
|
|
||||||
"net"
|
"net"
|
||||||
"net/http"
|
|
||||||
"net/url"
|
|
||||||
"os"
|
|
||||||
"regexp"
|
"regexp"
|
||||||
"strings"
|
"strings"
|
||||||
"sync"
|
"sync"
|
||||||
"time"
|
"time"
|
||||||
|
|
||||||
|
"github.com/9seconds/mtg/v2/ipblocklist/files"
|
||||||
"github.com/9seconds/mtg/v2/mtglib"
|
"github.com/9seconds/mtg/v2/mtglib"
|
||||||
"github.com/kentik/patricia"
|
"github.com/kentik/patricia"
|
||||||
"github.com/kentik/patricia/bool_tree"
|
"github.com/kentik/patricia/bool_tree"
|
||||||
@@ -41,20 +38,16 @@ var fireholRegexpComment = regexp.MustCompile(`\s*#.*?$`)
|
|||||||
// 127.0.0.1 # you can specify an IP
|
// 127.0.0.1 # you can specify an IP
|
||||||
// 10.0.0.0/8 # or cidr
|
// 10.0.0.0/8 # or cidr
|
||||||
type Firehol struct {
|
type Firehol struct {
|
||||||
ctx context.Context
|
ctx context.Context
|
||||||
ctxCancel context.CancelFunc
|
ctxCancel context.CancelFunc
|
||||||
logger mtglib.Logger
|
logger mtglib.Logger
|
||||||
|
|
||||||
updateMutex sync.RWMutex
|
updateMutex sync.RWMutex
|
||||||
|
|
||||||
remoteURLs []string
|
blocklists []files.File
|
||||||
localFiles []string
|
|
||||||
|
|
||||||
httpClient *http.Client
|
|
||||||
workerPool *ants.Pool
|
workerPool *ants.Pool
|
||||||
|
treeV4 *bool_tree.TreeV4
|
||||||
treeV4 *bool_tree.TreeV4
|
treeV6 *bool_tree.TreeV6
|
||||||
treeV6 *bool_tree.TreeV6
|
|
||||||
}
|
}
|
||||||
|
|
||||||
// Shutdown stop a background update process.
|
// Shutdown stop a background update process.
|
||||||
@@ -98,22 +91,14 @@ func (f *Firehol) Run(updateEach time.Duration) {
|
|||||||
}
|
}
|
||||||
}()
|
}()
|
||||||
|
|
||||||
if err := f.update(); err != nil {
|
f.update()
|
||||||
f.logger.WarningError("cannot update blocklist", err)
|
|
||||||
} else {
|
|
||||||
f.logger.Info("blocklist was updated")
|
|
||||||
}
|
|
||||||
|
|
||||||
for {
|
for {
|
||||||
select {
|
select {
|
||||||
case <-f.ctx.Done():
|
case <-f.ctx.Done():
|
||||||
return
|
return
|
||||||
case <-ticker.C:
|
case <-ticker.C:
|
||||||
if err := f.update(); err != nil {
|
f.update()
|
||||||
f.logger.WarningError("cannot update blocklist", err)
|
|
||||||
} else {
|
|
||||||
f.logger.Info("blocklist was updated")
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -138,121 +123,53 @@ func (f *Firehol) containsIPv6(addr net.IP) bool {
|
|||||||
return false
|
return false
|
||||||
}
|
}
|
||||||
|
|
||||||
func (f *Firehol) update() error { // nolint: funlen, cyclop
|
func (f *Firehol) update() {
|
||||||
ctx, cancel := context.WithCancel(f.ctx)
|
ctx, cancel := context.WithCancel(f.ctx)
|
||||||
defer cancel()
|
defer cancel()
|
||||||
|
|
||||||
wg := &sync.WaitGroup{}
|
wg := &sync.WaitGroup{}
|
||||||
wg.Add(len(f.remoteURLs) + len(f.localFiles))
|
wg.Add(len(f.blocklists))
|
||||||
|
|
||||||
treeMutex := &sync.Mutex{}
|
treeMutex := &sync.Mutex{}
|
||||||
v4tree := bool_tree.NewTreeV4()
|
v4tree := bool_tree.NewTreeV4()
|
||||||
v6tree := bool_tree.NewTreeV6()
|
v6tree := bool_tree.NewTreeV6()
|
||||||
|
|
||||||
errorChan := make(chan error, 1)
|
for _, v := range f.blocklists {
|
||||||
defer close(errorChan)
|
go func(file files.File) {
|
||||||
|
|
||||||
for _, v := range f.localFiles {
|
|
||||||
go func(filename string) {
|
|
||||||
defer wg.Done()
|
defer wg.Done()
|
||||||
|
|
||||||
if err := f.updateLocalFile(ctx, filename, treeMutex, v4tree, v6tree); err != nil {
|
logger := f.logger.BindStr("filename", file.String())
|
||||||
cancel()
|
|
||||||
f.logger.BindStr("filename", filename).WarningError("cannot update", err)
|
|
||||||
|
|
||||||
select {
|
fileContent, err := file.Open(ctx)
|
||||||
case errorChan <- err:
|
if err != nil {
|
||||||
default:
|
logger.WarningError("update has failed", err)
|
||||||
}
|
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
defer fileContent.Close()
|
||||||
|
|
||||||
|
if err := f.updateFromFile(treeMutex, v4tree, v6tree, bufio.NewScanner(fileContent)); err != nil {
|
||||||
|
logger.WarningError("update has failed", err)
|
||||||
}
|
}
|
||||||
}(v)
|
}(v)
|
||||||
}
|
}
|
||||||
|
|
||||||
for _, v := range f.remoteURLs {
|
|
||||||
value := v
|
|
||||||
|
|
||||||
f.workerPool.Submit(func() { // nolint: errcheck
|
|
||||||
defer wg.Done()
|
|
||||||
|
|
||||||
if err := f.updateRemoteURL(ctx, value, treeMutex, v4tree, v6tree); err != nil {
|
|
||||||
cancel()
|
|
||||||
f.logger.BindStr("url", value).WarningError("cannot update", err)
|
|
||||||
|
|
||||||
select {
|
|
||||||
case errorChan <- err:
|
|
||||||
default:
|
|
||||||
}
|
|
||||||
}
|
|
||||||
})
|
|
||||||
}
|
|
||||||
|
|
||||||
wg.Wait()
|
wg.Wait()
|
||||||
|
|
||||||
select {
|
|
||||||
case err := <-errorChan:
|
|
||||||
return fmt.Errorf("cannot update trees: %w", err)
|
|
||||||
default:
|
|
||||||
}
|
|
||||||
|
|
||||||
f.updateMutex.Lock()
|
f.updateMutex.Lock()
|
||||||
defer f.updateMutex.Unlock()
|
defer f.updateMutex.Unlock()
|
||||||
|
|
||||||
f.treeV4 = v4tree
|
f.treeV4 = v4tree
|
||||||
f.treeV6 = v6tree
|
f.treeV6 = v6tree
|
||||||
|
|
||||||
return nil
|
f.logger.Info("blocklist was updated")
|
||||||
}
|
}
|
||||||
|
|
||||||
func (f *Firehol) updateLocalFile(ctx context.Context, filename string,
|
func (f *Firehol) updateFromFile(mutex sync.Locker,
|
||||||
mutex sync.Locker,
|
|
||||||
v4tree *bool_tree.TreeV4, v6tree *bool_tree.TreeV6) error {
|
|
||||||
filefp, err := os.Open(filename)
|
|
||||||
if err != nil {
|
|
||||||
return fmt.Errorf("cannot open file: %w", err)
|
|
||||||
}
|
|
||||||
|
|
||||||
go func(ctx context.Context, closer io.Closer) {
|
|
||||||
<-ctx.Done()
|
|
||||||
closer.Close()
|
|
||||||
}(ctx, filefp)
|
|
||||||
|
|
||||||
defer filefp.Close()
|
|
||||||
|
|
||||||
return f.updateTrees(mutex, filefp, v4tree, v6tree)
|
|
||||||
}
|
|
||||||
|
|
||||||
func (f *Firehol) updateRemoteURL(ctx context.Context, url string,
|
|
||||||
mutex sync.Locker,
|
|
||||||
v4tree *bool_tree.TreeV4, v6tree *bool_tree.TreeV6) error {
|
|
||||||
req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
|
|
||||||
if err != nil {
|
|
||||||
return fmt.Errorf("cannot build a request: %w", err)
|
|
||||||
}
|
|
||||||
|
|
||||||
resp, err := f.httpClient.Do(req) // nolint: bodyclose
|
|
||||||
if err != nil {
|
|
||||||
return fmt.Errorf("cannot request a remote URL %s: %w", url, err)
|
|
||||||
}
|
|
||||||
|
|
||||||
go func(ctx context.Context, closer io.Closer) {
|
|
||||||
<-ctx.Done()
|
|
||||||
closer.Close()
|
|
||||||
}(ctx, resp.Body)
|
|
||||||
|
|
||||||
defer func(rc io.ReadCloser) {
|
|
||||||
io.Copy(io.Discard, rc) // nolint: errcheck
|
|
||||||
rc.Close()
|
|
||||||
}(resp.Body)
|
|
||||||
|
|
||||||
return f.updateTrees(mutex, resp.Body, v4tree, v6tree)
|
|
||||||
}
|
|
||||||
|
|
||||||
func (f *Firehol) updateTrees(mutex sync.Locker,
|
|
||||||
reader io.Reader,
|
|
||||||
v4tree *bool_tree.TreeV4,
|
v4tree *bool_tree.TreeV4,
|
||||||
v6tree *bool_tree.TreeV6) error {
|
v6tree *bool_tree.TreeV6,
|
||||||
scanner := bufio.NewScanner(reader)
|
scanner *bufio.Scanner) error {
|
||||||
|
|
||||||
for scanner.Scan() {
|
for scanner.Scan() {
|
||||||
text := scanner.Text()
|
text := scanner.Text()
|
||||||
text = fireholRegexpComment.ReplaceAllLiteralString(text, "")
|
text = fireholRegexpComment.ReplaceAllLiteralString(text, "")
|
||||||
@@ -271,7 +188,7 @@ func (f *Firehol) updateTrees(mutex sync.Locker,
|
|||||||
}
|
}
|
||||||
|
|
||||||
if scanner.Err() != nil {
|
if scanner.Err() != nil {
|
||||||
return fmt.Errorf("cannot parse a response: %w", scanner.Err())
|
return fmt.Errorf("cannot parse a file: %w", scanner.Err())
|
||||||
}
|
}
|
||||||
|
|
||||||
return nil
|
return nil
|
||||||
@@ -317,27 +234,36 @@ func (f *Firehol) updateAddToTrees(ip net.IP, cidr uint,
|
|||||||
// when it is necessary.
|
// when it is necessary.
|
||||||
func NewFirehol(logger mtglib.Logger, network mtglib.Network,
|
func NewFirehol(logger mtglib.Logger, network mtglib.Network,
|
||||||
downloadConcurrency uint,
|
downloadConcurrency uint,
|
||||||
remoteURLs []string,
|
urls []string,
|
||||||
localFiles []string) (*Firehol, error) {
|
localFiles []string) (*Firehol, error) {
|
||||||
for _, v := range remoteURLs {
|
blocklists := []files.File{}
|
||||||
parsed, err := url.Parse(v)
|
|
||||||
if err != nil {
|
|
||||||
return nil, fmt.Errorf("incorrect url %s: %w", v, err)
|
|
||||||
}
|
|
||||||
|
|
||||||
switch parsed.Scheme {
|
|
||||||
case "http", "https":
|
|
||||||
default:
|
|
||||||
return nil, fmt.Errorf("unsupported url %s", v)
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
for _, v := range localFiles {
|
for _, v := range localFiles {
|
||||||
if stat, err := os.Stat(v); os.IsNotExist(err) || stat.IsDir() || stat.Mode().Perm()&0o400 == 0 {
|
file, err := files.NewLocal(v)
|
||||||
return nil, fmt.Errorf("%s is not a readable file", v)
|
if err != nil {
|
||||||
|
return nil, fmt.Errorf("cannot create a local file %s: %w", v, err)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
blocklists = append(blocklists, file)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
httpClient := network.MakeHTTPClient(nil)
|
||||||
|
|
||||||
|
for _, v := range urls {
|
||||||
|
file, err := files.NewHTTP(httpClient, v)
|
||||||
|
if err != nil {
|
||||||
|
return nil, fmt.Errorf("cannot create a HTTP file %s: %w", v, err)
|
||||||
|
}
|
||||||
|
|
||||||
|
blocklists = append(blocklists, file)
|
||||||
|
}
|
||||||
|
|
||||||
|
return NewFireholFromFiles(logger, downloadConcurrency, blocklists)
|
||||||
|
}
|
||||||
|
|
||||||
|
func NewFireholFromFiles(logger mtglib.Logger,
|
||||||
|
downloadConcurrency uint,
|
||||||
|
blocklists []files.File) (*Firehol, error) {
|
||||||
if downloadConcurrency == 0 {
|
if downloadConcurrency == 0 {
|
||||||
downloadConcurrency = DefaultFireholDownloadConcurrency
|
downloadConcurrency = DefaultFireholDownloadConcurrency
|
||||||
}
|
}
|
||||||
@@ -349,11 +275,9 @@ func NewFirehol(logger mtglib.Logger, network mtglib.Network,
|
|||||||
ctx: ctx,
|
ctx: ctx,
|
||||||
ctxCancel: cancel,
|
ctxCancel: cancel,
|
||||||
logger: logger.Named("firehol"),
|
logger: logger.Named("firehol"),
|
||||||
httpClient: network.MakeHTTPClient(nil),
|
|
||||||
treeV4: bool_tree.NewTreeV4(),
|
treeV4: bool_tree.NewTreeV4(),
|
||||||
treeV6: bool_tree.NewTreeV6(),
|
treeV6: bool_tree.NewTreeV6(),
|
||||||
workerPool: workerPool,
|
workerPool: workerPool,
|
||||||
remoteURLs: remoteURLs,
|
blocklists: blocklists,
|
||||||
localFiles: localFiles,
|
|
||||||
}, nil
|
}, nil
|
||||||
}
|
}
|
||||||
|
|||||||
Reference in New Issue
Block a user