add exposefunc feature

expose.go

@@ -0,0 +1,130 @@
+package chromedp
+
+import (
+ "context"
+ "encoding/json"
+ "fmt"
+
+ "github.com/chromedp/cdproto/page"
+ "github.com/chromedp/cdproto/runtime"
+)
+
+// ExposedFunc is the function type that can be exposed to the browser env.
+type ExposedFunc func(args string) (string, error)
+
+// ExposeAction are actions which expose Go functions to the browser env.
+type ExposeAction Action
+
+// Expose is an action to add a function called fnName on the browser page's
+// window object. When called, the function executes fn in the Go env and
+// returns a Promise which resolves to the return value of fn.
+//
+// Note:
+// 1. This is the lite version of puppeteer's [page.exposeFunction].
+// 2. It adds "chromedpExposeFunc" to the page's window object too.
+// 3. The exposed function survives page navigation until the tab is closed.
+// 4. It exports the function to all frames on the current page.
+// 5. Avoid exposing multiple funcs with the same name.
+// 6. Maybe you just need runtime.AddBinding.
+//
+// [page.exposeFunction]: https://github.com/puppeteer/puppeteer/blob/v19.2.2/docs/api/puppeteer.page.exposefunction.md
+func Expose(fnName string, fn ExposedFunc) ExposeAction {
+ return ActionFunc(func(ctx context.Context) error {
+
+ expression := fmt.Sprintf(`chromedpExposeFunc.wrapBinding("exposedFun","%s");`, fnName)
+
+ err := Run(ctx,
+ runtime.AddBinding(fnName),
+ evaluateOnAllFrames(exposeJS),
+ evaluateOnAllFrames(expression),
+ // Make it effective after navigation.
+ addScriptToEvaluateOnNewDocument(exposeJS),
+ addScriptToEvaluateOnNewDocument(expression),
+ )
+ if err != nil {
+ return err
+ }
+
+ ListenTarget(ctx, func(ev interface{}) {
+ switch ev := ev.(type) {
+ case *runtime.EventBindingCalled:
+ if ev.Payload == "" {
+ return
+ }
+
+ var payload struct {
+ Type string `json:"type"`
+ Name string `json:"name"`
+ Seq int64 `json:"seq"`
+ Args string `json:"args"`
+ }
+
+ err := json.Unmarshal([]byte(ev.Payload), &payload)
+ if err != nil {
+ return
+ }
+
+ if payload.Type != "exposedFun" || payload.Name != fnName {
+ return
+ }
+
+ result, err := fn(payload.Args)
+
+ callback := "chromedpExposeFunc.deliverResult"
+ if err != nil {
+ result = err.Error()
+ callback = "chromedpExposeFunc.deliverError"
+ }
+
+ // Prevent the message from being processed by other functions
+ ev.Payload = ""
+
+ go func() {
+ err := Run(ctx,
+ CallFunctionOn(callback,
+ nil,
+ func(p *runtime.CallFunctionOnParams) *runtime.CallFunctionOnParams {
+ return p.WithExecutionContextID(ev.ExecutionContextID)
+ },
+ payload.Name,
+ payload.Seq,
+ result,
+ ),
+ )
+
+ if err != nil {
+ c := FromContext(ctx)
+ c.Browser.errf("failed to deliver result to exposed func %s: %s", fnName, err)
+ }
+ }()
+ }
+ })
+
+ return nil
+ })
+}
+
+func addScriptToEvaluateOnNewDocument(script string) Action {
+ return ActionFunc(func(ctx context.Context) error {
+ _, err := page.AddScriptToEvaluateOnNewDocument(script).Do(ctx)
+ return err
+ })
+}
+
+func evaluateOnAllFrames(script string) Action {
+ return ActionFunc(func(ctx context.Context) error {
+ c := FromContext(ctx)
+
+ c.Target.frameMu.RLock()
+ actions := make([]Action, 0, len(c.Target.execContexts))
+ for _, executionContextID := range c.Target.execContexts {
+ id := executionContextID
+ actions = append(actions, Evaluate(script, nil, func(p *runtime.EvaluateParams) *runtime.EvaluateParams {
+ return p.WithContextID(id)
+ }))
+ }
+ c.Target.frameMu.RUnlock()
+
+ return Tasks(actions).Do(ctx)
+ })
+}

expose_test.go

