Refactor firehol

This commit is contained in:
9seconds
2021-11-29 16:25:33 +03:00
parent cc101c9a47
commit 558fec60de
5 changed files with 69 additions and 140 deletions
+5 -1
View File
@@ -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
+1 -1
View File
@@ -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() {
+1
View File
@@ -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
} }
+7 -7
View File
@@ -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
View File
@@ -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
} }