optimize BypassUserAgentNames

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

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
 )

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=

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
 }