@@ -0,0 +1,178 @@
+package chromedp
+
+import (
+ "crypto/md5"
+ "encoding/base64"
+ "encoding/hex"
+ "fmt"
+ "testing"
+
+ "github.com/chromedp/cdproto/runtime"
+)
+
+func md5SumFunc(args string) (string, error) {
+ h := md5.New()
+ h.Write([]byte(args))
+ return hex.EncodeToString(h.Sum(nil)), nil
+}
+
+func base64EncodeFunc(args string) (string, error) {
+ return base64.StdEncoding.EncodeToString([]byte(testString)), nil
+}
+
+func echoFunc(str string) (string, error) {
+ return str, nil
+}
+
+const testString = "chromedp expose test"
+const testStringMd5 = "a93d69002a286b46c8aa114362afb7ac"
+const testStringBase64 = "Y2hyb21lZHAgZXhwb3NlIHRlc3Q="
+const testIFrameHTMLTitle = "page with an iframe"
+const testFormHTMLTitle = "this is form title"
+
+func TestExposeToAllFrames(t *testing.T) {
+ // allocate browser
+ ctx, cancel := testAllocate(t, "iframe.html")
+ defer cancel()
+
+ // expose echoFunc function as to browser current page and every frame
+ if err := Run(ctx, Expose("echo", echoFunc)); err != nil {
+ t.Fatal(err)
+ }
+
+ c := FromContext(ctx)
+
+ c.Target.frameMu.RLock()
+ executionContextIDs := make([]runtime.ExecutionContextID, 0, len(c.Target.execContexts))
+ for _, executionContextID := range c.Target.execContexts {
+ executionContextIDs = append(executionContextIDs, executionContextID)
+ }
+ c.Target.frameMu.RUnlock()
+
+ var res1 string
+ var res2 string
+ callEchoFunc := fmt.Sprintf(`%s(document.title);`, "echo")
+ for _, executionContextID := range executionContextIDs {
+ id := executionContextID
+ var res string
+ if err := Run(ctx, Evaluate(callEchoFunc, &res, func(p *runtime.EvaluateParams) *runtime.EvaluateParams {
+ return p.WithContextID(id).WithAwaitPromise(true)
+ })); err != nil {
+ t.Fatal(err)
+ }
+ if len(res1) == 0 {
+ res1 = res
+ } else {
+ res2 = res
+ }
+ }
+
+ // we expect res1 or res2 = testIFrameHTMLTitle or testFormHTMLTitle
+ if res1 == testIFrameHTMLTitle && res2 == testFormHTMLTitle || res1 == testFormHTMLTitle && res2 == testIFrameHTMLTitle {
+ // pass
+ } else {
+ t.Fatalf("res1: %s, res2: %s", res1, res2)
+ }
+}
+
+func TestExpose(t *testing.T) {
+ // allocate browser
+ ctx, cancel := testAllocate(t, "")
+ defer cancel()
+
+ // expose md5SumFunc function as md5 to browser current page and every frame
+ if err := Run(ctx, Expose("md5", md5SumFunc)); err != nil {
+ t.Fatal(err)
+ }
+
+ // expose base64EncodeFunc function as base64 to browser current page and every frame
+ if err := Run(ctx, Expose("base64", base64EncodeFunc)); err != nil {
+ t.Fatal(err)
+ }
+
+ // 1. When on the current page
+ var res string
+ callMd5 := fmt.Sprintf(`%s("%s");`, "md5", testString)
+ if err := Run(ctx, Evaluate(callMd5, &res, func(p *runtime.EvaluateParams) *runtime.EvaluateParams {
+ return p.WithAwaitPromise(true)
+ })); err != nil {
+ t.Fatal(err)
+ }
+
+ if res != testStringMd5 {
+ t.Fatalf("want: %s, got: %s", testStringMd5, res)
+ }
+
+ var res2 string
+ callBase64 := fmt.Sprintf(`%s("%s");`, "base64", testString)
+ if err := Run(ctx, Evaluate(callBase64, &res2, func(p *runtime.EvaluateParams) *runtime.EvaluateParams {
+ return p.WithAwaitPromise(true)
+ })); err != nil {
+ t.Fatal(err)
+ }
+ if res2 != testStringBase64 {
+ t.Fatalf("want: %s, got: %s", testStringBase64, res)
+ }
+
+ // 2. Navigate another page
+ if err := Run(ctx,
+ Navigate(testdataDir+"/child1.html"),
+ ); err != nil {
+ t.Fatal(err)
+ }
+
+ // we expect md5 can work properly.
+ if err := Run(ctx, Evaluate(callMd5, &res, func(p *runtime.EvaluateParams) *runtime.EvaluateParams {
+ return p.WithAwaitPromise(true)
+ })); err != nil {
+ t.Fatal(err)
+ }
+ if res != testStringMd5 {
+ t.Fatalf("want: %s, got: %s", testStringMd5, res)
+ }
+
+ // we expect base64 can work properly.
+ if err := Run(ctx, Evaluate(callBase64, &res2, func(p *runtime.EvaluateParams) *runtime.EvaluateParams {
+ return p.WithAwaitPromise(true)
+ })); err != nil {
+ t.Fatal(err)
+ }
+ if res2 != testStringBase64 {
+ t.Fatalf("want: %s, got: %s", testStringBase64, res)
+ }
+}
+
+func TestExposeMulti(t *testing.T) {
+ // allocate browser
+ ctx, cancel := testAllocate(t, "")
+ defer cancel()
+
+ // creates a new page. about:blank
+ Run(ctx)
+
+ // expose md5SumFunc function as sameFunc to browser current page and every frame
+ if err := Run(ctx, Expose("sameFunc", md5SumFunc)); err != nil {
+ t.Fatal(err)
+ }
+
+ // expose base64EncodeFunc function as sameFunc to browser current page and every frame
+ if err := Run(ctx, Expose("sameFunc", base64EncodeFunc)); err != nil {
+ t.Fatal(err)
+ }
+
+ // we expect first expose function to handle
+ var res string
+ sameFunc := fmt.Sprintf(`%s("%s");`, "sameFunc", testString)
+ if err := Run(ctx, Evaluate(sameFunc, &res, func(p *runtime.EvaluateParams) *runtime.EvaluateParams {
+ return p.WithAwaitPromise(true)
+ })); err != nil {
+ t.Fatal(err)
+ }
+
+ if res != testStringMd5 {
+ t.Fatalf("want md5SumFunc res:%s, got:%s", testStringMd5, res)
+ }
+ if res == testStringBase64 {
+ t.Fatalf("want md5SumFunc res:%s, got base64EncodeFunc res :%s", testStringMd5, res)
+ }
+}

