Skip to content

Commit

Permalink
use http.ResposnseController
Browse files Browse the repository at this point in the history
  • Loading branch information
ninedraft authored and AlexVulaj committed Jan 22, 2024
1 parent 4a5e66f commit 1e975a0
Show file tree
Hide file tree
Showing 2 changed files with 26 additions and 6 deletions.
10 changes: 4 additions & 6 deletions server.go
Original file line number Diff line number Diff line change
Expand Up @@ -174,13 +174,11 @@ func (u *Upgrader) Upgrade(w http.ResponseWriter, r *http.Request, responseHeade
}
}

h, ok := w.(http.Hijacker)
if !ok {
netConn, brw, err := http.NewResponseController(w).Hijack()
switch {
case errors.Is(err, errors.ErrUnsupported):
return u.returnError(w, r, http.StatusInternalServerError, "websocket: response does not implement http.Hijacker")
}
var brw *bufio.ReadWriter
netConn, brw, err := h.Hijack()
if err != nil {
case err != nil:
return u.returnError(w, r, http.StatusInternalServerError, err.Error())
}

Expand Down
22 changes: 22 additions & 0 deletions server_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -7,8 +7,10 @@ package websocket
import (
"bufio"
"bytes"
"errors"
"net"
"net/http"
"net/http/httptest"
"reflect"
"strings"
"testing"
Expand Down Expand Up @@ -152,3 +154,23 @@ func TestBufioReuse(t *testing.T) {
}
}
}

func TestHijack_NotSupported(t *testing.T) {
t.Parallel()

req := httptest.NewRequest(http.MethodGet, "http://example.com", nil)
req.Header.Set("Upgrade", "websocket")
req.Header.Set("Connection", "upgrade")
req.Header.Set("Sec-Websocket-Key", "dGhlIHNhbXBsZSBub25jZQ==")
req.Header.Set("Sec-Websocket-Version", "13")

recorder := httptest.NewRecorder()

upgrader := Upgrader{}
_, err := upgrader.Upgrade(recorder, req, nil)

if want := (HandshakeError{}); !errors.As(err, &want) || recorder.Code != http.StatusInternalServerError {
t.Errorf("want %T and status_code=%d", want, http.StatusInternalServerError)
t.Fatalf("got err=%T and status_code=%d", err, recorder.Code)
}
}

0 comments on commit 1e975a0

Please sign in to comment.