Bugfix/#597 headers panic (#598)

* Remodeled HTTPHeaders * Remodeled HTTPCookies
MontFerret · Mar 26, 2021 · d55bce3 · d55bce3
1 parent b3adea3
commit d55bce3
Show file tree

Hide file tree

Showing 34 changed files with 722 additions and 343 deletions.
diff --git a/e2e/tests/dynamic/page/cookies/delete.fqlx → e2e/tests/dynamic/page/cookies/delete.fql b/e2e/tests/dynamic/page/cookies/delete.fqlx → e2e/tests/dynamic/page/cookies/delete.fql
@@ -15,7 +15,7 @@ COOKIE_DEL(doc, COOKIE_GET(doc, "x-e2e"), "x-e2e-2")
 LET cookie1 = COOKIE_GET(doc, "x-e2e")
 LET cookie2 = COOKIE_GET(doc, "x-e2e-2")
 
-T::EQ(cookie1, "none")
-T::EQ(cookie2, "none")
+T::EQ(cookie1, NONE)
+T::EQ(cookie2, NONE)
 
 RETURN NONE
diff --git a/e2e/tests/dynamic/page/cookies/get.fqlx b/e2e/tests/dynamic/page/cookies/get.fqlx
@@ -7,6 +7,6 @@ LET cookiesPath = LENGTH(doc.cookies) > 0 ? "ok" : "false"
 LET cookie = COOKIE_GET(doc, "x-ferret")
 LET expected = "ok e2e"
 
-T::LEN(doc.cookies
+T::LEN(doc.cookies, 1)
 
 RETURN T::EQ(cookiesPath + " " + cookie.value, expected)
diff --git a/...tests/dynamic/page/cookies/load_with.fqlx → e2e/tests/dynamic/page/cookies/load_with.fql b/...tests/dynamic/page/cookies/load_with.fqlx → e2e/tests/dynamic/page/cookies/load_with.fql
@@ -7,8 +7,9 @@ LET doc = DOCUMENT(url, {
     }]
 })
 
-LET cookiesPath = LENGTH(doc.cookies) > 0 ? "ok" : "false"
 LET cookie = COOKIE_GET(doc, "x-e2e")
-LET expected = "ok test"
 
