homelab/traefik
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Add mirrorBody option to HTTP mirroring

Browse source
This commit is contained in:
Matteo Paier 2024-09-02 16:36:06 +02:00 • committed by GitHub
parent 51f7f610c9
commit eb99c8c785
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
19 changed files with 165 additions and 22 deletions
Download patch file Download diff file Expand all files Collapse all files

10
pkg/server/service/loadbalancer/mirror/mirror.go
View file

@ -25,6 +25,7 @@ type Mirroring struct {
rw http.ResponseWriter
routinePool *safe.Pool
mirrorBody bool
maxBodySize int64
wantsHealthCheck bool
@ -33,11 +34,12 @@ type Mirroring struct {
}
// New returns a new instance of *Mirroring.
func New(handler http.Handler, pool *safe.Pool, maxBodySize int64, hc *dynamic.HealthCheck) *Mirroring {
func New(handler http.Handler, pool *safe.Pool, mirrorBody bool, maxBodySize int64, hc *dynamic.HealthCheck) *Mirroring {
return &Mirroring{
routinePool: pool,
handler: handler,
rw: blackHoleResponseWriter{},
mirrorBody: mirrorBody,
maxBodySize: maxBodySize,
wantsHealthCheck: hc != nil,
}
@ -83,7 +85,7 @@ func (m *Mirroring) ServeHTTP(rw http.ResponseWriter, req *http.Request) {
}
logger := log.Ctx(req.Context())
rr, bytesRead, err := newReusableRequest(req, m.maxBodySize)
rr, bytesRead, err := newReusableRequest(req, m.mirrorBody, m.maxBodySize)
if err != nil && !errors.Is(err, errBodyTooLarge) {
http.Error(rw, fmt.Sprintf("%s: creating reusable request: %v",
http.StatusText(http.StatusInternalServerError), err), http.StatusInternalServerError)
@ -200,11 +202,11 @@ var errBodyTooLarge = errors.New("request body too large")
// if the returned error is errBodyTooLarge, newReusableRequest also returns the
// bytes that were already consumed from the request's body.
func newReusableRequest(req *http.Request, maxBodySize int64) (*reusableRequest, []byte, error) {
func newReusableRequest(req *http.Request, mirrorBody bool, maxBodySize int64) (*reusableRequest, []byte, error) {
if req == nil {
return nil, nil, errors.New("nil input request")
}
if req.Body == nil || req.ContentLength == 0 {
if req.Body == nil || req.ContentLength == 0 || !mirrorBody {
return &reusableRequest{req: req}, nil, nil
}

66
pkg/server/service/loadbalancer/mirror/mirror_test.go
View file

@ -21,7 +21,7 @@ func TestMirroringOn100(t *testing.T) {
rw.WriteHeader(http.StatusOK)
})
pool := safe.NewPool(context.Background())
mirror := New(handler, pool, defaultMaxBodySize, nil)
mirror := New(handler, pool, true, defaultMaxBodySize, nil)
err := mirror.AddMirror(http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
atomic.AddInt32(&countMirror1, 1)
}), 10)
@ -50,7 +50,7 @@ func TestMirroringOn10(t *testing.T) {
rw.WriteHeader(http.StatusOK)
})
pool := safe.NewPool(context.Background())
mirror := New(handler, pool, defaultMaxBodySize, nil)
mirror := New(handler, pool, true, defaultMaxBodySize, nil)
err := mirror.AddMirror(http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
atomic.AddInt32(&countMirror1, 1)
}), 10)
@ -74,7 +74,7 @@ func TestMirroringOn10(t *testing.T) {
}
func TestInvalidPercent(t *testing.T) {
mirror := New(http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {}), safe.NewPool(context.Background()), defaultMaxBodySize, nil)
mirror := New(http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {}), safe.NewPool(context.Background()), true, defaultMaxBodySize, nil)
err := mirror.AddMirror(nil, -1)
assert.Error(t, err)
@ -93,7 +93,7 @@ func TestHijack(t *testing.T) {
rw.WriteHeader(http.StatusOK)
})
pool := safe.NewPool(context.Background())
mirror := New(handler, pool, defaultMaxBodySize, nil)
mirror := New(handler, pool, true, defaultMaxBodySize, nil)
var mirrorRequest bool
err := mirror.AddMirror(http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
@ -117,7 +117,7 @@ func TestFlush(t *testing.T) {
rw.WriteHeader(http.StatusOK)
})
pool := safe.NewPool(context.Background())
mirror := New(handler, pool, defaultMaxBodySize, nil)
mirror := New(handler, pool, true, defaultMaxBodySize, nil)
var mirrorRequest bool
err := mirror.AddMirror(http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
@ -154,7 +154,7 @@ func TestMirroringWithBody(t *testing.T) {
rw.WriteHeader(http.StatusOK)
})
mirror := New(handler, pool, defaultMaxBodySize, nil)
mirror := New(handler, pool, true, defaultMaxBodySize, nil)
for range numMirrors {
err := mirror.AddMirror(http.HandlerFunc(func(rw http.ResponseWriter, r *http.Request) {
@ -177,13 +177,55 @@ func TestMirroringWithBody(t *testing.T) {
assert.Equal(t, numMirrors, int(val))
}
func TestMirroringWithIgnoredBody(t *testing.T) {
const numMirrors = 10
var (
countMirror int32
body = []byte(`body`)
emptyBody = []byte(``)
)
pool := safe.NewPool(context.Background())
handler := http.HandlerFunc(func(rw http.ResponseWriter, r *http.Request) {
assert.NotNil(t, r.Body)
bb, err := io.ReadAll(r.Body)
assert.NoError(t, err)
assert.Equal(t, body, bb)
rw.WriteHeader(http.StatusOK)
})
mirror := New(handler, pool, false, defaultMaxBodySize, nil)
for range numMirrors {
err := mirror.AddMirror(http.HandlerFunc(func(rw http.ResponseWriter, r *http.Request) {
assert.NotNil(t, r.Body)
bb, err := io.ReadAll(r.Body)
assert.NoError(t, err)
assert.Equal(t, emptyBody, bb)
atomic.AddInt32(&countMirror, 1)
}), 100)
assert.NoError(t, err)
}
req := httptest.NewRequest(http.MethodPost, "/", bytes.NewBuffer(body))
mirror.ServeHTTP(httptest.NewRecorder(), req)
pool.Stop()
val := atomic.LoadInt32(&countMirror)
assert.Equal(t, numMirrors, int(val))
}
func TestCloneRequest(t *testing.T) {
t.Run("http request body is nil", func(t *testing.T) {
req, err := http.NewRequest(http.MethodPost, "/", nil)
assert.NoError(t, err)
ctx := req.Context()
rr, _, err := newReusableRequest(req, defaultMaxBodySize)
rr, _, err := newReusableRequest(req, true, defaultMaxBodySize)
assert.NoError(t, err)
// first call
@ -208,7 +250,7 @@ func TestCloneRequest(t *testing.T) {
ctx := req.Context()
req.ContentLength = int64(contentLength)
rr, _, err := newReusableRequest(req, defaultMaxBodySize)
rr, _, err := newReusableRequest(req, true, defaultMaxBodySize)
assert.NoError(t, err)
// first call
@ -231,7 +273,7 @@ func TestCloneRequest(t *testing.T) {
req, err := http.NewRequest(http.MethodPost, "/", buf)
assert.NoError(t, err)
_, expectedBytes, err := newReusableRequest(req, 2)
_, expectedBytes, err := newReusableRequest(req, true, 2)
assert.Error(t, err)
assert.Equal(t, expectedBytes, bb[:3])
})
@ -243,7 +285,7 @@ func TestCloneRequest(t *testing.T) {
req, err := http.NewRequest(http.MethodPost, "/", buf)
assert.NoError(t, err)
rr, expectedBytes, err := newReusableRequest(req, 20)
rr, expectedBytes, err := newReusableRequest(req, true, 20)
assert.NoError(t, err)
assert.Nil(t, expectedBytes)
assert.Len(t, rr.body, 10)
@ -255,14 +297,14 @@ func TestCloneRequest(t *testing.T) {
req, err := http.NewRequest(http.MethodGet, "/", buf)
assert.NoError(t, err)
rr, expectedBytes, err := newReusableRequest(req, 20)
rr, expectedBytes, err := newReusableRequest(req, true, 20)
assert.NoError(t, err)
assert.Nil(t, expectedBytes)
assert.Empty(t, rr.body)
})
t.Run("no request given", func(t *testing.T) {
_, _, err := newReusableRequest(nil, defaultMaxBodySize)
_, _, err := newReusableRequest(nil, true, defaultMaxBodySize)
assert.Error(t, err)
})
}

12
pkg/server/service/service.go
View file

@ -34,7 +34,10 @@ import (
"google.golang.org/grpc/status"
)
const defaultMaxBodySize int64 = -1
const (
defaultMirrorBody = true
defaultMaxBodySize int64 = -1
)
// RoundTripperGetter is a roundtripper getter interface.
type RoundTripperGetter interface {
@ -197,11 +200,16 @@ func (m *Manager) getMirrorServiceHandler(ctx context.Context, config *dynamic.M
return nil, err
}
mirrorBody := defaultMirrorBody
if config.MirrorBody != nil {
mirrorBody = *config.MirrorBody
}
maxBodySize := defaultMaxBodySize
if config.MaxBodySize != nil {
maxBodySize = *config.MaxBodySize
}
handler := mirror.New(serviceHandler, m.routinePool, maxBodySize, config.HealthCheck)
handler := mirror.New(serviceHandler, m.routinePool, mirrorBody, maxBodySize, config.HealthCheck)
for _, mirrorConfig := range config.Mirrors {
mirrorHandler, err := m.BuildHTTP(ctx, mirrorConfig.Name)
if err != nil {