parent b45f889

author Tim Voronov <[email protected]> 1552622815 -0400 committer Владимир Фетисов <[email protected]> 1554919032 +0300 parent b45f889 author Tim Voronov <[email protected]> 1552622815 -0400 committer Владимир Фетисов <[email protected]> 1554919029 +0300 Updated style template Feature/MontFerret#236 cookies (MontFerret#242) * Added KeepCookies option to CDP driver * Added LoadDocumentParams * Added COOKIE_GET and COOKIE_SET methods Updated CHANGELOG Updated nav example fix order of arguments (MontFerret#269) Updated RAND function (MontFerret#271) * Updated RAND function * renamed a func Linter Cleanup (MontFerret#276) * linter cleanup * fix default case Bump github.com/mafredri/cdp from 0.22.0 to 0.23.0 (MontFerret#278) Bumps [github.com/mafredri/cdp](https://github.com/mafredri/cdp) from 0.22.0 to 0.23.0. - [Release notes](https://github.com/mafredri/cdp/releases) - [Commits](mafredri/cdp@v0.22.0...v0.23.0) Signed-off-by: dependabot[bot] <[email protected]> fix diff
3timeslazy · Apr 10, 2019 · f916f65 · f916f65
1 parent b45f889
commit f916f65
Show file tree

Hide file tree

