From cb64ab487dfe26ada79c3394abdda805e33b0d34 Mon Sep 17 00:00:00 2001
From: Yad Smood <y.s.inside@gmail.com>
Date: Mon, 24 Jul 2023 18:07:20 +0800
Subject: [PATCH] optimize BypassUserAgentNames

---
 cmd/serve/main.go | 17 +++++++++++------
 go.mod            |  2 +-
 go.sum            |  4 ++--
 service.go        |  9 ++++++---
 4 files changed, 20 insertions(+), 12 deletions(-)

diff --git a/cmd/serve/main.go b/cmd/serve/main.go
index dd10764..7bf78ad 100644
--- a/cmd/serve/main.go
+++ b/cmd/serve/main.go
@@ -16,8 +16,12 @@ func main() {
 	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")
+	var bypassUAs StringsFlag
+	flag.Var(&bypassUAs, "u", "bypass the specified user-agent names")
+
+	var blockList StringsFlag
+	flag.Var(&blockList, "b",
+		"block the requests that match the pattern, such as 'https://a.com/*', can set multiple ones")
 
 	flag.Parse()
 
@@ -28,7 +32,8 @@ func main() {
 	log.Printf("Bartender started %s -> %s\n", *port, *target)
 
 	b := bartender.New(*port, *target, *size)
-	b.BlockRequest(block...)
+	b.BlockRequests(blockList...)
+	b.BypassUserAgentNames(bypassUAs...)
 	b.MaxWait(*maxWait)
 	b.WarnUp()
 
@@ -38,13 +43,13 @@ func main() {
 	}
 }
 
-type BlockRequestsFlag []string
+type StringsFlag []string
 
-func (i *BlockRequestsFlag) String() string {
+func (i *StringsFlag) String() string {
 	return strings.Join(*i, ", ")
 }
 
-func (i *BlockRequestsFlag) Set(value string) error {
+func (i *StringsFlag) Set(value string) error {
 	*i = append(*i, value)
 
 	return nil
diff --git a/go.mod b/go.mod
index 82be161..3d97c48 100644
--- a/go.mod
+++ b/go.mod
@@ -3,7 +3,7 @@ module github.com/go-rod/bartender
 go 1.20
 
 require (
-	github.com/go-rod/rod v0.114.0
+	github.com/go-rod/rod v0.114.1
 	github.com/mileusna/useragent v1.3.3
 	github.com/ysmood/got v0.34.1
 )
diff --git a/go.sum b/go.sum
index e699e86..eae9a0d 100644
--- a/go.sum
+++ b/go.sum
@@ -1,5 +1,5 @@
-github.com/go-rod/rod v0.114.0 h1:P+zLOqsj+vKf4C86SfjP6ymyPl9VXoYKm+ceCeQms6Y=
-github.com/go-rod/rod v0.114.0/go.mod h1:aiedSEFg5DwG/fnNbUOTPMTTWX3MRj6vIs/a684Mthw=
+github.com/go-rod/rod v0.114.1 h1:osBWr88guzTXAIzwJWVmGZe3/utT9+lqKjkGSBsYMxw=
+github.com/go-rod/rod v0.114.1/go.mod h1:aiedSEFg5DwG/fnNbUOTPMTTWX3MRj6vIs/a684Mthw=
 github.com/mileusna/useragent v1.3.3 h1:hrIVmPevJY3ICS1Ob4yjqJToQiv2eD9iHaJBjxMihWY=
 github.com/mileusna/useragent v1.3.3/go.mod h1:3d8TOmwL/5I8pJjyVDteHtgDGcefrFUX4ccGOMKNYYc=
 github.com/ysmood/fetchup v0.2.3 h1:ulX+SonA0Vma5zUFXtv52Kzip/xe7aj4vqT5AJwQ+ZQ=
diff --git a/service.go b/service.go
index 92a178f..211e2fc 100644
--- a/service.go
+++ b/service.go
@@ -56,11 +56,14 @@ func New(addr, target string, poolSize int) *Bartender {
 	}
 }
 
-func (b *Bartender) BypassUserAgentNames(list map[string]bool) {
-	b.bypassList = list
+func (b *Bartender) BypassUserAgentNames(list ...string) {
+	b.bypassList = map[string]bool{}
+	for _, ua := range list {
+		b.bypassList[ua] = true
+	}
 }
 
-func (b *Bartender) BlockRequest(patterns ...string) {
+func (b *Bartender) BlockRequests(patterns ...string) {
 	b.blockRequests = patterns
 }