Przeglądaj źródła

Refactor firehol

tags/v2.1.3^2
9seconds 4 lat temu
rodzic
commit
558fec60de

+ 5
- 1
ipblocklist/files/http.go Wyświetl plik

22
 	response, err := h.http.Do(request)
22
 	response, err := h.http.Do(request)
23
 	if err != nil {
23
 	if err != nil {
24
 		if response != nil {
24
 		if response != nil {
25
-			io.Copy(io.Discard, response.Body)
25
+			io.Copy(io.Discard, response.Body) // nolint: errcheck
26
 			response.Body.Close()
26
 			response.Body.Close()
27
 		}
27
 		}
28
 
28
 
36
 	return response.Body, nil
36
 	return response.Body, nil
37
 }
37
 }
38
 
38
 
39
+func (h httpFile) String() string {
40
+	return h.url
41
+}
42
+
39
 func NewHTTP(client *http.Client, endpoint string) (File, error) {
43
 func NewHTTP(client *http.Client, endpoint string) (File, error) {
40
 	if client == nil {
44
 	if client == nil {
41
 		return nil, ErrBadHTTPClient
45
 		return nil, ErrBadHTTPClient

+ 1
- 1
ipblocklist/files/http_test.go Wyświetl plik

22
 }
22
 }
23
 
23
 
24
 func (suite *HTTPTestSuite) makeFile(path string) (files.File, error) {
24
 func (suite *HTTPTestSuite) makeFile(path string) (files.File, error) {
25
-	return files.NewHTTP(suite.httpClient, suite.httpServer.URL+"/"+path)
25
+	return files.NewHTTP(suite.httpClient, suite.httpServer.URL+"/"+path) // nolint: wrapcheck
26
 }
26
 }
27
 
27
 
28
 func (suite *HTTPTestSuite) SetupSuite() {
28
 func (suite *HTTPTestSuite) SetupSuite() {

+ 1
- 0
ipblocklist/files/init.go Wyświetl plik

10
 
10
 
11
 type File interface {
11
 type File interface {
12
 	Open(context.Context) (io.ReadCloser, error)
12
 	Open(context.Context) (io.ReadCloser, error)
13
+	String() string
13
 }
14
 }

+ 7
- 7
ipblocklist/files/local.go Wyświetl plik

4
 	"context"
4
 	"context"
5
 	"fmt"
5
 	"fmt"
6
 	"io"
6
 	"io"
7
-	"io/fs"
8
 	"os"
7
 	"os"
9
-	"path/filepath"
10
 )
8
 )
11
 
9
 
12
 type localFile struct {
10
 type localFile struct {
13
-	root fs.FS
14
-	name string
11
+	path string
15
 }
12
 }
16
 
13
 
17
 func (l localFile) Open(ctx context.Context) (io.ReadCloser, error) {
14
 func (l localFile) Open(ctx context.Context) (io.ReadCloser, error) {
18
-	return l.root.Open(l.name)
15
+	return os.Open(l.path) // nolint: wrapcheck
16
+}
17
+
18
+func (l localFile) String() string {
19
+	return l.path
19
 }
20
 }
20
 
21
 
21
 func NewLocal(path string) (File, error) {
22
 func NewLocal(path string) (File, error) {
24
 	}
25
 	}
25
 
26
 
26
 	return localFile{
27
 	return localFile{
27
-		root: os.DirFS(filepath.Dir(path)),
28
-		name: filepath.Base(path),
28
+		path: path,
29
 	}, nil
29
 	}, nil
30
 }
30
 }

+ 50
- 126
ipblocklist/firehol.go Wyświetl plik

4
 	"bufio"
4
 	"bufio"
5
 	"context"
5
 	"context"
6
 	"fmt"
6
 	"fmt"
7
-	"io"
8
 	"net"
7
 	"net"
9
-	"net/http"
10
-	"net/url"
11
-	"os"
12
 	"regexp"
8
 	"regexp"
13
 	"strings"
9
 	"strings"
14
 	"sync"
10
 	"sync"
15
 	"time"
11
 	"time"
16
 
12
 
13
+	"github.com/9seconds/mtg/v2/ipblocklist/files"
17
 	"github.com/9seconds/mtg/v2/mtglib"
