add max wait and block requests

Author: ysmood

.golangci.yml

@@ -4,6 +4,7 @@ run:
 linters:
   enable-all: true
   disable:
+    - wsl
     - contextcheck
     - gomnd
     - gochecknoinits

cmd/serve/main.go

@@ -4,6 +4,8 @@ import (
 	"flag"
 	"log"
 	"net/http"
+	"strings"
+	"time"
 
 	"github.com/go-rod/bartender"
 )
@@ -12,6 +14,10 @@ func main() {
 	port := flag.String("p", ":3001", "port to listen on")
 	target := flag.String("t", "", "target url to proxy")
 	size := flag.Int("s", 2, "size of the pool")
+	maxWait := flag.Duration("w", 3*time.Second, "max wait time for a page rendering")
+
+	var block BlockRequestsFlag
+	flag.Var(&block, "b", "block the requests that match the pattern, such as 'https://a.com/*', can set multiple ones")
 
 	flag.Parse()
 
@@ -21,8 +27,25 @@ func main() {
 
 	log.Printf("Bartender started %s -> %s\n", *port, *target)
 
-	err := http.ListenAndServe(*port, bartender.New(*port, *target, *size))
+	b := bartender.New(*port, *target, *size)
+	b.BlockRequest(block...)
+	b.MaxWait(*maxWait)
+	b.WarnUp()
+
+	err := http.ListenAndServe(*port, b)
 	if err != nil {
 		log.Fatalln(err)
 	}
 }
+
+type BlockRequestsFlag []string
+
+func (i *BlockRequestsFlag) String() string {
+	return strings.Join(*i, ", ")
+}
+
+func (i *BlockRequestsFlag) Set(value string) error {
+	*i = append(*i, value)
+
+	return nil
+}

service.go

@@ -8,17 +8,22 @@ import (
 	"net/http/httputil"
 	"net/url"
 	"strings"
+	"sync"
+	"time"
 
 	"github.com/go-rod/rod"
+	"github.com/go-rod/rod/lib/proto"
 	"github.com/mileusna/useragent"
 )
 
 type Bartender struct {
-	addr       string
-	target     *url.URL
-	proxy      *httputil.ReverseProxy
-	bypassList map[string]bool
-	pool       rod.PagePool
+	addr          string
+	target        *url.URL
+	proxy         *httputil.ReverseProxy
+	bypassList    map[string]bool
+	pool          rod.PagePool
+	blockRequests []string
+	maxWait       time.Duration
 }
 
 func New(addr, target string, poolSize int) *Bartender {
@@ -45,14 +50,53 @@ func New(addr, target string, poolSize int) *Bartender {
 			useragent.Edge:             true,
 			useragent.Vivaldi:          true,
 		},
-		pool: rod.NewPagePool(poolSize),
+		pool:          rod.NewPagePool(poolSize),
+		blockRequests: []string{},
+		maxWait:       3 * time.Second,
 	}
 }
 
 func (b *Bartender) BypassUserAgentNames(list map[string]bool) {
 	b.bypassList = list
 }
 
+func (b *Bartender) BlockRequest(patterns ...string) {
+	b.blockRequests = patterns
+}
+
+// MaxWait sets the max wait time for the headless browser to render the page.
+// If the max wait time is reached, bartender will stop waiting for page rendering and
+// immediately return the current html.
+func (b *Bartender) MaxWait(d time.Duration) {
+	b.maxWait = d
+}
+
+func (b *Bartender) newPage() *rod.Page {
+	page := rod.New().MustConnect().MustPage()
+
+	if len(b.blockRequests) > 0 {
+		router := page.HijackRequests()
+
+		for _, pattern := range b.blockRequests {
+			router.MustAdd(pattern, func(ctx *rod.Hijack) {
+				ctx.Response.Fail(proto.NetworkErrorReasonBlockedByClient)
+			})
+		}
+
+		go router.Run()
+	}
+
+	log.Println("headless browser started:", page.SessionID)
+
+	return page
+}
+
+func (b *Bartender) WarnUp() {
+	for i := 0; i < len(b.pool); i++ {
+		b.pool.Put(b.pool.Get(b.newPage))
+	}
+}
+
 func (b *Bartender) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 	ua := useragent.Parse(r.Header.Get("User-Agent"))
 	if r.Method != http.MethodGet || b.bypassList[ua.Name] {
@@ -80,11 +124,6 @@ func (b *Bartender) RenderPage(w http.ResponseWriter, r *http.Request) bool {
 
 	log.Println("headless render:", u)
 
-	page := b.pool.Get(func() *rod.Page { return rod.New().MustConnect().MustPage() })
-	defer b.pool.Put(page)
-
-	page.MustNavigate(u).MustWaitStable()
-
 	for k, vs := range resHeader {
 		if k == "Content-Length" {
 			continue
@@ -97,10 +136,34 @@ func (b *Bartender) RenderPage(w http.ResponseWriter, r *http.Request) bool {
 
 	w.WriteHeader(statusCode)
 
-	_, err := w.Write([]byte(page.MustHTML()))
-	if err != nil {
-		panic(err)
-	}
+	page := b.pool.Get(b.newPage)
+	defer b.pool.Put(page)
+
+	page, cancel := page.WithCancel()
+
+	once := sync.Once{}
+
+	go func() {
+		time.Sleep(b.maxWait)
+		log.Println("max wait time reached, return current html:", u)
+		once.Do(func() {
+			body, _ := page.HTML()
+			_, _ = w.Write([]byte(body))
+			cancel()
+		})
+	}()
+
+	_ = page.Context(r.Context()).Navigate(u)
+
+	_ = page.WaitStable(time.Second)
+
+	body, _ := page.HTML()
+
+	log.Println("headless render done:", u)
+
+	once.Do(func() {
+		_, _ = w.Write([]byte(body))
+	})
 
 	return true
 }