Showing 10 changed files with 293 additions and 282 deletions.
diff --git a/cli/repl.go b/cli/repl.go
@@ -9,8 +9,6 @@ import (
 
 	"github.com/MontFerret/ferret/pkg/parser/fql"
 
-	"github.com/MontFerret/ferret/pkg/parser/fql"
-
 	"github.com/MontFerret/ferret/pkg/compiler"
 	"github.com/MontFerret/ferret/pkg/runtime"
 	"github.com/MontFerret/ferret/pkg/runtime/logging"

diff --git a/pkg/drivers/cdp/document.go b/pkg/drivers/cdp/document.go
@@ -2,6 +2,7 @@ package cdp
 
 import (
 	"context"
+	"encoding/json"
 	"fmt"
 	"hash/fnv"
 	"sync"
@@ -19,6 +20,7 @@ import (
 	"github.com/mafredri/cdp"
 	"github.com/mafredri/cdp/protocol/dom"
 	"github.com/mafredri/cdp/protocol/input"
+	"github.com/mafredri/cdp/protocol/network"
 	"github.com/mafredri/cdp/protocol/page"
 	"github.com/mafredri/cdp/rpcc"
 	"github.com/pkg/errors"
@@ -49,21 +51,69 @@ func LoadHTMLDocument(
 	ctx context.Context,
 	conn *rpcc.Conn,
 	client *cdp.Client,
-	url string,
+	params drivers.LoadDocumentParams,
 ) (drivers.HTMLDocument, error) {
 	logger := logging.FromContext(ctx)
 
 	if conn == nil {
 		return nil, core.Error(core.ErrMissedArgument, "connection")
 	}
 
-	if url == "" {
+	if params.URL == "" {
 		return nil, core.Error(core.ErrMissedArgument, "url")
 	}
 
+	if params.Cookies != nil {
+		cookies := make([]network.CookieParam, 0, len(params.Cookies))
+
+		for _, c := range params.Cookies {
+			cookies = append(cookies, fromDriverCookie(params.URL, c))
+
+			logger.
+				Debug().
+				Timestamp().
+				Str("cookie", c.Name).
+				Msg("set cookie")
+		}
+
+		err := client.Network.SetCookies(
+			ctx,
+			network.NewSetCookiesArgs(cookies),
+		)
+
+		if err != nil {
+			return nil, err
+		}
+	}
+
+	if params.Header != nil {
+		j, err := json.Marshal(params.Header)
+
+		if err != nil {
+			return nil, err
+		}
+
+		for k := range params.Header {
+			logger.
+				Debug().
+				Timestamp().
+				Str("header", k).
+				Msg("set header")
+		}
+
+		err = client.Network.SetExtraHTTPHeaders(
+			ctx,
+			network.NewSetExtraHTTPHeadersArgs(network.Headers(j)),
+		)
+
+		if err != nil {
+			return nil, err
+		}
+	}
+
 	var err error
 
-	if url != BlankPageURL {
+	if params.URL != BlankPageURL {
 		err = waitForLoadEvent(ctx, client)
 
 		if err != nil {
@@ -109,7 +159,7 @@ func LoadHTMLDocument(
 		conn,
 		client,
 		broker,
-		values.NewString(url),
+		values.NewString(params.URL),
 		rootElement,
 	), nil
 }
@@ -316,6 +366,67 @@ func (doc *HTMLDocument) GetURL() core.Value {
 	return doc.url
 }
 
+func (doc *HTMLDocument) GetCookies(ctx context.Context) (*values.Array, error) {
+	doc.Lock()
+	defer doc.Unlock()
+
+	repl, err := doc.client.Network.GetAllCookies(ctx)
+
+	if err != nil {
+		return values.NewArray(0), err
+	}
+
+	if repl.Cookies == nil {
+		return values.NewArray(0), nil
+	}
+
+	cookies := values.NewArray(len(repl.Cookies))
+
+	for _, c := range repl.Cookies {
+		cookies.Push(toDriverCookie(c))
+	}
+
+	return cookies, nil
+}
+
+func (doc *HTMLDocument) SetCookies(ctx context.Context, cookies ...drivers.HTTPCookie) error {
+	doc.Lock()
+	defer doc.Unlock()
+
+	if len(cookies) == 0 {
+		return nil
+	}
+
+	params := make([]network.CookieParam, 0, len(cookies))
+
+	for _, c := range cookies {
+		params = append(params, fromDriverCookie(doc.url.String(), c))
+	}
+
+	return doc.client.Network.SetCookies(ctx, network.NewSetCookiesArgs(params))
+}
+
+func (doc *HTMLDocument) DeleteCookies(ctx context.Context, cookies ...drivers.HTTPCookie) error {
+	doc.Lock()
+	defer doc.Unlock()
+
+	if len(cookies) == 0 {
+		return nil
+	}
+
+	var err error
+
+	for _, c := range cookies {
+		err = doc.client.Network.DeleteCookies(ctx, fromDriverCookieDelete(doc.url.String(), c))
+
+		if err != nil {
+			break
+		}
+	}
+
+	return err
+}
+
 func (doc *HTMLDocument) SetURL(ctx context.Context, url values.String) error {
 	return doc.Navigate(ctx, url)
 }
@@ -852,6 +963,10 @@ func (doc *HTMLDocument) PrintToPDF(ctx context.Context, params drivers.PDFParam
 func (doc *HTMLDocument) CaptureScreenshot(ctx context.Context, params drivers.ScreenshotParams) (values.Binary, error) {
 	metrics, err := doc.client.Page.GetLayoutMetrics(ctx)
 
+	if err != nil {
+		return values.NewBinary(nil), err
+	}
+
 	if params.Format == drivers.ScreenshotFormatJPEG && params.Quality < 0 && params.Quality > 100 {
 		params.Quality = 100
 	}

diff --git a/pkg/drivers/cdp/events/wait.go b/pkg/drivers/cdp/events/wait.go
@@ -1,19 +1,21 @@
 package events
 
 import (
+	"context"
 	"time"
 
+	"github.com/MontFerret/ferret/pkg/drivers"
 	"github.com/MontFerret/ferret/pkg/drivers/cdp/eval"
 	"github.com/MontFerret/ferret/pkg/runtime/core"
 	"github.com/MontFerret/ferret/pkg/runtime/values"
 	"github.com/mafredri/cdp"
 )
 
 type (
-	Function func() (core.Value, error)
+	Function func(ctx context.Context) (core.Value, error)
+
 	WaitTask struct {
 		fun     Function
-		timeout time.Duration
 		polling time.Duration
 	}
 )
@@ -22,39 +24,31 @@ const DefaultPolling = time.Millisecond * time.Duration(200)
 
 func NewWaitTask(
 	fun Function,
-	timeout time.Duration,
 	polling time.Duration,
 ) *WaitTask {
 	return &WaitTask{
 		fun,
-		timeout,
 		polling,
 	}
 }
 
-func (task *WaitTask) Run() (core.Value, error) {
-	timer := time.NewTimer(task.timeout)
-
+func (task *WaitTask) Run(ctx context.Context) (core.Value, error) {
 	for {
 		select {
-		case <-timer.C:
+		case <-ctx.Done():
 			return values.None, core.ErrTimeout
 		default:
-			out, err := task.fun()
+			out, err := task.fun(ctx)
 
 			// expression failed
 			// terminating
 			if err != nil {
-				timer.Stop()
-
 				return values.None, err
 			}
 
 			// output is not empty
 			// terminating
 			if out != values.None {
-				timer.Stop()
-
 				return out, nil
 			}
 
@@ -67,19 +61,54 @@ func (task *WaitTask) Run() (core.Value, error) {
 func NewEvalWaitTask(
 	client *cdp.Client,
 	predicate string,
-	timeout time.Duration,
 	polling time.Duration,
 ) *WaitTask {
 	return NewWaitTask(
-		func() (core.Value, error) {
+		func(ctx context.Context) (core.Value, error) {
 			return eval.Eval(
+				ctx,
 				client,
 				predicate,
 				true,
 				false,
 			)
 		},
-		timeout,
 		polling,
 	)
 }
+
+func NewValueWaitTask(
+	when drivers.WaitEvent,
+	value core.Value,
+	getter Function,
+	polling time.Duration,
+) *WaitTask {
+	return &WaitTask{
+		func(ctx context.Context) (core.Value, error) {
+			current, err := getter(ctx)
+
+			if err != nil {
+				return values.None, err
+			}
+
+			if when == drivers.WaitEventPresence {
+				// Values appeared, exit
+				if current.Compare(value) == 0 {
+					// The value does not really matter if it's not None
+					// None indicates that operation needs to be repeated
+					return values.True, nil
+				}
+			} else {
+				// Value disappeared, exit
+				if current.Compare(value) != 0 {
+					// The value does not really matter if it's not None
+					// None indicates that operation needs to be repeated
+					return values.True, nil
+				}
+			}
+
+			return values.None, nil
+		},
+		polling,
+	}
+}
diff --git a/pkg/drivers/cdp/options.go b/pkg/drivers/cdp/options.go
@@ -2,9 +2,11 @@ package cdp
 
 type (
 	Options struct {
-		proxy     string
-		userAgent string
-		address   string
+		Name        string
+		Proxy       string
+		UserAgent   string
+		Address     string
+		KeepCookies bool
 	}
 
 	Option func(opts *Options)
@@ -14,7 +16,8 @@ const DefaultAddress = "http://127.0.0.1:9222"
 
 func newOptions(setters []Option) *Options {
 	opts := new(Options)
-	opts.address = DefaultAddress
+	opts.Name = DriverName
+	opts.Address = DefaultAddress
 
 	for _, setter := range setters {
 		setter(opts)
@@ -26,19 +29,31 @@ func newOptions(setters []Option) *Options {
 func WithAddress(address string) Option {
 	return func(opts *Options) {
 		if address != "" {
-			opts.address = address
+			opts.Address = address
 		}
 	}
 }
 
 func WithProxy(address string) Option {
 	return func(opts *Options) {
-		opts.proxy = address
+		opts.Proxy = address
 	}
 }
 
 func WithUserAgent(value string) Option {
 	return func(opts *Options) {
-		opts.userAgent = value
+		opts.UserAgent = value
+	}
+}
+
+func WithKeepCookies() Option {
+	return func(opts *Options) {
+		opts.KeepCookies = true
+	}
+}
+
+func WithCustomName(name string) Option {
+	return func(opts *Options) {
+		opts.Name = name
 	}
 }