Skip to content

Commit

Permalink
fix: default bypass
Browse files Browse the repository at this point in the history
  • Loading branch information
ysmood committed Jul 25, 2023
1 parent cb64ab4 commit 31329b3
Show file tree
Hide file tree
Showing 2 changed files with 30 additions and 20 deletions.
4 changes: 3 additions & 1 deletion cmd/serve/main.go
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,7 @@ 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 bypassUAs StringsFlag
var bypassUAs StringsFlag = bartender.DefaultBypassUserAgentNames
flag.Var(&bypassUAs, "u", "bypass the specified user-agent names")

var blockList StringsFlag
Expand All @@ -30,6 +30,8 @@ func main() {
}

log.Printf("Bartender started %s -> %s\n", *port, *target)
log.Printf("Block list: %v\n", blockList)
log.Printf("Bypass user-agent names: %v\n", bypassUAs)

b := bartender.New(*port, *target, *size)
b.BlockRequests(blockList...)
Expand Down
46 changes: 27 additions & 19 deletions service.go
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,19 @@ import (
"github.com/mileusna/useragent"
)

var DefaultBypassUserAgentNames = []string{
useragent.Opera,
useragent.OperaMini,
useragent.OperaTouch,
useragent.Chrome,
useragent.HeadlessChrome,
useragent.Firefox,
useragent.InternetExplorer,
useragent.Safari,
useragent.Edge,
useragent.Vivaldi,
}

type Bartender struct {
addr string
target *url.URL
Expand All @@ -35,32 +48,18 @@ func New(addr, target string, poolSize int) *Bartender {
proxy := httputil.NewSingleHostReverseProxy(u)

return &Bartender{
addr: addr,
target: u,
proxy: proxy,
bypassList: map[string]bool{
useragent.Opera: true,
useragent.OperaMini: true,
useragent.OperaTouch: true,
useragent.Chrome: true,
useragent.HeadlessChrome: true,
useragent.Firefox: true,
useragent.InternetExplorer: true,
useragent.Safari: true,
useragent.Edge: true,
useragent.Vivaldi: true,
},
addr: addr,
target: u,
proxy: proxy,
bypassList: strToMap(DefaultBypassUserAgentNames),
pool: rod.NewPagePool(poolSize),
blockRequests: []string{},
maxWait: 3 * time.Second,
}
}

func (b *Bartender) BypassUserAgentNames(list ...string) {
b.bypassList = map[string]bool{}
for _, ua := range list {
b.bypassList[ua] = true
}
b.bypassList = strToMap(list)
}

func (b *Bartender) BlockRequests(patterns ...string) {
Expand Down Expand Up @@ -198,3 +197,12 @@ func getHeader(ctx context.Context, u string) (int, http.Header) {
func htmlContentType(h http.Header) bool {
return strings.Contains(h.Get("Content-Type"), "text/html")
}

func strToMap(list []string) map[string]bool {
m := map[string]bool{}
for _, s := range list {
m[s] = true
}

return m
}

0 comments on commit 31329b3

Please sign in to comment.