From cbdbbd249239345f88bea031beb55e84c2f47688 Mon Sep 17 00:00:00 2001
From: Rafi Ramadhana <42462215+rafiramadhana@users.noreply.github.com>
Date: Thu, 23 Nov 2023 18:58:41 +0700
Subject: Add custom request header (#18)

---
 modules/headers/headers.go | 42 ++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 42 insertions(+)
 create mode 100644 modules/headers/headers.go

(limited to 'modules/headers/headers.go')

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{}
-- 
cgit v1.2.3