Merge pull request #204 from grafana/feat/105-k6-blockhostnames

Add request blocking based on k6 `blockHostnames` option

Author: Ivan Mirić

common/browser.go

@@ -230,9 +230,11 @@ func (b *Browser) onAttachedToTarget(ev *target.EventAttachedToTarget) {
 		return
 	}
 
+	session := b.connSessions.getSession(ev.SessionID)
+
 	switch evti.Type {
 	case "background_page":
-		p, err := NewPage(b.ctx, b.connSessions.getSession(ev.SessionID), browserCtx, evti.TargetID, nil, false, b.logger)
+		p, err := NewPage(b.ctx, session, browserCtx, evti.TargetID, nil, false, b.logger)
 		if err != nil {
 			isRunning := atomic.LoadInt64(&b.state) == BrowserStateOpen && b.IsConnected() //b.conn.isConnected()
 			if _, ok := err.(*websocket.CloseError); !ok && !isRunning {
@@ -272,7 +274,7 @@ func (b *Browser) onAttachedToTarget(ev *target.EventAttachedToTarget) {
 
 		b.logger.Debugf("Browser:onAttachedToTarget:page", "sid:%v tid:%v opener nil:%t", ev.SessionID, evti.TargetID, opener == nil)
 
-		p, err := NewPage(b.ctx, b.connSessions.getSession(ev.SessionID), browserCtx, evti.TargetID, opener, true, b.logger)
+		p, err := NewPage(b.ctx, session, browserCtx, evti.TargetID, opener, true, b.logger)
 		if err != nil {
 			isRunning := atomic.LoadInt64(&b.state) == BrowserStateOpen && b.IsConnected() //b.conn.isConnected()
 			if _, ok := err.(*websocket.CloseError); !ok && !isRunning {

common/browser_context.go

@@ -157,7 +157,7 @@ func (b *BrowserContext) ExposeBinding(name string, callback goja.Callable, opts
 }
 
 func (b *BrowserContext) ExposeFunction(name string, callback goja.Callable) {
-	k6Throw(b.ctx, "BrowserContext.newCDPSession(name, callback) has not been implemented yet")
+	k6Throw(b.ctx, "BrowserContext.exposeFunction(name, callback) has not been implemented yet")
 }
 
 // GrantPermissions enables the specified permissions, all others will be disabled.
@@ -245,7 +245,7 @@ func (b *BrowserContext) Pages() []api.Page {
 }
 
 func (b *BrowserContext) Route(url goja.Value, handler goja.Callable) {
-	k6Throw(b.ctx, "BrowserContext.setHTTPCredentials(httpCredentials) has not been implemented yet")
+	k6Throw(b.ctx, "BrowserContext.route(url, handler) has not been implemented yet")
 }
 
 // SetDefaultNavigationTimeout sets the default navigation timeout in milliseconds.
@@ -263,7 +263,7 @@ func (b *BrowserContext) SetDefaultTimeout(timeout int64) {
 }
 
 func (b *BrowserContext) SetExtraHTTPHeaders(headers map[string]string) {
-	k6Throw(b.ctx, "BrowserContext.setHTTPCredentials(httpCredentials) has not been implemented yet")
+	k6Throw(b.ctx, "BrowserContext.setExtraHTTPHeaders(headers) has not been implemented yet")
 }
 
 // SetGeolocation overrides the geo location of the user.

common/frame_manager.go

@@ -24,6 +24,7 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"strings"
 	"sync"
 	"sync/atomic"
 	"time"
@@ -573,7 +574,8 @@ func (m *FrameManager) NavigateFrame(frame *Frame, url string, opts goja.Value)
 		"fmid:%d fid:%v furl:%s url:%s", fmid, fid, furl, url)
 
 	rt := k6common.GetRuntime(m.ctx)
-	defaultReferer := m.page.mainFrameSession.getNetworkManager().extraHTTPHeaders["referer"]
+	netMgr := m.page.mainFrameSession.getNetworkManager()
+	defaultReferer := netMgr.extraHTTPHeaders["referer"]
 	parsedOpts := NewFrameGotoOptions(defaultReferer, time.Duration(m.timeoutSettings.navigationTimeout())*time.Second)
 	if err := parsedOpts.Parse(m.ctx, opts); err != nil {
 		k6common.Throw(rt, fmt.Errorf("failed parsing options: %w", err))
@@ -627,10 +629,16 @@ func (m *FrameManager) NavigateFrame(frame *Frame, url string, opts goja.Value)
 		if err != nil {
 			k6common.Throw(rt, err)
 		}
+
 		event = data.(*NavigationEvent)
 		if event.newDocument.documentID != newDocumentID {
-			k6common.Throw(rt, errors.New("navigation interrupted by another one"))
-		} else if event.err != nil {
+			m.logger.Debugf("FrameManager:NavigateFrame:interrupted",
+				"fmid:%d fid:%v furl:%s url:%s docID:%s newDocID:%s",
+				fmid, fid, furl, url, event.newDocument.documentID, newDocumentID)
+		} else if event.err != nil &&
+			// TODO: A more graceful way of avoiding Throw()?
+			!(netMgr.userReqInterceptionEnabled &&
+				strings.Contains(event.err.Error(), "ERR_BLOCKED_BY_CLIENT")) {
 			k6common.Throw(rt, event.err)
 		}
 	} else {

common/frame_session.go

@@ -112,26 +112,15 @@ func NewFrameSession(
 		},
 	}
 
+	var parentNM *NetworkManager
 	if fs.parent != nil {
-		fs.networkManager, err = NewNetworkManager(ctx, session, fs.manager, fs.parent.networkManager)
-		if err != nil {
-			logger.Debugf(
-				"NewFrameSession:NewNetworkManager",
-				"sid:%v tid:%v err:%v",
-				session.id, targetID, err)
-
-			return nil, err
-		}
-	} else {
-		fs.networkManager, err = NewNetworkManager(ctx, session, fs.manager, nil)
-		if err != nil {
-			logger.Debugf(
-				"NewFrameSession:NewNetworkManager#2",
-				"sid:%v tid:%v err:%v",
-				session.id, targetID, err)
-
-			return nil, err
-		}
+		parentNM = fs.parent.networkManager
+	}
+	fs.networkManager, err = NewNetworkManager(ctx, session, fs.manager, parentNM)
+	if err != nil {
+		logger.Debugf("NewFrameSession:NewNetworkManager", "sid:%v tid:%v err:%v",
+			session.id, targetID, err)
+		return nil, err
 	}
 
 	action := browser.GetWindowForTarget().WithTargetID(fs.targetID)
@@ -372,8 +361,11 @@ func (fs *FrameSession) initOptions() error {
 	fs.logger.Debugf("NewFrameSession:initOptions",
 		"sid:%v tid:%v", fs.session.id, fs.targetID)
 
-	opts := fs.manager.page.browserCtx.opts
-	optActions := []Action{}
+	var (
+		opts       = fs.manager.page.browserCtx.opts
+		optActions = []Action{}
+		state      = k6lib.GetState(fs.ctx)
+	)
 
 	if fs.isMainFrame() {
 		optActions = append(optActions, emulation.SetFocusEmulationEnabled(true))
@@ -413,9 +405,15 @@ func (fs *FrameSession) initOptions() error {
 		return err
 	}
 	fs.updateExtraHTTPHeaders(true)
-	if err := fs.updateRequestInterception(true); err != nil {
+
+	var reqIntercept bool
+	if state.Options.BlockedHostnames.Trie != nil {
+		reqIntercept = true
+	}
+	if err := fs.updateRequestInterception(reqIntercept); err != nil {
 		return err
 	}
+
 	fs.updateOffline(true)
 	fs.updateHttpCredentials(true)
 	if err := fs.updateEmulateMedia(true); err != nil {
@@ -990,10 +988,9 @@ func (fs *FrameSession) updateOffline(initial bool) {
 	}
 }
 
-func (fs *FrameSession) updateRequestInterception(initial bool) error {
-	fs.logger.Debugf("NewFrameSession:updateRequestInterception", "sid:%v tid:%v", fs.session.id, fs.targetID)
-
-	return fs.networkManager.setRequestInterception(fs.page.hasRoutes())
+func (fs *FrameSession) updateRequestInterception(enable bool) error {
+	fs.logger.Debugf("NewFrameSession:updateRequestInterception", "sid:%v tid:%v on:%v", fs.session.id, fs.targetID, enable)
+	return fs.networkManager.setRequestInterception(enable || fs.page.hasRoutes())
 }
 
 func (fs *FrameSession) updateViewport() error {

common/network_manager.go

@@ -24,6 +24,7 @@ import (
 	"context"
 	"fmt"
 	"net"
+	"net/url"
 	"strconv"
 	"sync"
 	"time"
@@ -36,6 +37,7 @@ import (
 	"github.com/dop251/goja"
 	k6common "go.k6.io/k6/js/common"
 	k6lib "go.k6.io/k6/lib"
+	k6types "go.k6.io/k6/lib/types"
 	k6stats "go.k6.io/k6/stats"
 )
 
@@ -48,7 +50,7 @@ type NetworkManager struct {
 
 	ctx          context.Context
 	logger       *Logger
-	session      *Session
+	session      session
 	parent       *NetworkManager
 	frameManager *FrameManager
 	credentials  *Credentials
@@ -235,10 +237,20 @@ func (m *NetworkManager) handleRequestRedirect(req *Request, redirectResponse *n
 }
 
 func (m *NetworkManager) initDomains() error {
-	action := network.Enable()
-	if err := action.Do(cdp.WithExecutor(m.ctx, m.session)); err != nil {
-		return fmt.Errorf("unable to execute %T: %w", action, err)
+	actions := []Action{network.Enable()}
+
+	// Only enable the Fetch domain if necessary, as it has a performance overhead.
+	if m.userReqInterceptionEnabled {
+		actions = append(actions,
+			network.SetCacheDisabled(true),
+			fetch.Enable().WithPatterns([]*fetch.RequestPattern{{URLPattern: "*"}}))
 	}
+	for _, action := range actions {
+		if err := action.Do(cdp.WithExecutor(m.ctx, m.session)); err != nil {
+			return fmt.Errorf("unable to execute %T: %w", action, err)
+		}
+	}
+
 	return nil
 }
 
@@ -250,6 +262,7 @@ func (m *NetworkManager) initEvents() {
 		cdproto.EventNetworkRequestWillBeSent,
 		cdproto.EventNetworkRequestServedFromCache,
 		cdproto.EventNetworkResponseReceived,
+		cdproto.EventFetchRequestPaused,
 	}, chHandler)
 
 	go func() {
@@ -279,6 +292,8 @@ func (m *NetworkManager) handleEvents(in <-chan Event) bool {
 			m.onRequestServedFromCache(ev)
 		case *network.EventResponseReceived:
 			m.onResponseReceived(ev)
+		case *fetch.EventRequestPaused:
+			m.onRequestPaused(ev)
 		}
 	}
 	return true
@@ -360,32 +375,58 @@ func (m *NetworkManager) onRequest(event *network.EventRequestWillBeSent, interc
 	m.reqsMu.Unlock()
 	m.emitRequestMetrics(req)
 	m.frameManager.requestStarted(req)
+}
 
-	if m.userReqInterceptionEnabled {
-		state := k6lib.GetState(m.ctx)
-		ip := net.ParseIP(req.url.Host)
-		blockedHosts := state.Options.BlockedHostnames.Trie
-		if blockedHosts != nil && ip == nil {
-			if match, blocked := blockedHosts.Contains(req.url.Host); blocked {
-				// Tell browser we've blocked this request.
-				fetch.FailRequest(fetch.RequestID(req.getID()), network.ErrorReasonBlockedByClient)
-
-				// Throw exception into JS runtime
-				rt := k6common.GetRuntime(m.ctx)
-				// TODO: create PR to make netext.BlockedHostError a public struct in k6 perhaps?
-				k6common.Throw(rt, fmt.Errorf("hostname (%s) is in a blocked pattern (%s)", req.url.Host, match))
-			}
+func (m *NetworkManager) onRequestPaused(event *fetch.EventRequestPaused) {
+	m.logger.Debugf("NetworkManager:onRequestPaused",
+		"sid:%s url:%v", m.session.ID(), event.Request.URL)
+	defer m.logger.Debugf("NetworkManager:onRequestPaused:return",
+		"sid:%s url:%v", m.session.ID(), event.Request.URL)
+
+	var (
+		failReason string
+		state      = k6lib.GetState(m.ctx)
+	)
+
+	defer func() { m.failOrContinueRequest(event, failReason) }()
+
+	purl, err := url.Parse(event.Request.URL)
+	if err != nil {
+		m.logger.Errorf("NetworkManager:onRequestPaused",
+			"error parsing URL: %s", err.Error())
+		return
+	}
+
+	failReason = handleBlockedHosts(purl, state.Options.BlockedHostnames.Trie)
+}
+
+func (m *NetworkManager) failOrContinueRequest(event *fetch.EventRequestPaused, failReason string) {
+	if failReason != "" {
+		action := fetch.FailRequest(event.RequestID, network.ErrorReasonBlockedByClient)
+		if err := action.Do(cdp.WithExecutor(m.ctx, m.session)); err != nil {
+			m.logger.Errorf("NetworkManager:onRequestPaused",
+				"error interrupting request: %s", err.Error())
+		} else {
+			m.logger.Warnf("NetworkManager:onRequestPaused",
+				"request %s %s was interrupted: %s", event.Request.Method, event.Request.URL, failReason)
+			return
 		}
+	}
+	action := fetch.ContinueRequest(event.RequestID)
+	if err := action.Do(cdp.WithExecutor(m.ctx, m.session)); err != nil {
+		m.logger.Errorf("NetworkManager:onRequestPaused",
+			"error continuing request: %s", err.Error())
+	}
+}
 
-		/*
-			TODO: is there a way to do IP filtering without requiring a lookup here?
-			for _, ipnet := range state.Options.BlacklistIPs {
-				if ipnet.Contains(ev.Request.URL) {
-					return "", netext.BlackListedIPError{ip: remote.IP, net: ipnet}
-				}
-			}
-		*/
+func handleBlockedHosts(u *url.URL, blockedHosts *k6types.HostnameTrie) string {
+	ip := net.ParseIP(u.Host)
+	if ip == nil {
+		if match, blocked := blockedHosts.Contains(u.Host); blocked {
+			return fmt.Sprintf("hostname %s is in a blocked pattern (%s)", u.Host, match)
+		}
 	}
+	return ""
 }
 
 func (m *NetworkManager) onRequestServedFromCache(event *network.EventRequestServedFromCache) {