feat: add panic recovery middleware #3687
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,30 @@ | ||
| package httpapi | ||
|
|
||
| import "net/http" | ||
|
|
||
| const ( | ||
| // XForwardedHostHeader is a header used by proxies to indicate the | ||
| // original host of the request. | ||
| XForwardedHostHeader = "X-Forwarded-Host" | ||
| ) | ||
|
|
||
| // RequestHost returns the name of the host from the request. It prioritizes | ||
| // 'X-Forwarded-Host' over r.Host since most requests are being proxied. | ||
| func RequestHost(r *http.Request) string { | ||
| host := r.Header.Get(XForwardedHostHeader) | ||
| if host != "" { | ||
| return host | ||
| } | ||
|
|
||
| return r.Host | ||
| } | ||
|
|
||
| func IsWebsocketUpgrade(r *http.Request) bool { | ||
| vs := r.Header.Values("Upgrade") | ||
| for _, v := range vs { | ||
| if v == "websocket" { | ||
| return true | ||
| } | ||
| } | ||
| return false | ||
| } |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change | ||||||
|---|---|---|---|---|---|---|---|---|
| @@ -0,0 +1,74 @@ | ||||||||
| package httpapi | ||||||||
|
|
||||||||
| import ( | ||||||||
| "bufio" | ||||||||
| "net" | ||||||||
| "net/http" | ||||||||
|
|
||||||||
| "golang.org/x/xerrors" | ||||||||
| ) | ||||||||
|
|
||||||||
| var _ http.ResponseWriter = (*StatusWriter)(nil) | ||||||||
| var _ http.Hijacker = (*StatusWriter)(nil) | ||||||||
|
|
||||||||
| // StatusWriter intercepts the status of the request and the response body up | ||||||||
| // to maxBodySize if Status >= 400. It is guaranteed to be the ResponseWriter | ||||||||
| // directly downstream from Middleware. | ||||||||
| type StatusWriter struct { | ||||||||
| http.ResponseWriter | ||||||||
| Status int | ||||||||
| Hijacked bool | ||||||||
| responseBody []byte | ||||||||
|
|
||||||||
| wroteHeader bool | ||||||||
| } | ||||||||
|
|
||||||||
| func (w *StatusWriter) WriteHeader(status int) { | ||||||||
| if !w.wroteHeader { | ||||||||
| w.Status = status | ||||||||
| w.wroteHeader = true | ||||||||
| } | ||||||||
| w.ResponseWriter.WriteHeader(status) | ||||||||
| } | ||||||||
|
|
||||||||
| func (w *StatusWriter) Write(b []byte) (int, error) { | ||||||||
| const maxBodySize = 4096 | ||||||||
|
|
||||||||
| if !w.wroteHeader { | ||||||||
| w.Status = http.StatusOK | ||||||||
|
There was a problem hiding this comment. Technically this will be true once we hit
Suggested change
|
||||||||
| w.wroteHeader = true | ||||||||
| } | ||||||||
|
|
||||||||
| if w.Status >= http.StatusBadRequest { | ||||||||
| // This is technically wrong as multiple calls to write | ||||||||
| // will simply overwrite w.ResponseBody but given that | ||||||||
| // we typically only write to the response body once | ||||||||
| // and this field is only used for logging I'm leaving | ||||||||
| // this as-is. | ||||||||
| w.responseBody = make([]byte, minInt(len(b), maxBodySize)) | ||||||||
| copy(w.responseBody, b) | ||||||||
| } | ||||||||
|
|
||||||||
| return w.ResponseWriter.Write(b) | ||||||||
| } | ||||||||
|
|
||||||||
| func minInt(a, b int) int { | ||||||||
| if a < b { | ||||||||
| return a | ||||||||
| } | ||||||||
| return b | ||||||||
| } | ||||||||
|
|
||||||||
| func (w *StatusWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) { | ||||||||
| hijacker, ok := w.ResponseWriter.(http.Hijacker) | ||||||||
| if !ok { | ||||||||
| return nil, nil, xerrors.Errorf("%T is not a http.Hijacker", w.ResponseWriter) | ||||||||
| } | ||||||||
| w.Hijacked = true | ||||||||
|
|
||||||||
| return hijacker.Hijack() | ||||||||
| } | ||||||||
|
|
||||||||
| func (w *StatusWriter) ResponseBody() []byte { | ||||||||
| return w.responseBody | ||||||||
| } | ||||||||
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,129 @@ | ||
| package httpapi_test | ||
|
|
||
| import ( | ||
| "bufio" | ||
| "crypto/rand" | ||
| "net" | ||
| "net/http" | ||
| "net/http/httptest" | ||
| "testing" | ||
|
|
||
| "github.com/stretchr/testify/require" | ||
| "golang.org/x/xerrors" | ||
|
|
||
| "github.com/coder/coder/coderd/httpapi" | ||
| ) | ||
|
|
||
| func TestStatusWriter(t *testing.T) { | ||
| t.Parallel() | ||
|
|
||
| t.Run("WriteHeader", func(t *testing.T) { | ||
| t.Parallel() | ||
|
|
||
| var ( | ||
| rec = httptest.NewRecorder() | ||
| w = &httpapi.StatusWriter{ResponseWriter: rec} | ||
| ) | ||
|
|
||
| w.WriteHeader(http.StatusOK) | ||
| require.Equal(t, http.StatusOK, w.Status) | ||
| // Validate that the code is written to the underlying Response. | ||
| require.Equal(t, http.StatusOK, rec.Code) | ||
| }) | ||
|
|
||
| t.Run("WriteHeaderTwice", func(t *testing.T) { | ||
| t.Parallel() | ||
|
|
||
| var ( | ||
| rec = httptest.NewRecorder() | ||
| w = &httpapi.StatusWriter{ResponseWriter: rec} | ||
| code = http.StatusNotFound | ||
| ) | ||
|
|
||
| w.WriteHeader(code) | ||
| w.WriteHeader(http.StatusOK) | ||
| // Validate that we only record the first status code. | ||
| require.Equal(t, code, w.Status) | ||
| // Validate that the code is written to the underlying Response. | ||
| require.Equal(t, code, rec.Code) | ||
| }) | ||
|
|
||
| t.Run("WriteNoHeader", func(t *testing.T) { | ||
| t.Parallel() | ||
| var ( | ||
| rec = httptest.NewRecorder() | ||
| w = &httpapi.StatusWriter{ResponseWriter: rec} | ||
| body = []byte("hello") | ||
| ) | ||
|
|
||
| _, err := w.Write(body) | ||
| require.NoError(t, err) | ||
|
|
||
| // Should set the status to OK. | ||
| require.Equal(t, http.StatusOK, w.Status) | ||
| // We don't record the body for codes <400. | ||
| require.Equal(t, []byte(nil), w.ResponseBody()) | ||
| require.Equal(t, body, rec.Body.Bytes()) | ||
| }) | ||
|
|
||
| t.Run("WriteAfterHeader", func(t *testing.T) { | ||
| t.Parallel() | ||
| var ( | ||
| rec = httptest.NewRecorder() | ||
| w = &httpapi.StatusWriter{ResponseWriter: rec} | ||
| body = []byte("hello") | ||
| code = http.StatusInternalServerError | ||
| ) | ||
|
|
||
| w.WriteHeader(code) | ||
| _, err := w.Write(body) | ||
| require.NoError(t, err) | ||
|
|
||
| require.Equal(t, code, w.Status) | ||
| require.Equal(t, body, w.ResponseBody()) | ||
| require.Equal(t, body, rec.Body.Bytes()) | ||
| }) | ||
|
|
||
| t.Run("WriteMaxBody", func(t *testing.T) { | ||
| t.Parallel() | ||
| var ( | ||
| rec = httptest.NewRecorder() | ||
| w = &httpapi.StatusWriter{ResponseWriter: rec} | ||
| // 8kb body. | ||
| body = make([]byte, 8<<10) | ||
| code = http.StatusInternalServerError | ||
| ) | ||
|
|
||
| _, err := rand.Read(body) | ||
| require.NoError(t, err) | ||
|
|
||
| w.WriteHeader(code) | ||
| _, err = w.Write(body) | ||
| require.NoError(t, err) | ||
|
|
||
| require.Equal(t, code, w.Status) | ||
| require.Equal(t, body, rec.Body.Bytes()) | ||
| require.Equal(t, body[:4096], w.ResponseBody()) | ||
| }) | ||
|
|
||
| t.Run("Hijack", func(t *testing.T) { | ||
| t.Parallel() | ||
| var ( | ||
| rec = httptest.NewRecorder() | ||
| ) | ||
|
|
||
| w := &httpapi.StatusWriter{ResponseWriter: hijacker{rec}} | ||
|
|
||
| _, _, err := w.Hijack() | ||
| require.Error(t, err) | ||
| require.Equal(t, "hijacked", err.Error()) | ||
| }) | ||
| } | ||
|
|
||
| type hijacker struct { | ||
| http.ResponseWriter | ||
| } | ||
|
|
||
| func (hijacker) Hijack() (net.Conn, *bufio.ReadWriter, error) { | ||
| return nil, nil, xerrors.New("hijacked") | ||
| } |
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Suggestion (non-blocking): I think it could be beneficial to make this a package-private type and expose the fields via functions.
Benefit is that a caller can't access the underlying ResponseWriter (bypassing StatusWriter), but can still test for the values via
if s, ok := rw.(interface{Status() int}); ok { s.Status() }There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I'm not sure I'm following. Do you just want to create a
StatusWriterinterface and unexport the current struct? Alternatively we could simply unexport theResponseWriteron the struct.There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Ah, sorry it was unclear, I was suggesting two changes:
ResponseWriter=>responseWriterStatus=>func(*statusWriter) Status() intNumber 1 was the main thing (to protect from direct ResponseWriter access), number 2 was a little extra which would allow the values to be inspected outside the package, if the need ever arises. But not necessary by any means.
In the case of 2, my example of
if s, ok := rw.(interface{Status() int}); ok { s.Status() }was hinting at that we don't need to export anything from this package. If that need ever arises (inspect value outside package), it could be done with an inline interface.I could've omitted that part though, I don't think it will be needed. 😄