mirror of
https://github.com/coder/coder.git
synced 2026-09-24 15:04:27 +08:00
feat: add rbac tracing (#4093)
This commit is contained in:
@@ -1,74 +0,0 @@
|
||||
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
|
||||
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
|
||||
}
|
||||
@@ -1,129 +0,0 @@
|
||||
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")
|
||||
}
|
||||
Reference in New Issue
Block a user