package lorca import ( "bufio" "encoding/json" "errors" "fmt" "io" "io/ioutil" "log" "os/exec" "regexp" "sync" "sync/atomic" "golang.org/x/net/websocket" ) type h = map[string]interface{} // Result is a struct for the resulting value of the JS expression or an error. type result struct { Value json.RawMessage Err error } type bindingFunc func(args []json.RawMessage) (interface{}, error) // Msg is a struct for incoming messages (results and async events) type msg struct { ID int `json:"id"` Result json.RawMessage `json:"result"` Error json.RawMessage `json:"error"` Method string `json:"method"` Params json.RawMessage `json:"params"` } type chrome struct { sync.Mutex cmd *exec.Cmd ws *websocket.Conn id int32 target string session string window int pending map[int]chan result bindings map[string]bindingFunc } func newChromeWithArgs(chromeBinary string, args ...string) (*chrome, error) { // The first two IDs are used internally during the initialization c := &chrome{ id: 2, pending: map[int]chan result{}, bindings: map[string]bindingFunc{}, } // Start chrome process c.cmd = exec.Command(chromeBinary, args...) pipe, err := c.cmd.StderrPipe() if err != nil { return nil, err } if err := c.cmd.Start(); err != nil { return nil, err } // Wait for websocket address to be printed to stderr re := regexp.MustCompile(`^DevTools listening on (ws://.*?)\r?\n$`) m, err := readUntilMatch(pipe, re) if err != nil { c.kill() return nil, err } wsURL := m[1] // Open a websocket c.ws, err = websocket.Dial(wsURL, "", "http://127.0.0.1") if err != nil { c.kill() return nil, err } // Find target and initialize session c.target, err = c.findTarget() if err != nil { c.kill() return nil, err } c.session, err = c.startSession(c.target) if err != nil { c.kill() return nil, err } go c.readLoop() for method, args := range map[string]h{ "Page.enable": nil, "Target.setAutoAttach": {"autoAttach": true, "waitForDebuggerOnStart": false}, "Network.enable": nil, "Runtime.enable": nil, "Security.enable": nil, "Performance.enable": nil, "Log.enable": nil, } { if _, err := c.send(method, args); err != nil { c.kill() c.cmd.Wait() return nil, err } } if !contains(args, "--headless") { win, err := c.getWindowForTarget(c.target) if err != nil { c.kill() return nil, err } c.window = win.WindowID } return c, nil } func (c *chrome) findTarget() (string, error) { err := websocket.JSON.Send(c.ws, h{ "id": 0, "method": "Target.setDiscoverTargets", "params": h{"discover": true}, }) if err != nil { return "", err } for { m := msg{} if err = websocket.JSON.Receive(c.ws, &m); err != nil { return "", err } else if m.Method == "Target.targetCreated" { target := struct { TargetInfo struct { Type string `json:"type"` ID string `json:"targetId"` } `json:"targetInfo"` }{} if err := json.Unmarshal(m.Params, &target); err != nil { return "", err } else if target.TargetInfo.Type == "page" { return target.TargetInfo.ID, nil } } } } func (c *chrome) startSession(target string) (string, error) { err := websocket.JSON.Send(c.ws, h{ "id": 1, "method": "Target.attachToTarget", "params": h{"targetId": target}, }) if err != nil { return "", err } for { m := msg{} if err = websocket.JSON.Receive(c.ws, &m); err != nil { return "", err } else if m.ID == 1 { if m.Error != nil { return "", errors.New("Target error: " + string(m.Error)) } session := struct { ID string `json:"sessionId"` }{} if err := json.Unmarshal(m.Result, &session); err != nil { return "", err } return session.ID, nil } } } // WindowState defines the state of the Chrome window, possible values are // "normal", "maximized", "minimized" and "fullscreen". type WindowState string const ( // WindowStateNormal defines a normal state of the browser window WindowStateNormal WindowState = "normal" // WindowStateMaximized defines a maximized state of the browser window WindowStateMaximized WindowState = "maximized" // WindowStateMinimized defines a minimized state of the browser window WindowStateMinimized WindowState = "minimized" // WindowStateFullscreen defines a fullscreen state of the browser window WindowStateFullscreen WindowState = "fullscreen" ) // Bounds defines settable window properties. type Bounds struct { Left int `json:"left"` Top int `json:"top"` Width int `json:"width"` Height int `json:"height"` WindowState WindowState `json:"windowState"` } type windowTargetMessage struct { WindowID int `json:"windowId"` Bounds Bounds `json:"bounds"` } func (c *chrome) getWindowForTarget(target string) (windowTargetMessage, error) { var m windowTargetMessage msg, err := c.send("Browser.getWindowForTarget", h{"targetId": target}) if err != nil { return m, err } err = json.Unmarshal(msg, &m) return m, err } type targetMessageTemplate struct { ID int `json:"id"` Method string `json:"method"` Params struct { Name string `json:"name"` Payload string `json:"payload"` ID int `json:"executionContextId"` Args []struct { Type string `json:"type"` Value interface{} `json:"value"` } `json:"args"` } `json:"params"` Error struct { Message string `json:"message"` } `json:"error"` Result json.RawMessage `json:"result"` } type targetMessage struct { targetMessageTemplate Result struct { Result struct { Type string `json:"type"` Subtype string `json:"subtype"` Description string `json:"description"` Value json.RawMessage `json:"value"` ObjectID string `json:"objectId"` } `json:"result"` Exception struct { Exception struct { Value json.RawMessage `json:"value"` } `json:"exception"` } `json:"exceptionDetails"` } `json:"result"` } func (c *chrome) readLoop() { for { m := msg{} if err := websocket.JSON.Receive(c.ws, &m); err != nil { return } if m.Method == "Target.receivedMessageFromTarget" { params := struct { SessionID string `json:"sessionId"` Message string `json:"message"` }{} json.Unmarshal(m.Params, ¶ms) if params.SessionID != c.session { continue } res := targetMessage{} json.Unmarshal([]byte(params.Message), &res) if res.ID == 0 && res.Method == "Runtime.consoleAPICalled" || res.Method == "Runtime.exceptionThrown" { log.Println(params.Message) } else if res.ID == 0 && res.Method == "Runtime.bindingCalled" { payload := struct { Name string `json:"name"` Seq int `json:"seq"` Args []json.RawMessage `json:"args"` }{} json.Unmarshal([]byte(res.Params.Payload), &payload) c.Lock() binding, ok := c.bindings[res.Params.Name] c.Unlock() if ok { jsString := func(v interface{}) string { b, _ := json.Marshal(v); return string(b) } go func() { result, error := "", `""` if r, err := binding(payload.Args); err != nil { error = jsString(err.Error()) } else if b, err := json.Marshal(r); err != nil { error = jsString(err.Error()) } else { result = string(b) } expr := fmt.Sprintf(` if (%[4]s) { window['%[1]s']['errors'].get(%[2]d)(%[4]s); } else { window['%[1]s']['callbacks'].get(%[2]d)(%[3]s); } window['%[1]s']['callbacks'].delete(%[2]d); window['%[1]s']['errors'].delete(%[2]d); `, payload.Name, payload.Seq, result, error) c.send("Runtime.evaluate", h{"expression": expr, "contextId": res.Params.ID}) }() } continue } c.Lock() resc, ok := c.pending[res.ID] delete(c.pending, res.ID) c.Unlock() if !ok { continue } if res.Error.Message != "" { resc <- result{Err: errors.New(res.Error.Message)} } else if res.Result.Exception.Exception.Value != nil { resc <- result{Err: errors.New(string(res.Result.Exception.Exception.Value))} } else if res.Result.Result.Type == "object" && res.Result.Result.Subtype == "error" { resc <- result{Err: errors.New(res.Result.Result.Description)} } else if res.Result.Result.Type != "" { resc <- result{Value: res.Result.Result.Value} } else { res := targetMessageTemplate{} json.Unmarshal([]byte(params.Message), &res) resc <- result{Value: res.Result} } } else if m.Method == "Target.targetDestroyed" { params := struct { TargetID string `json:"targetId"` }{} json.Unmarshal(m.Params, ¶ms) if params.TargetID == c.target { c.kill() return } } } } func (c *chrome) send(method string, params h) (json.RawMessage, error) { id := atomic.AddInt32(&c.id, 1) b, err := json.Marshal(h{"id": int(id), "method": method, "params": params}) if err != nil { return nil, err } resc := make(chan result) c.Lock() c.pending[int(id)] = resc c.Unlock() if err := websocket.JSON.Send(c.ws, h{ "id": int(id), "method": "Target.sendMessageToTarget", "params": h{"message": string(b), "sessionId": c.session}, }); err != nil { return nil, err } res := <-resc return res.Value, res.Err } func (c *chrome) load(url string) error { _, err := c.send("Page.navigate", h{"url": url}) return err } func (c *chrome) eval(expr string) (json.RawMessage, error) { return c.send("Runtime.evaluate", h{"expression": expr, "awaitPromise": true, "returnByValue": true}) } func (c *chrome) bind(name string, f bindingFunc) error { c.Lock() // check if binding already exists _, exists := c.bindings[name] c.bindings[name] = f c.Unlock() if exists { // Just replace callback and return, as the binding was already added to js // and adding it again would break it. return nil } if _, err := c.send("Runtime.addBinding", h{"name": name}); err != nil { return err } script := fmt.Sprintf(`(() => { const bindingName = '%s'; const binding = window[bindingName]; window[bindingName] = async (...args) => { const me = window[bindingName]; let errors = me['errors']; let callbacks = me['callbacks']; if (!callbacks) { callbacks = new Map(); me['callbacks'] = callbacks; } if (!errors) { errors = new Map(); me['errors'] = errors; } const seq = (me['lastSeq'] || 0) + 1; me['lastSeq'] = seq; const promise = new Promise((resolve, reject) => { callbacks.set(seq, resolve); errors.set(seq, reject); }); binding(JSON.stringify({name: bindingName, seq, args})); return promise; }})(); `, name) _, err := c.send("Page.addScriptToEvaluateOnNewDocument", h{"source": script}) if err != nil { return err } _, err = c.eval(script) return err } func (c *chrome) setBounds(b Bounds) error { if b.WindowState == "" { b.WindowState = WindowStateNormal } param := h{"windowId": c.window, "bounds": b} if b.WindowState != WindowStateNormal { param["bounds"] = h{"windowState": b.WindowState} } _, err := c.send("Browser.setWindowBounds", param) return err } func (c *chrome) bounds() (Bounds, error) { result, err := c.send("Browser.getWindowBounds", h{"windowId": c.window}) if err != nil { return Bounds{}, err } bounds := struct { Bounds Bounds `json:"bounds"` }{} err = json.Unmarshal(result, &bounds) return bounds.Bounds, err } func (c *chrome) pdf(width, height int) ([]byte, error) { result, err := c.send("Page.printToPDF", h{ "paperWidth": float32(width) / 96, "paperHeight": float32(height) / 96, }) if err != nil { return nil, err } pdf := struct { Data []byte `json:"data"` }{} err = json.Unmarshal(result, &pdf) return pdf.Data, err } func (c *chrome) png(x, y, width, height int, bg uint32, scale float32) ([]byte, error) { if x == 0 && y == 0 && width == 0 && height == 0 { // By default either use SVG size if it's an SVG, or use A4 page size bounds, err := c.eval(`document.rootElement ? [document.rootElement.x.baseVal.value, document.rootElement.y.baseVal.value, document.rootElement.width.baseVal.value, document.rootElement.height.baseVal.value] : [0,0,816,1056]`) if err != nil { return nil, err } rect := make([]int, 4) if err := json.Unmarshal(bounds, &rect); err != nil { return nil, err } x, y, width, height = rect[0], rect[1], rect[2], rect[3] } _, err := c.send("Emulation.setDefaultBackgroundColorOverride", h{ "color": h{ "r": (bg >> 16) & 0xff, "g": (bg >> 8) & 0xff, "b": bg & 0xff, "a": (bg >> 24) & 0xff, }, }) if err != nil { return nil, err } result, err := c.send("Page.captureScreenshot", h{ "clip": h{ "x": x, "y": y, "width": width, "height": height, "scale": scale, }, }) if err != nil { return nil, err } pdf := struct { Data []byte `json:"data"` }{} err = json.Unmarshal(result, &pdf) return pdf.Data, err } func (c *chrome) kill() error { if c.ws != nil { if err := c.ws.Close(); err != nil { return err } } // TODO: cancel all pending requests if state := c.cmd.ProcessState; state == nil || !state.Exited() { return c.cmd.Process.Kill() } return nil } func readUntilMatch(r io.ReadCloser, re *regexp.Regexp) ([]string, error) { br := bufio.NewReader(r) for { if line, err := br.ReadString('\n'); err != nil { r.Close() return nil, err } else if m := re.FindStringSubmatch(line); m != nil { go io.Copy(ioutil.Discard, br) return m, nil } } } func contains(arr []string, x string) bool { for _, n := range arr { if x == n { return true } } return false }