/
chrome.go
130 lines (108 loc) · 3.22 KB
/
chrome.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
package http
import (
"context"
"fmt"
"io"
"log"
"net/http"
"os"
"strings"
"time"
"github.com/chromedp/cdproto/network"
"github.com/chromedp/chromedp"
)
// ChromeClient is an http client which uses a headless chrome browser to make requests
// Note that this client requires a chrome browser to be installed on the machine
type ChromeClient struct{}
// Do performs the http request using a headless chrome browser
func (c *ChromeClient) Do(req *http.Request) (*http.Response, error) {
if req.Method != http.MethodGet {
return nil, fmt.Errorf("only GET requests are supported")
}
dir, err := os.MkdirTemp("", "chromedp-tmp")
if err != nil {
return nil, err
}
defer os.RemoveAll(dir)
opts := append(chromedp.DefaultExecAllocatorOptions[:],
chromedp.DisableGPU,
chromedp.UserDataDir(dir),
chromedp.UserAgent(req.UserAgent()),
chromedp.WindowSize(1920, 1080),
)
allocCtx, cancel1 := chromedp.NewExecAllocator(context.Background(), opts...)
defer cancel1()
// also set up a custom logger
taskCtx, cancel2 := chromedp.NewContext(allocCtx, chromedp.WithLogf(log.Printf))
defer cancel2()
// Set up a timeout to prevent the script from running indefinitely
ctx, cancel3 := context.WithTimeout(taskCtx, 10*time.Second)
defer cancel3()
// Set up custom headers
headers := make(map[string]interface{})
for k, v := range req.Header {
headers[k] = strings.Join(v, ",")
}
//Listen for the response
var (
statusCode int64
responseHeaders network.Headers
)
url := req.URL.String()
chromedp.ListenTarget(ctx, func(event interface{}) {
switch responseReceivedEvent := event.(type) {
case *network.EventResponseReceived:
response := responseReceivedEvent.Response
if response.URL == url {
statusCode = response.Status
responseHeaders = response.Headers
}
}
})
// Navigate to the website
var body string
err = chromedp.Run(ctx,
chromedp.Tasks{
network.Enable(),
network.SetExtraHTTPHeaders(network.Headers(headers)),
chromedp.Navigate(url),
&myQueryAction{&body, &responseHeaders},
})
if err != nil {
return nil, err
}
return buildResponse(statusCode, responseHeaders, body), nil
}
// buildResponse builds an http response from the given status code, headers, and body
func buildResponse(statusCode int64, h network.Headers, body string) *http.Response {
headers := make(map[string][]string)
for k, v := range h {
headers[k] = []string{v.(string)}
}
resp := &http.Response{
StatusCode: int(statusCode),
Header: headers,
Body: io.NopCloser(strings.NewReader(string(body))),
ContentLength: int64(len(body)),
}
resp.Header.Set("Content-Length", fmt.Sprintf("%d", len(body)))
return resp
}
func isJsonResp(h network.Headers) bool {
for k, v := range h {
if strings.ToLower(k) == "content-type" {
return strings.Contains(strings.ToLower(v.(string)), "json")
}
}
return false
}
type myQueryAction struct {
body *string
responseHeaders *network.Headers
}
func (m *myQueryAction) Do(ctx context.Context) error {
if isJsonResp(*m.responseHeaders) {
return chromedp.InnerHTML(`pre`, m.body, chromedp.NodeVisible, chromedp.ByQuery).Do(ctx)
}
return chromedp.OuterHTML(`html`, m.body, chromedp.NodeVisible, chromedp.ByQuery).Do(ctx)
}