diff options
Diffstat (limited to 'modules/headers/headers.go')
| -rw-r--r-- | modules/headers/headers.go | 42 |
1 files changed, 42 insertions, 0 deletions
diff --git a/modules/headers/headers.go b/modules/headers/headers.go new file mode 100644 index 0000000..877b370 --- /dev/null +++ b/modules/headers/headers.go @@ -0,0 +1,42 @@ +// This Source Code Form is subject to the terms of the Mozilla Public +// License, v. 2.0. If a copy of the MPL was not distributed with this +// file, You can obtain one at http://mozilla.org/MPL/2.0/. + +package headers + +import ( + "net/http" + + "github.com/philippta/flyscrape" +) + +func init() { + flyscrape.RegisterModule(Module{}) +} + +type Module struct { + Headers map[string]string `json:"headers"` +} + +func (Module) ModuleInfo() flyscrape.ModuleInfo { + return flyscrape.ModuleInfo{ + ID: "headers", + New: func() flyscrape.Module { return new(Module) }, + } +} + +func (m Module) AdaptTransport(t http.RoundTripper) http.RoundTripper { + if len(m.Headers) == 0 { + return t + } + + return flyscrape.RoundTripFunc(func(r *http.Request) (*http.Response, error) { + for k, v := range m.Headers { + r.Header.Set(k, v) + } + + return t.RoundTrip(r) + }) +} + +var _ flyscrape.TransportAdapter = Module{} |