14
 	"github.com/9seconds/mtg/v2/mtglib"
18
 	"github.com/kentik/patricia"
15
 	"github.com/kentik/patricia"
19
 	"github.com/kentik/patricia/bool_tree"
16
 	"github.com/kentik/patricia/bool_tree"
41
 //     127.0.0.1   # you can specify an IP
38
 //     127.0.0.1   # you can specify an IP
42
 //     10.0.0.0/8  # or cidr
39
 //     10.0.0.0/8  # or cidr
43
 type Firehol struct {
40
 type Firehol struct {
44
-	ctx       context.Context
45
-	ctxCancel context.CancelFunc
46
-	logger    mtglib.Logger
47
-
41
+	ctx         context.Context
42
+	ctxCancel   context.CancelFunc
43
+	logger      mtglib.Logger
48
 	updateMutex sync.RWMutex
44
 	updateMutex sync.RWMutex
49
 
45
 
50
-	remoteURLs []string
51
-	localFiles []string
46
+	blocklists []files.File
52
 
47
 
53
-	httpClient *http.Client
54
 	workerPool *ants.Pool
48
 	workerPool *ants.Pool
55
-
56
-	treeV4 *bool_tree.TreeV4
57
-	treeV6 *bool_tree.TreeV6
49
+	treeV4     *bool_tree.TreeV4
50
+	treeV6     *bool_tree.TreeV6
58
 }
51
 }
59
 
52
 
60
 // Shutdown stop a background update process.
53
 // Shutdown stop a background update process.
98
 		}
91
 		}
99
 	}()
92
 	}()
100
 
93
 
101
-	if err := f.update(); err != nil {
102
-		f.logger.WarningError("cannot update blocklist", err)
103
-	} else {
104
-		f.logger.Info("blocklist was updated")
105
-	}
94
+	f.update()
106
 
95
 
107
 	for {
96
 	for {
108
 		select {
97
 		select {
109
 		case <-f.ctx.Done():
98
 		case <-f.ctx.Done():
110
 			return
99
 			return
111
 		case <-ticker.C:
100
 		case <-ticker.C:
112
-			if err := f.update(); err != nil {
113
-				f.logger.WarningError("cannot update blocklist", err)
114
-			} else {
115
-				f.logger.Info("blocklist was updated")
116
-			}
101
+			f.update()
117
 		}
102
 		}
118
 	}
103
 	}
119
 }
104
 }
138
 	return false
123
 	return false
139
 }
124
 }
140
 
125
 
