Skip to content

Commit 43c9e93

Browse files
ysmoodysmood
committed
bypass resources other than html
1 parent b831542 commit 43c9e93

File tree

5 files changed

+80
-14
lines changed

5 files changed

+80
-14
lines changed

‎.golangci.yml

+1
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@ run:
44
linters:
55
enable-all: true
66
disable:
7+
- contextcheck
78
- gomnd
89
- gochecknoinits
910
- paralleltest

‎go.mod

+1-1
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@ module github.com/go-rod/bartender
33
go 1.20
44

55
require (
6-
github.com/go-rod/rod v0.113.4
6+
github.com/go-rod/rod v0.114.0
77
github.com/mileusna/useragent v1.3.3
88
github.com/ysmood/got v0.34.1
99
)

‎go.sum

+2-2
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
1-
github.com/go-rod/rod v0.113.4 h1:O5a/VTl1h6a08ecXkOF9oVJBi8fjBnF0FtyMSgd3KaY=
2-
github.com/go-rod/rod v0.113.4/go.mod h1:aiedSEFg5DwG/fnNbUOTPMTTWX3MRj6vIs/a684Mthw=
1+
github.com/go-rod/rod v0.114.0 h1:P+zLOqsj+vKf4C86SfjP6ymyPl9VXoYKm+ceCeQms6Y=
2+
github.com/go-rod/rod v0.114.0/go.mod h1:aiedSEFg5DwG/fnNbUOTPMTTWX3MRj6vIs/a684Mthw=
33
github.com/mileusna/useragent v1.3.3 h1:hrIVmPevJY3ICS1Ob4yjqJToQiv2eD9iHaJBjxMihWY=
44
github.com/mileusna/useragent v1.3.3/go.mod h1:3d8TOmwL/5I8pJjyVDteHtgDGcefrFUX4ccGOMKNYYc=
55
github.com/ysmood/fetchup v0.2.3 h1:ulX+SonA0Vma5zUFXtv52Kzip/xe7aj4vqT5AJwQ+ZQ=

‎service.go

+64-9
Original file line numberDiff line numberDiff line change
@@ -2,10 +2,12 @@
22
package bartender
33

44
import (
5+
"context"
56
"log"
67
"net/http"
78
"net/http/httputil"
89
"net/url"
10+
"strings"
911

1012
"github.com/go-rod/rod"
1113
"github.com/mileusna/useragent"
@@ -53,27 +55,80 @@ func (b *Bartender) BypassUserAgentNames(list map[string]bool) {
5355

5456
func (b *Bartender) ServeHTTP(w http.ResponseWriter, r *http.Request) {
5557
ua := useragent.Parse(r.Header.Get("User-Agent"))
56-
if b.bypassList[ua.Name] {
58+
if r.Method != http.MethodGet || b.bypassList[ua.Name] {
5759
b.proxy.ServeHTTP(w, r)
5860

5961
return
6062
}
6163

62-
b.RenderPage(w, r)
64+
if b.RenderPage(w, r) {
65+
return
66+
}
67+
68+
b.proxy.ServeHTTP(w, r)
6369
}
6470

65-
func (b *Bartender) RenderPage(w http.ResponseWriter, r *http.Request) {
66-
log.Println("headless render:", r.URL.String())
71+
// RenderPage returns true if the page is rendered by the headless browser.
72+
func (b *Bartender) RenderPage(w http.ResponseWriter, r *http.Request) bool {
73+
u := b.getTargetURL(r.URL)
74+
75+
statusCode, resHeader := getHeader(r.Context(), u)
76+
77+
if !htmlContentType(resHeader) {
78+
return false
79+
}
80+
81+
log.Println("headless render:", u)
82+
83+
page := b.pool.Get(func() *rod.Page { return rod.New().MustConnect().MustPage() })
84+
defer b.pool.Put(page)
85+
86+
page.MustNavigate(u).MustWaitStable()
87+
88+
for k, vs := range resHeader {
89+
if k == "Content-Length" {
90+
continue
91+
}
92+
93+
for _, v := range vs {
94+
w.Header().Add(k, v)
95+
}
96+
}
97+
98+
w.WriteHeader(statusCode)
6799

68-
u := *r.URL
100+
_, err := w.Write([]byte(page.MustHTML()))
101+
if err != nil {
102+
panic(err)
103+
}
104+
105+
return true
106+
}
107+
108+
func (b *Bartender) getTargetURL(reqURL *url.URL) string {
109+
u := *reqURL
69110
u.Scheme = b.target.Scheme
70111
u.Host = b.target.Host
71112

72-
page := b.pool.Get(func() *rod.Page { return rod.New().MustConnect().MustPage() })
113+
return u.String()
114+
}
73115

74-
page.MustNavigate(u.String()).MustWaitStable()
116+
func getHeader(ctx context.Context, u string) (int, http.Header) {
117+
req, err := http.NewRequestWithContext(ctx, http.MethodGet, u, nil)
118+
if err != nil {
119+
panic(err)
120+
}
75121

76-
w.Header().Set("Content-Type", "text/html; charset=utf-8")
122+
res, err := http.DefaultClient.Do(req)
123+
if err != nil {
124+
panic(err)
125+
}
126+
127+
_ = res.Body.Close()
128+
129+
return res.StatusCode, res.Header
130+
}
77131

78-
_, _ = w.Write([]byte(page.MustHTML()))
132+
func htmlContentType(h http.Header) bool {
133+
return strings.Contains(h.Get("Content-Type"), "text/html")
79134
}

‎service_test.go

+12-2
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,10 @@ import (
1111
func TestBasic(t *testing.T) {
1212
g := got.T(t)
1313

14-
website := g.Serve().Route("/", ".html", `<html>
14+
website := g.Serve()
15+
16+
website.Route("/a.png", ".png", "image")
17+
website.Route("/", ".html", `<html>
1518
<body></body>
1619
<script>
1720
window.onload = () => {
@@ -27,8 +30,9 @@ func TestBasic(t *testing.T) {
2730
proxy.Mux.HandleFunc("/", bt.ServeHTTP)
2831

2932
{
33+
//nolint: lll
3034
// browser
31-
res := g.Req("", proxy.URL("/test?q=ok"), http.Header{"Accept-Language": {"en"}})
35+
res := g.Req("", proxy.URL("/test?q=ok"), http.Header{"User-Agent": {"Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/114.0.0.0 Safari/537.36"}})
3236
g.Has(res.String(), "<body></body>")
3337
}
3438

@@ -37,4 +41,10 @@ func TestBasic(t *testing.T) {
3741
res := g.Req("", proxy.URL("/test?q=ok"))
3842
g.Has(res.String(), "/test?q=ok")
3943
}
44+
45+
{
46+
// can get image
47+
res := g.Req("", proxy.URL("/a.png"))
48+
g.Has(res.String(), "image")
49+
}
4050
}

0 commit comments

Comments
 (0)