Skip to content

Commit

Permalink
update messages
Browse files Browse the repository at this point in the history
  • Loading branch information
hslam committed Jan 17, 2021
1 parent 87d2bf3 commit 41d0961
Show file tree
Hide file tree
Showing 9 changed files with 29 additions and 118 deletions.
2 changes: 1 addition & 1 deletion conn.go
Original file line number Diff line number Diff line change
Expand Up @@ -68,7 +68,7 @@ func (c *Conn) Read(b []byte) (n int, err error) {
c.connBuffer = c.connBuffer[:num]
return len(b), nil
}
f, err := c.readFrame()
f, err := c.readFrame(nil)
if err != nil {
return 0, err
}
Expand Down
Binary file removed example/.DS_Store
Binary file not shown.
25 changes: 0 additions & 25 deletions example/client/client.go

This file was deleted.

25 changes: 0 additions & 25 deletions example/client/client.html

This file was deleted.

29 changes: 0 additions & 29 deletions example/server/server.go

This file was deleted.

27 changes: 0 additions & 27 deletions example/server_poll/server_poll.go

This file was deleted.

23 changes: 20 additions & 3 deletions frame.go
Original file line number Diff line number Diff line change
Expand Up @@ -62,7 +62,7 @@ func (c *Conn) putFrame(f *frame) {
framePool.Put(f)
}

func (c *Conn) readFrame() (f *frame, err error) {
func (c *Conn) readFrame(buf []byte) (f *frame, err error) {
f = c.getFrame()
for {
length := uint64(len(c.buffer))
Expand All @@ -78,7 +78,17 @@ func (c *Conn) readFrame() (f *frame, err error) {
} else if offset == 0 {
goto read
} else {
c.buffer = c.buffer[offset:]
msgLength := len(f.PayloadData)
var p []byte
if cap(buf) >= msgLength {
p = buf[:msgLength]
} else {
p = make([]byte, msgLength)
}
copy(p, f.PayloadData)
f.PayloadData = p
n := copy(c.buffer, c.buffer[offset:])
c.buffer = c.buffer[:n]
return
}
}
Expand All @@ -105,7 +115,14 @@ func (c *Conn) readFrame() (f *frame, err error) {
}
return nil, err
} else if n > 0 {
c.buffer = append(c.buffer, readBuffer[:n]...)
length := len(c.buffer)
size := length + n
if cap(c.buffer) >= size {
c.buffer = c.buffer[:size]
copy(c.buffer[length:], readBuffer[:n])
} else {
c.buffer = append(c.buffer, readBuffer[:n]...)
}
if c.shared {
c.readPool.Put(readBuffer)
}
Expand Down
8 changes: 4 additions & 4 deletions messages.go
Original file line number Diff line number Diff line change
Expand Up @@ -25,7 +25,7 @@ func (c *Conn) ReceiveMessage(v interface{}) (err error) {
defer c.reading.Unlock()
c.buffer = c.buffer[:0]
c.connBuffer = c.connBuffer[:0]
f, err := c.readFrame()
f, err := c.readFrame(nil)
if err != nil {
return err
}
Expand Down Expand Up @@ -94,12 +94,12 @@ func (c *Conn) SendMessage(v interface{}) (err error) {
}

// ReadMessage reads single message from ws.
func (c *Conn) ReadMessage() (p []byte, err error) {
func (c *Conn) ReadMessage(buf []byte) (p []byte, err error) {
c.reading.Lock()
defer c.reading.Unlock()
c.buffer = c.buffer[:0]
c.connBuffer = c.connBuffer[:0]
f, err := c.readFrame()
f, err := c.readFrame(buf)
if err != nil {
return nil, err
}
Expand Down Expand Up @@ -128,7 +128,7 @@ func (c *Conn) ReadTextMessage() (p string, err error) {
defer c.reading.Unlock()
c.buffer = c.buffer[:0]
c.connBuffer = c.connBuffer[:0]
f, err := c.readFrame()
f, err := c.readFrame(nil)
if err != nil {
return "", err
}
Expand Down
8 changes: 4 additions & 4 deletions websocket_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,7 @@ func TestWebsocket(t *testing.T) {
addr := ":8080"
Serve := func(conn *Conn) {
for {
msg, err := conn.ReadMessage()
msg, err := conn.ReadMessage(nil)
if err != nil {
break
}
Expand All @@ -35,7 +35,7 @@ func TestWebsocket(t *testing.T) {
if err := conn.WriteMessage([]byte(msg)); err != nil {
t.Error(err)
}
data, err := conn.ReadMessage()
data, err := conn.ReadMessage(nil)
if err != nil {
t.Error(err)
} else if string(data) != msg {
Expand All @@ -50,7 +50,7 @@ func TestUpgrade(t *testing.T) {
addr := ":8080"
Serve := func(conn *Conn) {
for {
msg, err := conn.ReadMessage()
msg, err := conn.ReadMessage(nil)
if err != nil {
break
}
Expand Down Expand Up @@ -90,7 +90,7 @@ func TestUpgrade(t *testing.T) {
if err := conn.WriteMessage([]byte(msg)); err != nil {
t.Error(err)
}
data, err := conn.ReadMessage()
data, err := conn.ReadMessage(nil)
if err != nil {
t.Error(err)
} else if string(data) != msg {
Expand Down

0 comments on commit 41d0961

Please sign in to comment.