From 558fec60de4c32e009a3dded49481c389af0844e Mon Sep 17 00:00:00 2001 From: 9seconds Date: Mon, 29 Nov 2021 16:25:33 +0300 Subject: [PATCH] Refactor firehol --- ipblocklist/files/http.go | 6 +- ipblocklist/files/http_test.go | 2 +- ipblocklist/files/init.go | 1 + ipblocklist/files/local.go | 14 +-- ipblocklist/firehol.go | 186 ++++++++++----------------------- 5 files changed, 69 insertions(+), 140 deletions(-) diff --git a/ipblocklist/files/http.go b/ipblocklist/files/http.go index 69b4cfc..6a60ede 100644 --- a/ipblocklist/files/http.go +++ b/ipblocklist/files/http.go @@ -22,7 +22,7 @@ func (h httpFile) Open(ctx context.Context) (io.ReadCloser, error) { response, err := h.http.Do(request) if err != nil { if response != nil { - io.Copy(io.Discard, response.Body) + io.Copy(io.Discard, response.Body) // nolint: errcheck response.Body.Close() } @@ -36,6 +36,10 @@ func (h httpFile) Open(ctx context.Context) (io.ReadCloser, error) { return response.Body, nil } +func (h httpFile) String() string { + return h.url +} + func NewHTTP(client *http.Client, endpoint string) (File, error) { if client == nil { return nil, ErrBadHTTPClient diff --git a/ipblocklist/files/http_test.go b/ipblocklist/files/http_test.go index 294739e..6559969 100644 --- a/ipblocklist/files/http_test.go +++ b/ipblocklist/files/http_test.go @@ -22,7 +22,7 @@ type HTTPTestSuite struct { } 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() { diff --git a/ipblocklist/files/init.go b/ipblocklist/files/init.go index 520b2c5..97570af 100644 --- a/ipblocklist/files/init.go +++ b/ipblocklist/files/init.go @@ -10,4 +10,5 @@ var ErrBadHTTPClient = errors.New("incorrect http client") type File interface { Open(context.Context) (io.ReadCloser, error) + String() string } diff --git a/ipblocklist/files/local.go b/ipblocklist/files/local.go index 8e3fe64..3cd08c7 100644 --- a/ipblocklist/files/local.go +++ b/ipblocklist/files/local.go @@ -4,18 +4,19 @@ import ( "context" "fmt" "io" - "io/fs" "os" - "path/filepath" ) type localFile struct { - root fs.FS - name string + path string } 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) { @@ -24,7 +25,6 @@ func NewLocal(path string) (File, error) { } return localFile{ - root: os.DirFS(filepath.Dir(path)), - name: filepath.Base(path), + path: path, }, nil } diff --git a/ipblocklist/firehol.go b/ipblocklist/firehol.go index 23f4cdd..41726a7 100644 --- a/ipblocklist/firehol.go +++ b/ipblocklist/firehol.go @@ -4,16 +4,13 @@ import ( "bufio" "context" "fmt" - "io" "net" - "net/http" - "net/url" - "os" "regexp" "strings" "sync" "time" + "github.com/9seconds/mtg/v2/ipblocklist/files" "github.com/9seconds/mtg/v2/mtglib" "github.com/kentik/patricia" "github.com/kentik/patricia/bool_tree" @@ -41,20 +38,16 @@ var fireholRegexpComment = regexp.MustCompile(`\s*#.*?$`) // 127.0.0.1 # you can specify an IP // 10.0.0.0/8 # or cidr type Firehol struct { - ctx context.Context - ctxCancel context.CancelFunc - logger mtglib.Logger - + ctx context.Context + ctxCancel context.CancelFunc + logger mtglib.Logger updateMutex sync.RWMutex - remoteURLs []string - localFiles []string + blocklists []files.File - httpClient *http.Client workerPool *ants.Pool - - treeV4 *bool_tree.TreeV4 - treeV6 *bool_tree.TreeV6 + treeV4 *bool_tree.TreeV4 + treeV6 *bool_tree.TreeV6 } // Shutdown stop a background update process. @@ -98,22 +91,14 @@ func (f *Firehol) Run(updateEach time.Duration) { } }() - if err := f.update(); err != nil { - f.logger.WarningError("cannot update blocklist", err) - } else { - f.logger.Info("blocklist was updated") - } + f.update() for { select { case <-f.ctx.Done(): return case <-ticker.C: - if err := f.update(); err != nil { - f.logger.WarningError("cannot update blocklist", err) - } else { - f.logger.Info("blocklist was updated") - } + f.update() } } } @@ -138,121 +123,53 @@ func (f *Firehol) containsIPv6(addr net.IP) bool { return false } -func (f *Firehol) update() error { // nolint: funlen, cyclop +func (f *Firehol) update() { ctx, cancel := context.WithCancel(f.ctx) defer cancel() wg := &sync.WaitGroup{} - wg.Add(len(f.remoteURLs) + len(f.localFiles)) + wg.Add(len(f.blocklists)) treeMutex := &sync.Mutex{} v4tree := bool_tree.NewTreeV4() v6tree := bool_tree.NewTreeV6() - errorChan := make(chan error, 1) - defer close(errorChan) - - for _, v := range f.localFiles { - go func(filename string) { + for _, v := range f.blocklists { + go func(file files.File) { defer wg.Done() - if err := f.updateLocalFile(ctx, filename, treeMutex, v4tree, v6tree); err != nil { - cancel() - f.logger.BindStr("filename", filename).WarningError("cannot update", err) + logger := f.logger.BindStr("filename", file.String()) - select { - case errorChan <- err: - default: - } + fileContent, err := file.Open(ctx) + if err != nil { + 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) } - 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() - select { - case err := <-errorChan: - return fmt.Errorf("cannot update trees: %w", err) - default: - } - f.updateMutex.Lock() defer f.updateMutex.Unlock() f.treeV4 = v4tree f.treeV6 = v6tree - return nil + f.logger.Info("blocklist was updated") } -func (f *Firehol) updateLocalFile(ctx context.Context, filename string, - 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, +func (f *Firehol) updateFromFile(mutex sync.Locker, v4tree *bool_tree.TreeV4, - v6tree *bool_tree.TreeV6) error { - scanner := bufio.NewScanner(reader) - + v6tree *bool_tree.TreeV6, + scanner *bufio.Scanner) error { for scanner.Scan() { text := scanner.Text() text = fireholRegexpComment.ReplaceAllLiteralString(text, "") @@ -271,7 +188,7 @@ func (f *Firehol) updateTrees(mutex sync.Locker, } 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 @@ -317,27 +234,36 @@ func (f *Firehol) updateAddToTrees(ip net.IP, cidr uint, // when it is necessary. func NewFirehol(logger mtglib.Logger, network mtglib.Network, downloadConcurrency uint, - remoteURLs []string, + urls []string, localFiles []string) (*Firehol, error) { - for _, v := range remoteURLs { - 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) - } - } + blocklists := []files.File{} for _, v := range localFiles { - if stat, err := os.Stat(v); os.IsNotExist(err) || stat.IsDir() || stat.Mode().Perm()&0o400 == 0 { - return nil, fmt.Errorf("%s is not a readable file", v) + file, err := files.NewLocal(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 { downloadConcurrency = DefaultFireholDownloadConcurrency } @@ -349,11 +275,9 @@ func NewFirehol(logger mtglib.Logger, network mtglib.Network, ctx: ctx, ctxCancel: cancel, logger: logger.Named("firehol"), - httpClient: network.MakeHTTPClient(nil), treeV4: bool_tree.NewTreeV4(), treeV6: bool_tree.NewTreeV6(), workerPool: workerPool, - remoteURLs: remoteURLs, - localFiles: localFiles, + blocklists: blocklists, }, nil }