-RETURN T::EQ(cookiesPath + " " + cookie.value, expected)
+T::NOT::NONE(cookie)
+T::EQ(cookie.value, "test")
+
+RETURN NONE
diff --git a/e2e/tests/dynamic/page/cookies/set.fqlx → e2e/tests/dynamic/page/cookies/set.fql b/e2e/tests/dynamic/page/cookies/set.fqlx → e2e/tests/dynamic/page/cookies/set.fql
diff --git a/examples/headers.fql b/examples/headers.fql
@@ -0,0 +1,5 @@
+LET proxy_header = {"Proxy-Authorization": ["Basic e40b7d5eff464a4fb51efed2d1a19a24"]}
+
+LET doc = DOCUMENT("https://google.com", { headers: proxy_header})
+
+RETURN doc
diff --git a/pkg/drivers/cdp/driver.go b/pkg/drivers/cdp/driver.go
@@ -36,7 +36,7 @@ type Driver struct {
 
 func NewDriver(opts ...Option) *Driver {
 	drv := new(Driver)
-	drv.options = newOptions(opts)
+	drv.options = NewOptions(opts)
 	drv.dev = devtool.New(drv.options.Address)
 
 	return drv
@@ -137,43 +137,11 @@ func (drv *Driver) createConnection(ctx context.Context, keepCookies bool) (*rpc
 }
 
 func (drv *Driver) setDefaultParams(params drivers.Params) drivers.Params {
-	if params.UserAgent == "" {
-		params.UserAgent = drv.options.UserAgent
-	}
-
 	if params.Viewport == nil {
 		params.Viewport = defaultViewport
 	}
 
-	if drv.options.Headers != nil && params.Headers == nil {
-		params.Headers = make(drivers.HTTPHeaders)
-	}
-
-	// set default headers
-	for k, v := range drv.options.Headers {
-		_, exists := params.Headers[k]
-
-		// do not override user's set values
-		if !exists {
-			params.Headers[k] = v
-		}
-	}
-
-	if drv.options.Cookies != nil && params.Cookies == nil {
-		params.Cookies = make(drivers.HTTPCookies)
-	}
-
-	// set default cookies
-	for k, v := range drv.options.Cookies {
-		_, exists := params.Cookies[k]
-
-		// do not override user's set values
-		if !exists {
-			params.Cookies[k] = v
-		}
-	}
-
-	return params
+	return drivers.SetDefaultParams(drv.options.Options, params)
 }
 
 func (drv *Driver) init(ctx context.Context) error {

diff --git a/pkg/drivers/cdp/helpers.go b/pkg/drivers/cdp/helpers.go
@@ -51,12 +51,6 @@ func enableFeatures(ctx context.Context, client *cdp.Client, params drivers.Para
 		func() error {
 			ua := common.GetUserAgent(params.UserAgent)
 
-			//logger.
-			//	Debug().
-			//	Timestamp().
-			//	Str("user-agent", ua).
-			//	Msg("using User-Agent")
-
 			// do not use custom user agent
 			if ua == "" {
 				return nil

diff --git a/pkg/drivers/cdp/network/manager.go b/pkg/drivers/cdp/network/manager.go
@@ -34,7 +34,7 @@ type (
 		mu                 sync.Mutex
 		logger             *zerolog.Logger
 		client             *cdp.Client
-		headers            drivers.HTTPHeaders
+		headers            *drivers.HTTPHeaders
 		eventLoop          *events.Loop
 		cancel             context.CancelFunc
 		responseListenerID events.ListenerID
@@ -53,20 +53,20 @@ func New(
 	m := new(Manager)
 	m.logger = logger
 	m.client = client
-	m.headers = make(drivers.HTTPHeaders)
+	m.headers = drivers.NewHTTPHeaders()
 	m.eventLoop = events.NewLoop()
 	m.cancel = cancel
 	m.response = new(sync.Map)
 
-	if len(options.Cookies) > 0 {
+	if options.Cookies != nil && len(options.Cookies) > 0 {
 		for url, cookies := range options.Cookies {
 			if err := m.setCookiesInternal(ctx, url, cookies); err != nil {
 				return nil, err
 			}
 		}
 	}
 
-	if len(options.Headers) > 0 {
+	if options.Headers != nil && options.Headers.Length() > 0 {
 		if err := m.setHeadersInternal(ctx, options.Headers); err != nil {
 			return nil, err
 		}
@@ -104,7 +104,7 @@ func New(
 
 	m.responseListenerID = m.eventLoop.AddListener(responseReceived, m.onResponse)
 
-	if len(options.Filter.Patterns) > 0 {
+	if options.Filter != nil && len(options.Filter.Patterns) > 0 {
 		el2 := events.NewLoop()
 
 		err = m.client.Fetch.Enable(ctx, toFetchArgs(options.Filter.Patterns))
@@ -147,87 +147,100 @@ func (m *Manager) Close() error {
 	return nil
 }
 
-func (m *Manager) GetCookies(ctx context.Context) (drivers.HTTPCookies, error) {
+func (m *Manager) GetCookies(ctx context.Context) (*drivers.HTTPCookies, error) {
 	repl, err := m.client.Network.GetAllCookies(ctx)
 
 	if err != nil {
 		return nil, errors.Wrap(err, "failed to get cookies")
 	}
 
-	cookies := make(drivers.HTTPCookies)
+	cookies := drivers.NewHTTPCookies()
 
 	if repl.Cookies == nil {
 		return cookies, nil
 	}
 
 	for _, c := range repl.Cookies {
-		cookies[c.Name] = toDriverCookie(c)
+		cookies.Set(toDriverCookie(c))
 	}
 
 	return cookies, nil
 }
 
-func (m *Manager) SetCookies(ctx context.Context, url string, cookies drivers.HTTPCookies) error {
+func (m *Manager) SetCookies(ctx context.Context, url string, cookies *drivers.HTTPCookies) error {
 	m.mu.Lock()
 	defer m.mu.Unlock()
 
 	return m.setCookiesInternal(ctx, url, cookies)
 }
 
-func (m *Manager) setCookiesInternal(ctx context.Context, url string, cookies drivers.HTTPCookies) error {
-	if len(cookies) == 0 {
+func (m *Manager) setCookiesInternal(ctx context.Context, url string, cookies *drivers.HTTPCookies) error {
+	if cookies == nil {
+		return errors.Wrap(core.ErrMissedArgument, "cookies")
+	}
+
+	if cookies.Length() == 0 {
 		return nil
 	}
 
-	params := make([]network.CookieParam, 0, len(cookies))
+	params := make([]network.CookieParam, 0, cookies.Length())
 
-	for _, c := range cookies {
-		params = append(params, fromDriverCookie(url, c))
-	}
+	cookies.ForEach(func(value drivers.HTTPCookie, _ values.String) bool {
+		params = append(params, fromDriverCookie(url, value))
+
+		return true
+	})
 
 	return m.client.Network.SetCookies(ctx, network.NewSetCookiesArgs(params))
 }
 
-func (m *Manager) DeleteCookies(ctx context.Context, url string, cookies drivers.HTTPCookies) error {
+func (m *Manager) DeleteCookies(ctx context.Context, url string, cookies *drivers.HTTPCookies) error {
 	m.mu.Lock()
 	defer m.mu.Unlock()
 
-	if len(cookies) == 0 {
+	if cookies == nil {
+		return errors.Wrap(core.ErrMissedArgument, "cookies")
+	}
+
+	if cookies.Length() == 0 {
 		return nil
 	}
 
 	var err error
 
-	for _, c := range cookies {
-		err = m.client.Network.DeleteCookies(ctx, fromDriverCookieDelete(url, c))
+	cookies.ForEach(func(value drivers.HTTPCookie, _ values.String) bool {
+		err = m.client.Network.DeleteCookies(ctx, fromDriverCookieDelete(url, value))
 
 		if err != nil {
-			break
+			return false
 		}
-	}
+
+		return true
+	})
 
 	return err
 }
 
-func (m *Manager) GetHeaders(_ context.Context) (drivers.HTTPHeaders, error) {
-	copied := make(drivers.HTTPHeaders)
+func (m *Manager) GetHeaders(_ context.Context) (*drivers.HTTPHeaders, error) {
+	m.mu.Lock()
+	defer m.mu.Unlock()
 
-	for k, v := range m.headers {
-		copied[k] = v
+	if m.headers == nil {
+		return drivers.NewHTTPHeaders(), nil
 	}
 
-	return copied, nil
+	return m.headers.Clone().(*drivers.HTTPHeaders), nil
 }
 
-func (m *Manager) SetHeaders(ctx context.Context, headers drivers.HTTPHeaders) error {
+func (m *Manager) SetHeaders(ctx context.Context, headers *drivers.HTTPHeaders) error {
 	m.mu.Lock()
 	defer m.mu.Unlock()
 
 	return m.setHeadersInternal(ctx, headers)
 }
 
-func (m *Manager) setHeadersInternal(ctx context.Context, headers drivers.HTTPHeaders) error {
-	if len(headers) == 0 {
+func (m *Manager) setHeadersInternal(ctx context.Context, headers *drivers.HTTPHeaders) error {
+	if headers.Length() == 0 {
 		return nil
 	}
 
@@ -461,7 +474,7 @@ func (m *Manager) onResponse(_ context.Context, message interface{}) (out bool)
 	response := drivers.HTTPResponse{
 		StatusCode: msg.Response.Status,
 		Status:     msg.Response.StatusText,
-		Headers:    make(drivers.HTTPHeaders),
+		Headers:    drivers.NewHTTPHeaders(),
 	}
 
 	deserialized := make(map[string]string)

diff --git a/pkg/drivers/cdp/network/options.go b/pkg/drivers/cdp/network/options.go
@@ -6,16 +6,16 @@ import (
 )
 
 type (
-	Cookies map[string]drivers.HTTPCookies
+	Cookies map[string]*drivers.HTTPCookies
 
 	Filter struct {
 		Patterns []drivers.ResourceFilter
 	}
 
 	Options struct {
 		Cookies Cookies
-		Headers drivers.HTTPHeaders
-		Filter  Filter
+		Headers *drivers.HTTPHeaders
+		Filter  *Filter
 	}
 )