Skip to content

Commit

Permalink
errors: allow retrieve rejection status code (gobwas#139)
Browse files Browse the repository at this point in the history
  • Loading branch information
gobwas authored Jul 10, 2021
1 parent d2c3dba commit 272ac76
Show file tree
Hide file tree
Showing 2 changed files with 19 additions and 15 deletions.
30 changes: 17 additions & 13 deletions errors.go
Original file line number Diff line number Diff line change
Expand Up @@ -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
}
4 changes: 2 additions & 2 deletions server.go
Original file line number Diff line number Diff line change
Expand Up @@ -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
}
Expand Down Expand Up @@ -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
}
Expand Down

0 comments on commit 272ac76

Please sign in to comment.