add exposefunc feature

expose.go

@@ -0,0 +1,111 @@
+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. (iframe?)
+// 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),
+			Evaluate(exposeJS, nil),
+			Evaluate(expression, nil),
+			// 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
+	})
+}

expose_test.go

@@ -0,0 +1,127 @@
+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
+}
+
+const testString = "chromedp expose test"
+const testStringMd5 = "a93d69002a286b46c8aa114362afb7ac"
+const testStringBase64 = "Y2hyb21lZHAgZXhwb3NlIHRlc3Q="
+
+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 });
+        });
+      },
+    });
+  },
+};