141
-func (f *Firehol) update() error { // nolint: funlen, cyclop
126
+func (f *Firehol) update() {
142
 	ctx, cancel := context.WithCancel(f.ctx)
127
 	ctx, cancel := context.WithCancel(f.ctx)
143
 	defer cancel()
128
 	defer cancel()
144
 
129
 
145
 	wg := &sync.WaitGroup{}
130
 	wg := &sync.WaitGroup{}
146
-	wg.Add(len(f.remoteURLs) + len(f.localFiles))
131
+	wg.Add(len(f.blocklists))
147
 
132
 
148
 	treeMutex := &sync.Mutex{}
133
 	treeMutex := &sync.Mutex{}
149
 	v4tree := bool_tree.NewTreeV4()
134
 	v4tree := bool_tree.NewTreeV4()
150
 	v6tree := bool_tree.NewTreeV6()
135
 	v6tree := bool_tree.NewTreeV6()
151
 
136
 
152
-	errorChan := make(chan error, 1)
153
-	defer close(errorChan)
154
-
155
-	for _, v := range f.localFiles {
156
-		go func(filename string) {
137
+	for _, v := range f.blocklists {
138
+		go func(file files.File) {
157
 			defer wg.Done()
139
 			defer wg.Done()
158
 
140
 
159
-			if err := f.updateLocalFile(ctx, filename, treeMutex, v4tree, v6tree); err != nil {
160
-				cancel()
161
-				f.logger.BindStr("filename", filename).WarningError("cannot update", err)
141
+			logger := f.logger.BindStr("filename", file.String())
162
 
142
 
163
-				select {
164
-				case errorChan <- err:
165
-				default:
166
-				}
167
-			}
168
-		}(v)
169
-	}
143
+			fileContent, err := file.Open(ctx)
144
+			if err != nil {
145
+				logger.WarningError("update has failed", err)
170
 
146
 
171
-	for _, v := range f.remoteURLs {
172
-		value := v
173
-
174
-		f.workerPool.Submit(func() { // nolint: errcheck
175
-			defer wg.Done()
147
+				return
148
+			}
176
 
149
 
177
-			if err := f.updateRemoteURL(ctx, value, treeMutex, v4tree, v6tree); err != nil {
178
-				cancel()
179
-				f.logger.BindStr("url", value).WarningError("cannot update", err)
150
+			defer fileContent.Close()
180
 
151
 
181
-				select {
182
-				case errorChan <- err:
183
-				default:
184
-				}
152
+			if err := f.updateFromFile(treeMutex, v4tree, v6tree, bufio.NewScanner(fileContent)); err != nil {
153
+				logger.WarningError("update has failed", err)
185
 			}
154
 			}
186
-		})
155
+		}(v)
187
 	}
156
 	}
188
 
157
 
189
 	wg.Wait()
158
 	wg.Wait()
190
 
159
 
191
-	select {
192
-	case err := <-errorChan:
193
-		return fmt.Errorf("cannot update trees: %w", err)
194
-	default:
195
-	}
196
-
197
 	f.updateMutex.Lock()
160
 	f.updateMutex.Lock()
198
 	defer f.updateMutex.Unlock()
161
 	defer f.updateMutex.Unlock()
199
 
162
 
200
 	f.treeV4 = v4tree
163
 	f.treeV4 = v4tree
201
 	f.treeV6 = v6tree
164
 	f.treeV6 = v6tree
202
 
165
 
203
-	return nil
204
-}
205
-
206
-func (f *Firehol) updateLocalFile(ctx context.Context, filename string,
207
-	mutex sync.Locker,
208
-	v4tree *bool_tree.TreeV4, v6tree *bool_tree.TreeV6) error {
209
-	filefp, err := os.Open(filename)
210
-	if err != nil {
211
-		return fmt.Errorf("cannot open file: %w", err)
212
-	}
213
-
214
-	go func(ctx context.Context, closer io.Closer) {
215
-		<-ctx.Done()
216
-		closer.Close()
217
-	}(ctx, filefp)
218
-
219
-	defer filefp.Close()
220
-
221
-	return f.updateTrees(mutex, filefp, v4tree, v6tree)
166
+	f.logger.Info("blocklist was updated")
222
 }
167
 }
223
 
168
 
