From d521901f820f35ee86196017961f9e85f41315e6 Mon Sep 17 00:00:00 2001 From: Dylan Cant Date: Sun, 5 Feb 2023 17:20:32 -0500 Subject: [PATCH 1/3] added a Hijack Route for upgrade --- server.go | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) diff --git a/server.go b/server.go index ce530ae..5ded5ff 100644 --- a/server.go +++ b/server.go @@ -54,6 +54,9 @@ var ( RejectionStatus(http.StatusBadRequest), RejectionReason(fmt.Sprintf("handshake error: bad %q header", headerSecVersion)), ) + ErrHandshakeHijack = RejectConnectionError( + RejectionStatus(http.StatusTeapot), + ) ) // ErrMalformedResponse is returned by Dialer to indicate that server response @@ -340,6 +343,8 @@ type Upgrader struct { // sent with appropriate HTTP error code and body set to error message. // // RejectConnectionError could be used to get more control on response. + // + // Use ErrHandshakeHijack to return without writing to connection Negotiate func(httphead.Option) (httphead.Option, error) // Header is an optional HandshakeHeader instance that could be used to @@ -360,6 +365,8 @@ type Upgrader struct { // sent with appropriate HTTP error code and body set to error message. // // RejectConnectionError could be used to get more control on response. + // + // Use ErrHandshakeHijack to return without writing to connection OnRequest func(uri []byte) error // OnHost is a callback that will be called after "Host" header successful @@ -375,6 +382,8 @@ type Upgrader struct { // sent with appropriate HTTP error code and body set to error message. // // RejectConnectionError could be used to get more control on response. + // + // Use ErrHandshakeHijack to return without writing to connection OnHost func(host []byte) error // OnHeader is a callback that will be called after successful parsing of @@ -388,6 +397,8 @@ type Upgrader struct { // sent with appropriate HTTP error code and body set to error message. // // RejectConnectionError could be used to get more control on response. + // + // Use ErrHandshakeHijack to return without writing to connection OnHeader func(key, value []byte) error // OnBeforeUpgrade is a callback that will be called before sending @@ -403,6 +414,8 @@ type Upgrader struct { // sent with appropriate HTTP error code and body set to error message. // // RejectConnectionError could be used to get more control on response. + // + // Use ErrHandshakeHijack to return without writing to connection OnBeforeUpgrade func() (header HandshakeHeader, err error) } @@ -632,6 +645,11 @@ func (u Upgrader) Upgrade(conn io.ReadWriter) (hs Handshake, err error) { var code int if rej, ok := err.(*ConnectionRejectedError); ok { code = rej.code + if code == http.StatusTeapot { + // Hijacked! Implementation knows what they are doing. + // Return current err without writing to connection. + return + } header[1] = rej.header } if code == 0 { From 30f8677082e031ef12bc1c87c9bf78dac45fbbec Mon Sep 17 00:00:00 2001 From: Dylan Cant Date: Sat, 18 Feb 2023 19:39:20 -0500 Subject: [PATCH 2/3] fixed escape hatch --- errors.go | 4 ++++ server.go | 40 +++++++++++++++++----------------------- 2 files changed, 21 insertions(+), 23 deletions(-) diff --git a/errors.go b/errors.go index f5668b2..e66aee3 100644 --- a/errors.go +++ b/errors.go @@ -57,3 +57,7 @@ func (r *ConnectionRejectedError) Error() string { func (r *ConnectionRejectedError) StatusCode() int { return r.code } + +type EscapeHatch interface { + Escape() bool +} diff --git a/server.go b/server.go index 5ded5ff..6379c0c 100644 --- a/server.go +++ b/server.go @@ -54,9 +54,6 @@ var ( RejectionStatus(http.StatusBadRequest), RejectionReason(fmt.Sprintf("handshake error: bad %q header", headerSecVersion)), ) - ErrHandshakeHijack = RejectConnectionError( - RejectionStatus(http.StatusTeapot), - ) ) // ErrMalformedResponse is returned by Dialer to indicate that server response @@ -343,8 +340,6 @@ type Upgrader struct { // sent with appropriate HTTP error code and body set to error message. // // RejectConnectionError could be used to get more control on response. - // - // Use ErrHandshakeHijack to return without writing to connection Negotiate func(httphead.Option) (httphead.Option, error) // Header is an optional HandshakeHeader instance that could be used to @@ -365,8 +360,6 @@ type Upgrader struct { // sent with appropriate HTTP error code and body set to error message. // // RejectConnectionError could be used to get more control on response. - // - // Use ErrHandshakeHijack to return without writing to connection OnRequest func(uri []byte) error // OnHost is a callback that will be called after "Host" header successful @@ -382,9 +375,7 @@ type Upgrader struct { // sent with appropriate HTTP error code and body set to error message. // // RejectConnectionError could be used to get more control on response. - // - // Use ErrHandshakeHijack to return without writing to connection - OnHost func(host []byte) error + OnHost func(err error, host []byte) error // OnHeader is a callback that will be called after successful parsing of // header, that is not used during WebSocket handshake procedure. That is, @@ -397,9 +388,7 @@ type Upgrader struct { // sent with appropriate HTTP error code and body set to error message. // // RejectConnectionError could be used to get more control on response. - // - // Use ErrHandshakeHijack to return without writing to connection - OnHeader func(key, value []byte) error + OnHeader func(err error, key, value []byte) error // OnBeforeUpgrade is a callback that will be called before sending // successful upgrade response. @@ -414,8 +403,6 @@ type Upgrader struct { // sent with appropriate HTTP error code and body set to error message. // // RejectConnectionError could be used to get more control on response. - // - // Use ErrHandshakeHijack to return without writing to connection OnBeforeUpgrade func() (header HandshakeHeader, err error) } @@ -516,7 +503,7 @@ func (u Upgrader) Upgrade(conn io.ReadWriter) (hs Handshake, err error) { nonce = make([]byte, nonceSize) ) - for err == nil { + for { line, e := readLine(br) if e != nil { return hs, e @@ -536,31 +523,35 @@ func (u Upgrader) Upgrade(conn io.ReadWriter) (hs Handshake, err error) { case headerHostCanonical: headerSeen |= headerSeenHost if onHost := u.OnHost; onHost != nil { - err = onHost(v) + err = onHost(err, v) } case headerUpgradeCanonical: headerSeen |= headerSeenUpgrade if !bytes.Equal(v, specHeaderValueUpgrade) && !bytes.EqualFold(v, specHeaderValueUpgrade) { err = ErrHandshakeBadUpgrade + break } case headerConnectionCanonical: headerSeen |= headerSeenConnection if !bytes.Equal(v, specHeaderValueConnection) && !btsHasToken(v, specHeaderValueConnectionLower) { err = ErrHandshakeBadConnection + break } case headerSecVersionCanonical: headerSeen |= headerSeenSecVersion if !bytes.Equal(v, specHeaderValueSecVersion) { err = ErrHandshakeUpgradeRequired + break } case headerSecKeyCanonical: headerSeen |= headerSeenSecKey if len(v) != nonceSize { err = ErrHandshakeBadSecKey + break } else { copy(nonce[:], v) } @@ -575,12 +566,16 @@ func (u Upgrader) Upgrade(conn io.ReadWriter) (hs Handshake, err error) { } if !ok { err = ErrMalformedRequest + break } } case headerSecExtensionsCanonical: if f := u.Negotiate; err == nil && f != nil { hs.Extensions, err = negotiateExtensions(v, hs.Extensions, f) + if err != nil { + break + } } // DEPRECATED path. if custom, check := u.ExtensionCustom, u.Extension; u.Negotiate == nil && (custom != nil || check != nil) { @@ -592,12 +587,13 @@ func (u Upgrader) Upgrade(conn io.ReadWriter) (hs Handshake, err error) { } if !ok { err = ErrMalformedRequest + break } } default: if onHeader := u.OnHeader; onHeader != nil { - err = onHeader(k, v) + err = onHeader(err, k, v) } } } @@ -643,13 +639,11 @@ func (u Upgrader) Upgrade(conn io.ReadWriter) (hs Handshake, err error) { } if err != nil { var code int + if t, ok := err.(EscapeHatch); ok && t.Escape() { + return + } if rej, ok := err.(*ConnectionRejectedError); ok { code = rej.code - if code == http.StatusTeapot { - // Hijacked! Implementation knows what they are doing. - // Return current err without writing to connection. - return - } header[1] = rej.header } if code == 0 { From 0928bdfc7eab81e3aa294fe30cd53a1bd8d4bffc Mon Sep 17 00:00:00 2001 From: dylan cant Date: Thu, 1 Jun 2023 15:03:30 -0400 Subject: [PATCH 3/3] changed upgrade handler --- server.go | 7 ++----- 1 file changed, 2 insertions(+), 5 deletions(-) diff --git a/server.go b/server.go index 0771ac4..bb84f80 100644 --- a/server.go +++ b/server.go @@ -534,10 +534,8 @@ func (u Upgrader) Upgrade(conn io.ReadWriter) (hs Handshake, err error) { } case headerConnectionCanonical: - headerSeen |= headerSeenConnection - if !bytes.Equal(v, specHeaderValueConnection) && !btsHasToken(v, specHeaderValueConnectionLower) { - err = ErrHandshakeBadConnection - break + if bytes.Equal(v, specHeaderValueConnection) || btsHasToken(v, specHeaderValueConnectionLower) { + headerSeen |= headerSeenConnection } case headerSecVersionCanonical: @@ -633,7 +631,6 @@ func (u Upgrader) Upgrade(conn io.ReadWriter) (hs Handshake, err error) { default: panic("unknown headers state") } - case err == nil && u.OnBeforeUpgrade != nil: header[1], err = u.OnBeforeUpgrade() }