Skip to content

Commit 058e736

Browse files
committed
Add msgType parameter to NetConn adapter
Closes #113
1 parent 5028f22 commit 058e736

File tree

2 files changed

+14
-11
lines changed

2 files changed

+14
-11
lines changed

netconn.go

Lines changed: 12 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@ package websocket
22

33
import (
44
"context"
5+
"fmt"
56
"io"
67
"math"
78
"net"
@@ -17,8 +18,8 @@ import (
1718
// correctly and so provided in the library.
1819
// See https://github.com/nhooyr/websocket/issues/100.
1920
//
20-
// Every Write to the net.Conn will correspond to a binary message
21-
// write on *webscoket.Conn.
21+
// Every Write to the net.Conn will correspond to a message write of
22+
// the given type on *websocket.Conn.
2223
//
2324
// Close will close the *websocket.Conn with StatusNormalClosure.
2425
//
@@ -30,9 +31,10 @@ import (
3031
// and "websocket/unknown-addr" for String.
3132
//
3233
// A received StatusNormalClosure close frame will be translated to EOF when reading.
33-
func NetConn(c *Conn) net.Conn {
34+
func NetConn(c *Conn, msgType MessageType) net.Conn {
3435
nc := &netConn{
35-
c: c,
36+
c: c,
37+
msgType: msgType,
3638
}
3739

3840
var cancel context.CancelFunc
@@ -52,7 +54,8 @@ func NetConn(c *Conn) net.Conn {
5254
}
5355

5456
type netConn struct {
55-
c *Conn
57+
c *Conn
58+
msgType MessageType
5659

5760
writeTimer *time.Timer
5861
writeContext context.Context
@@ -71,7 +74,7 @@ func (c *netConn) Close() error {
7174
}
7275

7376
func (c *netConn) Write(p []byte) (int, error) {
74-
err := c.c.Write(c.writeContext, MessageBinary, p)
77+
err := c.c.Write(c.writeContext, c.msgType, p)
7578
if err != nil {
7679
return 0, err
7780
}
@@ -93,9 +96,9 @@ func (c *netConn) Read(p []byte) (int, error) {
9396
}
9497
return 0, err
9598
}
96-
if typ != MessageBinary {
97-
c.c.Close(StatusUnsupportedData, "can only accept binary messages")
98-
return 0, xerrors.Errorf("unexpected frame type read for net conn adapter (expected %v): %v", MessageBinary, typ)
99+
if typ != c.msgType {
100+
c.c.Close(StatusUnsupportedData, fmt.Sprintf("can only accept %v messages", c.msgType))
101+
return 0, xerrors.Errorf("unexpected frame type read for net conn adapter (expected %v): %v", c.msgType, typ)
99102
}
100103
c.reader = r
101104
}

websocket_test.go

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -127,7 +127,7 @@ func TestHandshake(t *testing.T) {
127127
}
128128
defer c.Close(websocket.StatusInternalError, "")
129129

130-
nc := websocket.NetConn(c)
130+
nc := websocket.NetConn(c, websocket.MessageBinary)
131131
defer nc.Close()
132132

133133
nc.SetWriteDeadline(time.Time{})
@@ -152,7 +152,7 @@ func TestHandshake(t *testing.T) {
152152
}
153153
defer c.Close(websocket.StatusInternalError, "")
154154

155-
nc := websocket.NetConn(c)
155+
nc := websocket.NetConn(c, websocket.MessageBinary)
156156
defer nc.Close()
157157

158158
nc.SetReadDeadline(time.Time{})

0 commit comments

Comments
 (0)