errors: allow retrieve rejection status code (#139)

Fixes #136

Author: gobwas

errors.go

@@ -2,54 +2,58 @@ package ws
 
 // RejectOption represents an option used to control the way connection is
 // rejected.
-type RejectOption func(*rejectConnectionError)
+type RejectOption func(*ConnectionRejectedError)
 
 // RejectionReason returns an option that makes connection to be rejected with
 // given reason.
 func RejectionReason(reason string) RejectOption {
-	return func(err *rejectConnectionError) {
+	return func(err *ConnectionRejectedError) {
 		err.reason = reason
 	}
 }
 
 // RejectionStatus returns an option that makes connection to be rejected with
 // given HTTP status code.
 func RejectionStatus(code int) RejectOption {
-	return func(err *rejectConnectionError) {
+	return func(err *ConnectionRejectedError) {
 		err.code = code
 	}
 }
 
 // RejectionHeader returns an option that makes connection to be rejected with
 // given HTTP headers.
 func RejectionHeader(h HandshakeHeader) RejectOption {
-	return func(err *rejectConnectionError) {
+	return func(err *ConnectionRejectedError) {
 		err.header = h
 	}
 }
 
-// RejectConnectionError constructs an error that could be used to control the way
-// handshake is rejected by Upgrader.
+// RejectConnectionError constructs an error that could be used to control the
+// way handshake is rejected by Upgrader.
 func RejectConnectionError(options ...RejectOption) error {
-	err := new(rejectConnectionError)
+	err := new(ConnectionRejectedError)
 	for _, opt := range options {
 		opt(err)
 	}
 	return err
 }
 
-// rejectConnectionError represents a rejection of connection during WebSocket
-// handshake error.
+// ConnectionRejectedError represents a rejection of connection during
+// WebSocket handshake error.
 //
-// It can be returned by Upgrader's On* hooks to control the way WebSocket
-// handshake is rejected.
-type rejectConnectionError struct {
+// It can be returned by Upgrader's On* hooks to indicate that WebSocket
+// handshake should be rejected.
+type ConnectionRejectedError struct {
 	reason string
 	code   int
 	header HandshakeHeader
 }
 
 // Error implements error interface.
-func (r *rejectConnectionError) Error() string {
+func (r *ConnectionRejectedError) Error() string {
 	return r.reason
 }
+
+func (r *ConnectionRejectedError) StatusCode() int {
+	return r.code
+}

server.go

@@ -251,7 +251,7 @@ func (u HTTPUpgrader) Upgrade(r *http.Request, w http.ResponseWriter) (conn net.
 		err = rw.Writer.Flush()
 	} else {
 		var code int
-		if rej, ok := err.(*rejectConnectionError); ok {
+		if rej, ok := err.(*ConnectionRejectedError); ok {
 			code = rej.code
 			header[1] = rej.header
 		}
@@ -630,7 +630,7 @@ func (u Upgrader) Upgrade(conn io.ReadWriter) (hs Handshake, err error) {
 	}
 	if err != nil {
 		var code int
-		if rej, ok := err.(*rejectConnectionError); ok {
+		if rej, ok := err.(*ConnectionRejectedError); ok {
 			code = rej.code
 			header[1] = rej.header
 		}