finish updating websocket/read/send logic

Author: wirepair

Gopkg.lock

@@ -25,16 +25,14 @@ THE SOFTWARE.
 package gcd
 
 import (
-	"io"
+	"context"
 	"log"
-	"net"
 	"sync"
 	"sync/atomic"
 	"time"
 
 	"github.com/wirepair/gcd/gcdapi"
 	"github.com/wirepair/gcd/gcdmessage"
-	"golang.org/x/net/websocket"
 )
 
 // TargetInfo defines the 'tab' or target for this chrome instance,
@@ -57,13 +55,14 @@ type TargetInfo struct {
 // Events are handled by mapping the method name to a function which takes a target and byte output.
 // For now, callers will need to unmarshall the types themselves.
 type ChromeTarget struct {
+	ctx    context.Context
 	sendId int64 // An Id which is atomically incremented per request.
 	// must be at top because of alignement and atomic usage
 	replyLock       sync.RWMutex                           // lock for dispatching responses
 	replyDispatcher map[int64]chan *gcdmessage.Message     // Replies to synch methods using a non-buffered channel
 	eventLock       sync.RWMutex                           // lock for dispatching events
 	eventDispatcher map[string]func(*ChromeTarget, []byte) // calls the function when events match the subscribed method
-	conn            *websocket.Conn                        // the connection to the chrome debugger service for this tab/process
+	conn            *wsConn                                // the connection to the chrome debugger service for this tab/process
 
 	// Chrome Debugger Domains
 	Accessibility        *gcdapi.Accessibility
@@ -122,8 +121,8 @@ type ChromeTarget struct {
 }
 
 // openChromeTarget creates a new Chrome Target by connecting to the service given the URL taken from initial connection.
-func openChromeTarget(addr string, target *TargetInfo) (*ChromeTarget, error) {
-	conn, err := wsConnection(addr, target.WebSocketDebuggerUrl)
+func openChromeTarget(ctx context.Context, addr string, target *TargetInfo) (*ChromeTarget, error) {
+	conn, err := wsConnection(ctx, target.WebSocketDebuggerUrl)
 	if err != nil {
 		return nil, err
 	}
@@ -135,6 +134,7 @@ func openChromeTarget(addr string, target *TargetInfo) (*ChromeTarget, error) {
 	doneCh := make(chan struct{})
 	chromeTarget := &ChromeTarget{conn: conn, Target: target, sendCh: sendCh, replyDispatcher: replier, eventDispatcher: eventer, doneCh: doneCh, sendId: 0}
 	chromeTarget.apiTimeout = 120 * time.Second // default 120 seconds to wait for chrome to respond to us
+	chromeTarget.ctx = ctx
 	chromeTarget.Init()
 	chromeTarget.listen()
 	return chromeTarget, nil
@@ -262,9 +262,10 @@ func (c *ChromeTarget) listenWrite() {
 			c.replyLock.Unlock()
 
 			c.debugf("%d sending to chrome. %s\n", msg.Id, msg.Data)
-
-			err := websocket.Message.Send(c.conn, string(msg.Data))
+			log.Printf("%d sending to chrome. %s\n", msg.Id, msg.Data)
+			err := c.conn.Write(c.ctx, msg.Data)
 			if err != nil {
+				log.Printf("error sending message: %s\n", err)
 				c.debugf("error sending message: %s\n", err)
 				return
 			}
@@ -277,25 +278,42 @@ func (c *ChromeTarget) listenWrite() {
 
 // Listens for responses coming in from the Chrome Debugger Service.
 func (c *ChromeTarget) listenRead() {
+	readCh := make(chan []byte, 1)
+	writeClosed := make(chan struct{})
+	go func() {
+		for {
+			var msg []byte
+			err := c.conn.Read(c.ctx, &msg)
+			if err != nil {
+				c.debugf("error in ws read: %s\n", err)
+				close(writeClosed)
+				return
+			} else {
+				select {
+				case <-c.ctx.Done():
+					return
+				case readCh <- msg:
+				}
+			}
+		}
+	}()
+
 	for {
 		select {
+		case <-writeClosed:
+			return
 		// receive done from listenWrite
 		case <-c.doneCh:
 			return
-		// read data from websocket connection
-		default:
-			var msg string
-			err := websocket.Message.Receive(c.conn, &msg)
-			if err == io.EOF {
-				c.debugf("error io.EOF in websocket read")
-				return
-			} else if err != nil {
-				c.debugf("error in ws read: %s\n", err)
-			} else {
-				go c.dispatchResponse([]byte(msg))
+		case <-c.ctx.Done():
+			return
+		case msg := <-readCh:
+			if len(msg) != 0 {
+				c.dispatchResponse(msg)
 			}
 		}
 	}
+
 }
 
 type responseHeader struct {
@@ -374,20 +392,11 @@ func (c *ChromeTarget) checkTargetDisconnected(method string) {
 }
 
 // Connects to the tab/process for sending/recv'ing debug events
-func wsConnection(addr, url string) (*websocket.Conn, error) {
-	conn, err := net.Dial("tcp", addr)
+func wsConnection(ctx context.Context, url string) (*wsConn, error) {
+	client, err := newWsConnDial(ctx, url)
 	if err != nil {
 		return nil, err
 	}
-
-	config, errConfig := websocket.NewConfig(url, "http://localhost")
-	if errConfig != nil {
-		return nil, errConfig
-	}
-	client, errWS := websocket.NewClient(config, conn)
-	if errWS != nil {
-		return nil, errWS
-	}
 	return client, nil
 }
 

Gopkg.toml

@@ -25,6 +25,7 @@ THE SOFTWARE.
 package gcd
 
 import (
+	"context"
 	"errors"
 	"fmt"
 	"io/ioutil"
@@ -83,6 +84,7 @@ type Gcd struct {
 	flags             []string
 	env               []string
 	chomeApiVersion   string
+	ctx               context.Context
 }
 
 // Give it a friendly name.
@@ -94,6 +96,7 @@ func NewChromeDebugger() *Gcd {
 	c.terminatedHandler = nil
 	c.flags = make([]string, 0)
 	c.env = make([]string, 0)
+	c.ctx = context.Background()
 	return c
 }
 
@@ -261,7 +264,7 @@ func (c *Gcd) GetNewTargets(knownIds map[string]struct{}) ([]*ChromeTarget, erro
 	chromeTargets := make([]*ChromeTarget, 0)
 	for _, v := range connectableTargets {
 		if _, ok := knownIds[v.Id]; !ok {
-			target, err := openChromeTarget(c.addr, v)
+			target, err := openChromeTarget(c.ctx, c.addr, v)
 			if err != nil {
 				return nil, err
 			}
@@ -316,7 +319,7 @@ func (c *Gcd) NewTab() (*ChromeTarget, error) {
 	if err != nil {
 		return nil, &GcdDecodingErr{Message: err.Error()}
 	}
-	return openChromeTarget(c.addr, tabTarget)
+	return openChromeTarget(c.ctx, c.addr, tabTarget)
 }
 
 // GetFirstTab returns the first tab created, to be called when
@@ -328,7 +331,7 @@ func (c *Gcd) GetFirstTab() (*ChromeTarget, error) {
 	}
 	for _, tabTarget := range connectableTargets {
 		if tabTarget.Type == "page" {
-			return openChromeTarget(c.addr, tabTarget)
+			return openChromeTarget(c.ctx, c.addr, tabTarget)
 		}
 	}
 	return nil, ErrNoTabAvailable

chrome_target.go

@@ -3,8 +3,13 @@ module github.com/wirepair/gcd
 go 1.14
 
 require (
+	github.com/gobwas/httphead v0.0.0-20180130184737-2c6c146eadee // indirect
+	github.com/gobwas/pool v0.2.0 // indirect
+	github.com/gobwas/ws v1.0.3
 	github.com/json-iterator/go v1.1.9
 	github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd // indirect
 	github.com/modern-go/reflect2 v1.0.1 // indirect
-	golang.org/x/net v0.0.0-20181207154023-610586996380
+	github.com/stretchr/testify v1.4.0 // indirect
+	golang.org/x/sys v0.0.0-20200116001909-b77594299b42 // indirect
+	gopkg.in/yaml.v2 v2.2.8 // indirect
 )

gcd.go

@@ -1,5 +1,12 @@
 github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
+github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
 github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
+github.com/gobwas/httphead v0.0.0-20180130184737-2c6c146eadee h1:s+21KNqlpePfkah2I+gwHF8xmJWRjooY+5248k6m4A0=
+github.com/gobwas/httphead v0.0.0-20180130184737-2c6c146eadee/go.mod h1:L0fX3K22YWvt/FAX9NnzrNzcI4wNYi9Yku4O0LKYflo=
+github.com/gobwas/pool v0.2.0 h1:QEmUOlnSjWtnpRGHF3SauEiOsy82Cup83Vf2LcMlnc8=
+github.com/gobwas/pool v0.2.0/go.mod h1:q8bcK0KcYlCgd9e7WYLm9LpyS+YeLd8JVDW6WezmKEw=
+github.com/gobwas/ws v1.0.3 h1:ZOigqf7iBxkA4jdQ3am7ATzdlOFp9YzA6NmuvEEZc9g=
+github.com/gobwas/ws v1.0.3/go.mod h1:szmBTxLgaFppYjEmNtny/v3w89xOydFnnZMcgRRu/EM=
 github.com/google/gofuzz v1.0.0/go.mod h1:dBl0BpW6vV/+mYPU4Po3pmUjxk6FQPldtuIdl/M65Eg=
 github.com/json-iterator/go v1.1.9 h1:9yzud/Ht36ygwatGx56VwCZtlI/2AD15T1X2sjSuGns=
 github.com/json-iterator/go v1.1.9/go.mod h1:KdQUCv79m/52Kvf8AW2vK1V8akMuk1QjK/uOdHXbAo4=
@@ -11,8 +18,16 @@ github.com/modern-go/reflect2 v0.0.0-20180701023420-4b7aa43c6742 h1:Esafd1046DLD
 github.com/modern-go/reflect2 v0.0.0-20180701023420-4b7aa43c6742/go.mod h1:bx2lNnkwVCuqBIxFjflWJWanXIb3RllmbCylyMrvgv0=
 github.com/modern-go/reflect2 v1.0.1 h1:9f412s+6RmYXLWZSEzVVgPGK7C2PphHj5RJrvfx9AWI=
 github.com/modern-go/reflect2 v1.0.1/go.mod h1:bx2lNnkwVCuqBIxFjflWJWanXIb3RllmbCylyMrvgv0=
+github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
 github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
 github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
 github.com/stretchr/testify v1.3.0/go.mod h1:M5WIy9Dh21IEIfnGCwXGc5bZfKNJtfHm1UVUgZn+9EI=
-golang.org/x/net v0.0.0-20181207154023-610586996380 h1:zPQexyRtNYBc7bcHmehl1dH6TB3qn8zytv8cBGLDNY0=
-golang.org/x/net v0.0.0-20181207154023-610586996380/go.mod h1:mL1N/T3taQHkDXs73rZJwtUhF3w3ftmwwsq0BUmARs4=
+github.com/stretchr/testify v1.4.0 h1:2E4SXV/wtOkTonXsotYi4li6zVWxYlZuYNCXe9XRJyk=
+github.com/stretchr/testify v1.4.0/go.mod h1:j7eGeouHqKxXV5pUuKE4zz7dFj8WfuZ+81PSLYec5m4=
+golang.org/x/sys v0.0.0-20200116001909-b77594299b42 h1:vEOn+mP2zCOVzKckCZy6YsCtDblrpj/w7B9nxGNELpg=
+golang.org/x/sys v0.0.0-20200116001909-b77594299b42/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
+gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM=
+gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
+gopkg.in/yaml.v2 v2.2.2/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI=
+gopkg.in/yaml.v2 v2.2.8 h1:obN1ZagJSUGI0Ek/LBmuj4SNLPfIny3KsKFopxRdj10=
+gopkg.in/yaml.v2 v2.2.8/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI=

go.mod

@@ -0,0 +1,92 @@
+package gcd
+
+import (
+	"bytes"
+	"context"
+	"fmt"
+	"io"
+	"net"
+	"net/url"
+	"time"
+
+	"github.com/gobwas/ws"
+	"github.com/gobwas/ws/wsutil"
+)
+
+const writeSize = 15000000 // 15mb
+
+// adapted from https://github.com/chromedp/chromedp/blob/8e0a16689423d48d8907c62a543c7ea468059228/conn.go
+type wsConn struct {
+	conn   net.Conn
+	writer wsutil.Writer
+}
+
+func newWsConnDial(ctx context.Context, url string) (*wsConn, error) {
+	wconn := &wsConn{}
+	conn, br, _, err := ws.Dial(ctx, url)
+	if err != nil {
+		return nil, err
+	}
+	if br != nil {
+		panic("br should be nil")
+	}
+	wconn.conn = conn
+	wconn.writer = *wsutil.NewWriterBufferSize(conn, ws.StateClientSide, ws.OpText, 1<<20)
+	return wconn, nil
+}
+
+func formatURL(toFormat string) string {
+	u, err := url.Parse(toFormat)
+	if err != nil {
+		return ""
+	}
+	host, port, err := net.SplitHostPort(u.Host)
+	if err != nil {
+		return ""
+	}
+	addr, err := net.ResolveIPAddr("ip", host)
+	if err != nil {
+		return ""
+	}
+	u.Host = net.JoinHostPort(addr.IP.String(), port)
+	return u.String()
+}
+
+func (c *wsConn) Read(ctx context.Context, msg *[]byte) error {
+	// get websocket reader
+	c.conn.SetReadDeadline(time.Now().Add(10 * time.Second))
+	reader := wsutil.NewReader(c.conn, ws.StateClientSide)
+	h, err := reader.NextFrame()
+	if err != nil {
+		return err
+	}
+
+	if h.OpCode == ws.OpClose {
+		return io.EOF
+	}
+
+	if h.OpCode != ws.OpText {
+		return fmt.Errorf("InvalidWebsocketMessage")
+	}
+
+	var b bytes.Buffer
+	if _, err := b.ReadFrom(reader); err != nil {
+		return err
+	}
+
+	*msg = b.Bytes()
+	return nil
+}
+
+// Write writes a message.
+func (c *wsConn) Write(_ context.Context, msg []byte) error {
+	c.writer.Reset(c.conn, ws.StateClientSide, ws.OpText)
+	if _, err := c.writer.Write(msg); err != nil {
+		return err
+	}
+	return c.writer.Flush()
+}
+
+func (c *wsConn) Close() error {
+	return c.Close()
+}