224
-func (f *Firehol) updateRemoteURL(ctx context.Context, url string,
225
-	mutex sync.Locker,
226
-	v4tree *bool_tree.TreeV4, v6tree *bool_tree.TreeV6) error {
227
-	req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
228
-	if err != nil {
229
-		return fmt.Errorf("cannot build a request: %w", err)
230
-	}
231
-
232
-	resp, err := f.httpClient.Do(req) // nolint: bodyclose
233
-	if err != nil {
234
-		return fmt.Errorf("cannot request a remote URL %s: %w", url, err)
235
-	}
236
-
237
-	go func(ctx context.Context, closer io.Closer) {
238
-		<-ctx.Done()
239
-		closer.Close()
240
-	}(ctx, resp.Body)
241
-
242
-	defer func(rc io.ReadCloser) {
243
-		io.Copy(io.Discard, rc) // nolint: errcheck
244
-		rc.Close()
245
-	}(resp.Body)
246
-
247
-	return f.updateTrees(mutex, resp.Body, v4tree, v6tree)
248
-}
249
-
250
-func (f *Firehol) updateTrees(mutex sync.Locker,
251
-	reader io.Reader,
169
+func (f *Firehol) updateFromFile(mutex sync.Locker,
252
 	v4tree *bool_tree.TreeV4,
170
 	v4tree *bool_tree.TreeV4,
253
-	v6tree *bool_tree.TreeV6) error {
254
-	scanner := bufio.NewScanner(reader)
255
-
171
+	v6tree *bool_tree.TreeV6,
172
+	scanner *bufio.Scanner) error {
256
 	for scanner.Scan() {
173
 	for scanner.Scan() {
257
 		text := scanner.Text()
174
 		text := scanner.Text()
258
 		text = fireholRegexpComment.ReplaceAllLiteralString(text, "")
175
 		text = fireholRegexpComment.ReplaceAllLiteralString(text, "")
271
 	}
188
 	}
272
 
189
 
273
 	if scanner.Err() != nil {
190
 	if scanner.Err() != nil {
274
-		return fmt.Errorf("cannot parse a response: %w", scanner.Err())
191
+		return fmt.Errorf("cannot parse a file: %w", scanner.Err())
275
 	}
192
 	}
276
 
193
 
277
 	return nil
194
 	return nil
317
 // when it is necessary.
234
 // when it is necessary.
318
 func NewFirehol(logger mtglib.Logger, network mtglib.Network,
235
 func NewFirehol(logger mtglib.Logger, network mtglib.Network,
319
 	downloadConcurrency uint,
236
 	downloadConcurrency uint,
320
-	remoteURLs []string,
237
+	urls []string,
321
 	localFiles []string) (*Firehol, error) {
238
 	localFiles []string) (*Firehol, error) {
322
-	for _, v := range remoteURLs {
323
-		parsed, err := url.Parse(v)
239
+	blocklists := []files.File{}
240
+
241
+	for _, v := range localFiles {
242
+		file, err := files.NewLocal(v)
324
 		if err != nil {
243
 		if err != nil {
325
-			return nil, fmt.Errorf("incorrect url %s: %w", v, err)
244
+			return nil, fmt.Errorf("cannot create a local file %s: %w", v, err)
326
 		}
245
 		}
327
 
246
 
328
-		switch parsed.Scheme {
329
-		case "http", "https":
330
-		default:
331
-			return nil, fmt.Errorf("unsupported url %s", v)
332
-		}
247
+		blocklists = append(blocklists, file)
333
 	}
248
 	}
334
 
249
 
335
-	for _, v := range localFiles {
336
-		if stat, err := os.Stat(v); os.IsNotExist(err) || stat.IsDir() || stat.Mode().Perm()&0o400 == 0 {
337
-			return nil, fmt.Errorf("%s is not a readable file", v)
250
+	httpClient := network.MakeHTTPClient(nil)
251
+
252
+	for _, v := range urls {
253
+		file, err := files.NewHTTP(httpClient, v)
254
+		if err != nil {
255
+			return nil, fmt.Errorf("cannot create a HTTP file %s: %w", v, err)
338
 		}
256
 		}
257
+
258
+		blocklists = append(blocklists, file)
339
 	}
259
 	}
340
 
260
 
261
+	return NewFireholFromFiles(logger, downloadConcurrency, blocklists)
262
+}
263
+
264
+func NewFireholFromFiles(logger mtglib.Logger,
265
+	downloadConcurrency uint,
266
+	blocklists []files.File) (*Firehol, error) {
341
 	if downloadConcurrency == 0 {
267
 	if downloadConcurrency == 0 {
342
 		downloadConcurrency = DefaultFireholDownloadConcurrency
268
 		downloadConcurrency = DefaultFireholDownloadConcurrency
343
 	}
269
 	}
349
 		ctx:        ctx,
275
 		ctx:        ctx,
350
 		ctxCancel:  cancel,
276
 		ctxCancel:  cancel,
351
 		logger:     logger.Named("firehol"),
277
 		logger:     logger.Named("firehol"),
352
-		httpClient: network.MakeHTTPClient(nil),
353
 		treeV4:     bool_tree.NewTreeV4(),
278
 		treeV4:     bool_tree.NewTreeV4(),
354
 		treeV6:     bool_tree.NewTreeV6(),
279
 		treeV6:     bool_tree.NewTreeV6(),
355
 		workerPool: workerPool,
280
 		workerPool: workerPool,
356
-		remoteURLs: remoteURLs,
357
-		localFiles: localFiles,
281
+		blocklists: blocklists,
358
 	}, nil
282
 	}, nil
359
 }
283
 }

Ładowanie…
Anuluj
Zapisz