js.go

@@ -56,4 +56,11 @@ var (
  // It's modified to make mutation polling respect timeout even when there is not a DOM mutation.
  //go:embed js/waitForPredicatePageFunction.js
  waitForPredicatePageFunction string
+
+ // exposedJS is a javascript snippet that wraps the function (CDP binding)
+ // It refers to puppeteer. See
+ // https://github.com/puppeteer/puppeteer/blob/main/packages/puppeteer-core/src/common/util.ts#L248-L327
+ // It's modified to make BindingFunc takes exactly one argument, this argument should be string
+ //go:embed js/expose.js
+ exposeJS string
 )

js/expose.js

@@ -0,0 +1,43 @@
+var chromedpExposeFunc = chromedpExposeFunc || {
+ bindings: {},
+ deliverError: function (name, seq, message) {
+ const error = new Error(message);
+ chromedpExposeFunc.bindings[name].callbacks.get(seq).reject(error);
+ chromedpExposeFunc.bindings[name].callbacks.delete(seq);
+ },
+ deliverResult: function (name, seq, result) {
+ chromedpExposeFunc.bindings[name].callbacks.get(seq).resolve(result);
+ chromedpExposeFunc.bindings[name].callbacks.delete(seq);
+ },
+ wrapBinding: function (type, name) {
+ // Store the binding function added by the call of runtime.AddBinding.
+ chromedpExposeFunc.bindings[name] = window[name];
+
+ // Replace the binding function.
+ Object.assign(window, {
+ [name](args) {
+ if (typeof args != 'string') {
+ return Promise.reject(
+ new Error(
+ 'function takes exactly one argument, this argument should be string'
+ )
+ );
+ }
+
+ const binding = chromedpExposeFunc.bindings[name];
+
+ binding.callbacks ??= new Map();
+
+ const seq = (binding.lastSeq ?? 0) + 1;
+ binding.lastSeq = seq;
+
+ // Call the binding function to trigger runtime.EventBindingCalled.
+ binding(JSON.stringify({ type, name, seq, args }));
+
+ return new Promise((resolve, reject) => {
+ binding.callbacks.set(seq, { resolve, reject });
+ });
+ },
+ });
+ },
+};