Support screenshot with headless browser remote allocator

cmd/screenshot/main.go

@@ -16,9 +16,11 @@ import (
 func main() {
 	var timeout uint64
 	var format string
+	var remoteAddr string
 
-	flag.Uint64Var(&timeout, "timeout", 60, "Screenshot timeout")
-	flag.StringVar(&format, "format", "png", "Screenshot file format, default: png")
+	flag.Uint64Var(&timeout, "timeout", 60, "Screenshot timeout.")
+	flag.StringVar(&format, "format", "png", "Screenshot file format.")
+	flag.StringVar(&remoteAddr, "remote-addr", "", "Headless browser remote address, such as 127.0.0.1:9222")
 
 	flag.Parse()
 
@@ -41,7 +43,18 @@ func main() {
 
 	urls := helper.MatchURL(str)
 
-	shots, err := screenshot.Screenshot(ctx, urls, screenshot.ScaleFactor(1), screenshot.Format(format))
+	var err error
+	var shots []screenshot.Screenshots
+	if remoteAddr != "" {
+		remote, err := screenshot.NewChromeRemoteScreenshoter(remoteAddr)
+		if err != nil {
+			fmt.Println(err)
+			return
+		}
+		shots, err = remote.Screenshot(ctx, urls, screenshot.ScaleFactor(1))
+	} else {
+		shots, err = screenshot.Screenshot(ctx, urls, screenshot.ScaleFactor(1), screenshot.Format(format))
+	}
 	if err != nil {
 		if err == context.DeadlineExceeded {
 			fmt.Println(err.Error())

go.mod

@@ -3,8 +3,8 @@ module github.com/wabarc/screenshot
 go 1.16
 
 require (
-	github.com/chromedp/cdproto v0.0.0-20210126020047-7ec7357d1463
-	github.com/chromedp/chromedp v0.6.5
+	github.com/chromedp/cdproto v0.0.0-20210323015217-0942afbea50e
+	github.com/chromedp/chromedp v0.6.10
 	github.com/wabarc/helper v0.0.0-20210127120855-10af37cc2616
-	golang.org/x/sys v0.0.0-20210124154548-22da62e12c0c // indirect
+	golang.org/x/sys v0.0.0-20210415045647-66c3f260301c // indirect
 )

go.sum

@@ -1,8 +1,7 @@
-github.com/chromedp/cdproto v0.0.0-20210122124816-7a656c010d57/go.mod h1:55pim6Ht4LJKdVLlyFJV/g++HsEA1hQxPbB5JyNdZC0=
-github.com/chromedp/cdproto v0.0.0-20210126020047-7ec7357d1463 h1:jLOnhBJ0rDgLdxTpAT6KlbQWcPTdlsxs+9d1RvX9u1w=
-github.com/chromedp/cdproto v0.0.0-20210126020047-7ec7357d1463/go.mod h1:55pim6Ht4LJKdVLlyFJV/g++HsEA1hQxPbB5JyNdZC0=
-github.com/chromedp/chromedp v0.6.5 h1:hPaDYBpvD2WFicln0ByzV+XRhSOtLgAgsu39O455iWY=
-github.com/chromedp/chromedp v0.6.5/go.mod h1:/Q6h52DkrFuvOgmCuR6O3xT5g0bZYoPqjANKBEvQGEY=
+github.com/chromedp/cdproto v0.0.0-20210323015217-0942afbea50e h1:UimnzLuARNkGi2XsNznUoOLFP/noktdUMrr7fcb3D4U=
+github.com/chromedp/cdproto v0.0.0-20210323015217-0942afbea50e/go.mod h1:At5TxYYdxkbQL0TSefRjhLE3Q0lgvqKKMSFUglJ7i1U=
+github.com/chromedp/chromedp v0.6.10 h1:Yd4X6ngkWbn6A+hv6mUzV9kVHrPn7L4+vf2uyNbze2s=
+github.com/chromedp/chromedp v0.6.10/go.mod h1:Q8L2uDLH9YFYbThK5fqPpyWa3CT4y9dqHLxaQr+Yhl8=
 github.com/chromedp/sysutil v1.0.0 h1:+ZxhTpfpZlmchB58ih/LBHX52ky7w2VhQVKQMucy3Ic=
 github.com/chromedp/sysutil v1.0.0/go.mod h1:kgWmDdq8fTzXYcKIBqIYvRRTnYb9aNS9moAV0xufSww=
 github.com/gobwas/httphead v0.1.0 h1:exrUm0f4YX0L7EBwZHuCF4GDp8aJfVeBrlLQrs6NqWU=
@@ -13,11 +12,10 @@ github.com/gobwas/ws v1.0.4 h1:5eXU1CZhpQdq5kXbKb+sECH5Ia5KiO6CYzIzdlVx6Bs=
 github.com/gobwas/ws v1.0.4/go.mod h1:szmBTxLgaFppYjEmNtny/v3w89xOydFnnZMcgRRu/EM=
 github.com/josharian/intern v1.0.0 h1:vlS4z54oSdjm0bgjRigI+G1HpF+tI+9rE5LLzOg8HmY=
 github.com/josharian/intern v1.0.0/go.mod h1:5DoeVV0s6jJacbCEi61lwdGj/aVlrQvzHFFd8Hwg//Y=
-github.com/mailru/easyjson v0.7.6 h1:8yTIVnZgCoiM1TgqoeTl+LfU5Jg6/xL3QhGQnimLYnA=
-github.com/mailru/easyjson v0.7.6/go.mod h1:xzfreul335JAWq5oZzymOObrkdz5UnU4kGfJJLY9Nlc=
+github.com/mailru/easyjson v0.7.7 h1:UGYAvKxe3sBsEDzO8ZeWOSlIQfWFlxbzLZe7hwFURr0=
+github.com/mailru/easyjson v0.7.7/go.mod h1:xzfreul335JAWq5oZzymOObrkdz5UnU4kGfJJLY9Nlc=
 github.com/wabarc/helper v0.0.0-20210127120855-10af37cc2616 h1:wZ5HtpmZAVUq0Im5Sm92ycJrTeLJk5lB/Kvh55Rd+Ps=
 github.com/wabarc/helper v0.0.0-20210127120855-10af37cc2616/go.mod h1:N9P4r7Rn46p4nkWtXV6ztN3p5ACVnp++bgfwjTqSxQ8=
-golang.org/x/sys v0.0.0-20210122093101-04d7465088b8 h1:de2yTH1xuxjmGB7i6Z5o2z3RCHVa0XlpSZzjd8Fe6bE=
-golang.org/x/sys v0.0.0-20210122093101-04d7465088b8/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
-golang.org/x/sys v0.0.0-20210124154548-22da62e12c0c h1:VwygUrnw9jn88c4u8GD3rZQbqrP/tgas88tPUbBxQrk=
-golang.org/x/sys v0.0.0-20210124154548-22da62e12c0c/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
+golang.org/x/sys v0.0.0-20210320140829-1e4c9ba3b0c4/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
+golang.org/x/sys v0.0.0-20210415045647-66c3f260301c h1:6L+uOeS3OQt/f4eFHXZcTxeZrGCuz+CLElgEBjbcTA4=
+golang.org/x/sys v0.0.0-20210415045647-66c3f260301c/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=

screenshot.go

@@ -2,9 +2,13 @@ package screenshot
 
 import (
 	"context"
+	"encoding/json"
+	"fmt"
 	"log"
 	"math"
+	"net/http"
 	"os"
+	"strings"
 	"sync"
 	"time"
 
@@ -29,6 +33,43 @@ type chromeRemoteScreenshoter struct {
 	url string
 }
 
+// NewChromeRemoteScreenshoter creates a Screenshoter backed by Chrome DevTools Protocol.
+// The addr is the headless chrome websocket debugger endpoint, such as 127.0.0.1:9222.
+func NewChromeRemoteScreenshoter(addr string) (Screenshoter, error) {
+	// Due to issue#505 (https://github.com/chromedp/chromedp/issues/505),
+	// chrome restricts the host must be IP or localhost, we should rewrite the url.
+	req, err := http.NewRequest(http.MethodGet, fmt.Sprintf("http://%s/json/version", addr), nil)
+	if err != nil {
+		return nil, err
+	}
+	req.Host = "localhost"
+	resp, err := http.DefaultClient.Do(req)
+	if err != nil {
+		return nil, err
+	}
+	defer resp.Body.Close()
+
+	var result map[string]interface{}
+	if err = json.NewDecoder(resp.Body).Decode(&result); err != nil {
+		return nil, err
+	}
+
+	return &chromeRemoteScreenshoter{
+		url: strings.Replace(result["webSocketDebuggerUrl"].(string), "localhost", addr, 1),
+	}, nil
+}
+
+func (s *chromeRemoteScreenshoter) Screenshot(ctx context.Context, urls []string, options ...ScreenshotOption) ([]Screenshots, error) {
+	if s.url == "" {
+		return nil, fmt.Errorf("Can't connect to headless browser")
+	}
+
+	allocCtx, cancel := chromedp.NewRemoteAllocator(ctx, s.url)
+	defer cancel()
+
+	return screenshotStart(allocCtx, urls, options...)
+}
+
 func Screenshot(ctx context.Context, urls []string, options ...ScreenshotOption) ([]Screenshots, error) {
 	// https://github.com/chromedp/chromedp/blob/b56cd66f9cebd6a1fa1283847bbf507409d48225/allocate.go#L53
 	var allocOpts = chromedp.DefaultExecAllocatorOptions[:]
@@ -41,6 +82,10 @@ func Screenshot(ctx context.Context, urls []string, options ...ScreenshotOption)
 	allocCtx, cancel := chromedp.NewExecAllocator(ctx, allocOpts...)
 	defer cancel()
 
+	return screenshotStart(allocCtx, urls, options...)
+}
+
+func screenshotStart(allocCtx context.Context, urls []string, options ...ScreenshotOption) ([]Screenshots, error) {
 	var browserOpts []chromedp.ContextOption
 	if debug := os.Getenv("CHROMEDP_DEBUG"); debug != "" && debug != "false" {
 		browserOpts = append(browserOpts, chromedp.WithDebugf(log.Printf))

screenshot_test.go

@@ -5,6 +5,7 @@ import (
 	"io"
 	"net/http"
 	"net/http/httptest"
+	"os/exec"
 	"reflect"
 	"strings"
 	"testing"
@@ -72,6 +73,73 @@ func TestScreenshot(t *testing.T) {
 	}
 }
 
+func TestScreenshotWithRemote(t *testing.T) {
+	ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second)
+	defer cancel()
+
+	ts := newServer()
+	defer ts.Close()
+
+	cmd := exec.Command("chromium-browser", "--headless", "--disable-gpu", "--no-sandbox", "--remote-debugging-port=9222")
+	if err := cmd.Start(); err != nil {
+		t.Fatalf("Start Chromium headless failed: %v", err)
+	}
+	done := make(chan error, 1)
+	go func() {
+		done <- cmd.Wait()
+		select {
+		case <-time.After(60 * time.Second):
+			if err := cmd.Process.Kill(); err != nil {
+				t.Errorf("Failed to kill process: %v", err)
+			}
+			t.Log("Process killed as timeout reached")
+		case err := <-done:
+			if err != nil {
+				t.Errorf("Process finished with error: %v", err)
+			}
+			if err := cmd.Process.Kill(); err != nil {
+				t.Errorf("Failed to kill process: %v", err)
+			}
+			t.Log("Process finished successfully")
+		}
+	}()
+	time.Sleep(3 * time.Second)
+
+	urls := []string{ts.URL}
+	remote, err := NewChromeRemoteScreenshoter("127.0.0.1:9222")
+	if err != nil {
+		t.Fatal(err)
+	}
+	shots, err := remote.Screenshot(ctx, urls, ScaleFactor(1))
+	if err != nil {
+		t.Log(urls)
+		if err == context.DeadlineExceeded {
+			t.Error(err.Error(), http.StatusRequestTimeout)
+			done <- err
+			return
+		}
+		t.Error(err.Error(), http.StatusServiceUnavailable)
+		done <- err
+		return
+	}
+
+	for _, shot := range shots {
+		if reflect.TypeOf(shot) != reflect.TypeOf(Screenshots{}) {
+			t.Fail()
+		}
+
+		if shot.Title != "Example Domain" {
+			t.Log(shot.Title)
+			t.Fail()
+		}
+
+		if shot.Data == nil {
+			t.Fail()
+		}
+	}
+	done <- nil
+}
+
 func TestScreenshotFormat(t *testing.T) {
 	ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second)
 	defer cancel()