Add WithDebugf() context option

Adds the high level WithDebugf() context option, and associated lower
level browser and dial options for setting a protocol wire debugger.
Additionally changes the conn.Conn.Read/Write implementations to be more
efficient, using direct easyjson.{Marshal,Unmarshal} calls and logging
to debug func when available.
This commit is contained in:
Kenneth Shaw 2019-04-09 15:11:14 +07:00 committed by Daniel Martí
parent b481eeac51
commit e8122e4a26
4 changed files with 110 additions and 30 deletions

View File

@ -36,6 +36,7 @@ type Browser struct {
// logging funcs
logf func(string, ...interface{})
errf func(string, ...interface{})
dbgf func(string, ...interface{})
// The optional fields below are helpful for some tests.
@ -60,32 +61,28 @@ type cmdJob struct {
// NewBrowser creates a new browser.
func NewBrowser(ctx context.Context, urlstr string, opts ...BrowserOption) (*Browser, error) {
conn, err := DialContext(ctx, ForceIP(urlstr))
if err != nil {
return nil, err
}
b := &Browser{
conn: conn,
tabQueue: make(chan newTab, 1),
tabResult: make(chan *Target, 1),
cmdQueue: make(chan cmdJob),
logf: log.Printf,
cmdQueue: make(chan cmdJob),
logf: log.Printf,
}
// apply options
for _, o := range opts {
o(b)
}
// ensure errf is set
if b.errf == nil {
b.errf = func(s string, v ...interface{}) { b.logf("ERROR: "+s, v...) }
}
// dial
var err error
b.conn, err = DialContext(ctx, ForceIP(urlstr), WithConnDebugf(b.dbgf))
if err != nil {
return nil, err
}
go b.run(ctx)
return b, nil
}
@ -332,6 +329,12 @@ func WithBrowserErrorf(f func(string, ...interface{})) BrowserOption {
return func(b *Browser) { b.errf = f }
}
// WithBrowserDebugf is a browser option to specify a func to log actual
// websocket messages.
func WithBrowserDebugf(f func(string, ...interface{})) BrowserOption {
return func(b *Browser) { b.dbgf = f }
}
// WithConsolef is a browser option to specify a func to receive chrome log events.
//
// Note: NOT YET IMPLEMENTED.

View File

@ -232,6 +232,11 @@ func WithErrorf(f func(string, ...interface{})) ContextOption {
return WithBrowserOption(WithBrowserErrorf(f))
}
// WithDebugf is a shortcut for WithBrowserOption(WithBrowserDebugf(f)).
func WithDebugf(f func(string, ...interface{})) ContextOption {
return WithBrowserOption(WithBrowserDebugf(f))
}
// WithBrowserOption allows passing a number of browser options to the allocator
// when allocating a new browser. As such, this context option can only be used
// when NewContext is allocating a new browser.
@ -248,7 +253,6 @@ func WithBrowserOption(opts ...BrowserOption) ContextOption {
func Targets(ctx context.Context) ([]*target.Info, error) {
// Don't rely on Run, as that needs to be able to call Targets, and we
// don't want cyclic func calls.
c := FromContext(ctx)
if c == nil || c.Allocator == nil {
return nil, ErrInvalidContext

102
conn.go
View File

@ -3,11 +3,13 @@ package chromedp
import (
"context"
"io"
"io/ioutil"
"net"
"strings"
"github.com/chromedp/cdproto"
"github.com/gorilla/websocket"
"github.com/mailru/easyjson"
)
var (
@ -28,24 +30,11 @@ type Transport interface {
// Conn wraps a gorilla/websocket.Conn connection.
type Conn struct {
*websocket.Conn
}
// Read reads the next message.
func (c *Conn) Read() (*cdproto.Message, error) {
msg := new(cdproto.Message)
if err := c.ReadJSON(msg); err != nil {
return nil, err
}
return msg, nil
}
// Write writes a message.
func (c *Conn) Write(msg *cdproto.Message) error {
return c.WriteJSON(msg)
dbgf func(string, ...interface{})
}
// DialContext dials the specified websocket URL using gorilla/websocket.
func DialContext(ctx context.Context, urlstr string) (*Conn, error) {
func DialContext(ctx context.Context, urlstr string, opts ...DialOption) (*Conn, error) {
d := &websocket.Dialer{
ReadBufferSize: DefaultReadBufferSize,
WriteBufferSize: DefaultWriteBufferSize,
@ -57,7 +46,78 @@ func DialContext(ctx context.Context, urlstr string) (*Conn, error) {
return nil, err
}
return &Conn{conn}, nil
// apply opts
c := &Conn{
Conn: conn,
}
for _, o := range opts {
o(c)
}
return c, nil
}
// Read reads the next message.
func (c *Conn) Read() (*cdproto.Message, error) {
// get websocket reader
typ, r, err := c.NextReader()
if err != nil {
return nil, err
}
if typ != websocket.TextMessage {
return nil, ErrInvalidWebsocketMessage
}
// when dbgf defined, buffer, log, unmarshal
if c.dbgf != nil {
// buffer output
buf, err := ioutil.ReadAll(r)
if err != nil {
return nil, err
}
c.dbgf("<- %s", string(buf))
msg := new(cdproto.Message)
if err = easyjson.Unmarshal(buf, msg); err != nil {
return nil, err
}
return msg, nil
}
// unmarshal direct from reader
msg := new(cdproto.Message)
if err = easyjson.UnmarshalFromReader(r, msg); err != nil {
return nil, err
}
return msg, nil
}
// Write writes a message.
func (c *Conn) Write(msg *cdproto.Message) error {
w, err := c.NextWriter(websocket.TextMessage)
if err != nil {
return err
}
if c.dbgf != nil {
var buf []byte
buf, err = easyjson.Marshal(msg)
if err != nil {
return err
}
c.dbgf("-> %s", string(buf))
_, err = w.Write(buf)
if err != nil {
return err
}
} else {
// direct marshal
_, err = easyjson.MarshalToWriter(msg, w)
if err != nil {
return err
}
}
return w.Close()
}
// ForceIP forces the host component in urlstr to be an IP address.
@ -80,3 +140,13 @@ func ForceIP(urlstr string) string {
}
return urlstr
}
// DialOption is a dial option.
type DialOption func(*Conn)
// WithConnDebugf is a dial option to set a protocol logger.
func WithConnDebugf(f func(string, ...interface{})) DialOption {
return func(c *Conn) {
c.dbgf = f
}
}

View File

@ -10,6 +10,9 @@ func (err Error) Error() string {
// Error types.
const (
// ErrInvalidWebsocketMessage is the invalid websocket message.
ErrInvalidWebsocketMessage Error = "invalid websocket message"
// ErrInvalidDimensions is the invalid dimensions error.
ErrInvalidDimensions Error = "invalid dimensions"