diff --git a/browser.go b/browser.go index add65ae..636177f 100644 --- a/browser.go +++ b/browser.go @@ -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. diff --git a/chromedp.go b/chromedp.go index 3a56d0a..e9bf091 100644 --- a/chromedp.go +++ b/chromedp.go @@ -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 diff --git a/conn.go b/conn.go index fa17be3..5347237 100644 --- a/conn.go +++ b/conn.go @@ -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 + } +} diff --git a/errors.go b/errors.go index b4bcb9b..ce03779 100644 --- a/errors.go +++ b/errors.go @@ -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"