66 "errors"
77 "fmt"
88 "log"
9- "nhooyr.io/websocket/internal/wsframe"
9+ "nhooyr.io/websocket/internal/bufpool"
10+ "time"
1011)
1112
1213// StatusCode represents a WebSocket status code.
@@ -75,21 +76,101 @@ func CloseStatus(err error) StatusCode {
7576 return - 1
7677}
7778
79+ // Close closes the WebSocket connection with the given status code and reason.
80+ //
81+ // It will write a WebSocket close frame with a timeout of 5s and then wait 5s for
82+ // the peer to send a close frame.
83+ // Thus, it implements the full WebSocket close handshake.
84+ // All data messages received from the peer during the close handshake
85+ // will be discarded.
86+ //
87+ // The connection can only be closed once. Additional calls to Close
88+ // are no-ops.
89+ //
90+ // The maximum length of reason must be 125 bytes otherwise an internal
91+ // error will be sent to the peer. For this reason, you should avoid
92+ // sending a dynamic reason.
93+ //
94+ // Close will unblock all goroutines interacting with the connection once
95+ // complete.
96+ func (c * Conn ) Close (code StatusCode , reason string ) error {
97+ err := c .closeHandshake (code , reason )
98+ if err != nil {
99+ return fmt .Errorf ("failed to close websocket: %w" , err )
100+ }
101+ return nil
102+ }
103+
104+ func (c * Conn ) closeHandshake (code StatusCode , reason string ) error {
105+ err := c .cw .sendClose (code , reason )
106+ if err != nil {
107+ return err
108+ }
109+
110+ return c .cr .waitClose ()
111+ }
112+
113+ func (cw * connWriter ) error (code StatusCode , err error ) {
114+ cw .c .setCloseErr (err )
115+ cw .sendClose (code , err .Error ())
116+ cw .c .close (nil )
117+ }
118+
119+ func (cw * connWriter ) sendClose (code StatusCode , reason string ) error {
120+ ce := CloseError {
121+ Code : code ,
122+ Reason : reason ,
123+ }
124+
125+ cw .c .setCloseErr (fmt .Errorf ("sent close frame: %w" , ce ))
126+
127+ var p []byte
128+ if ce .Code != StatusNoStatusRcvd {
129+ p = ce .bytes ()
130+ }
131+
132+ return cw .control (context .Background (), opClose , p )
133+ }
134+
135+ func (cr * connReader ) waitClose () error {
136+ defer cr .c .close (nil )
137+
138+ return nil
139+
140+ ctx , cancel := context .WithTimeout (context .Background (), time .Second * 5 )
141+ defer cancel ()
142+
143+ err := cr .mu .Lock (ctx )
144+ if err != nil {
145+ return err
146+ }
147+ defer cr .mu .Unlock ()
148+
149+ b := bufpool .Get ()
150+ buf := b .Bytes ()
151+ buf = buf [:cap (buf )]
152+ defer bufpool .Put (b )
153+
154+ for {
155+ // TODO
156+ return nil
157+ }
158+ }
159+
78160func parseClosePayload (p []byte ) (CloseError , error ) {
79161 if len (p ) == 0 {
80162 return CloseError {
81163 Code : StatusNoStatusRcvd ,
82164 }, nil
83165 }
84166
85- code , reason , err := wsframe .ParseClosePayload (p )
86- if err != nil {
87- return CloseError {}, err
167+ if len (p ) < 2 {
168+ return CloseError {}, fmt .Errorf ("close payload %q too small, cannot even contain the 2 byte status code" , p )
88169 }
89170
90171 ce := CloseError {
91- Code : StatusCode (code ),
92- Reason : reason ,
172+ Code : StatusCode (binary . BigEndian . Uint16 ( p ) ),
173+ Reason : string ( p [ 2 :]) ,
93174 }
94175
95176 if ! validWireCloseCode (ce .Code ) {
@@ -129,11 +210,13 @@ func (ce CloseError) bytes() []byte {
129210 return p
130211}
131212
213+ const maxCloseReason = maxControlPayload - 2
214+
132215func (ce CloseError ) bytesErr () ([]byte , error ) {
133- const maxReason = maxControlPayload - 2
134- if len (ce .Reason ) > maxReason {
135- return nil , fmt .Errorf ("reason string max is %v but got %q with length %v" , maxReason , ce .Reason , len (ce .Reason ))
216+ if len (ce .Reason ) > maxCloseReason {
217+ return nil , fmt .Errorf ("reason string max is %v but got %q with length %v" , maxCloseReason , ce .Reason , len (ce .Reason ))
136218 }
219+
137220 if ! validWireCloseCode (ce .Code ) {
138221 return nil , fmt .Errorf ("status code %v cannot be set" , ce .Code )
139222 }
@@ -144,34 +227,6 @@ func (ce CloseError) bytesErr() ([]byte, error) {
144227 return buf , nil
145228}
146229
147- // CloseRead will start a goroutine to read from the connection until it is closed or a data message
148- // is received. If a data message is received, the connection will be closed with StatusPolicyViolation.
149- // Since CloseRead reads from the connection, it will respond to ping, pong and close frames.
150- // After calling this method, you cannot read any data messages from the connection.
151- // The returned context will be cancelled when the connection is closed.
152- //
153- // Use this when you do not want to read data messages from the connection anymore but will
154- // want to write messages to it.
155- func (c * Conn ) CloseRead (ctx context.Context ) context.Context {
156- ctx , cancel := context .WithCancel (ctx )
157- go func () {
158- defer cancel ()
159- c .Reader (ctx )
160- c .Close (StatusPolicyViolation , "unexpected data message" )
161- }()
162- return ctx
163- }
164-
165- // SetReadLimit sets the max number of bytes to read for a single message.
166- // It applies to the Reader and Read methods.
167- //
168- // By default, the connection has a message read limit of 32768 bytes.
169- //
170- // When the limit is hit, the connection will be closed with StatusMessageTooBig.
171- func (c * Conn ) SetReadLimit (n int64 ) {
172- c .r .lr .limit .Store (n )
173- }
174-
175230func (c * Conn ) setCloseErr (err error ) {
176231 c .closeMu .Lock ()
177232 c .setCloseErrNoLock (err )
0 commit comments