experimental: UseRequest/ResponseBody methods

experimental/transaction.go

@@ -0,0 +1,36 @@
+// Copyright 2024 Juan Pablo Tosso and the OWASP Coraza contributors
+// SPDX-License-Identifier: Apache-2.0
+
+package experimental
+
+import (
+	"github.com/corazawaf/coraza/v3/types"
+)
+
+type Transaction interface {
+	types.Transaction
+
+	// UseRequestBody directly sets the provided byte slice as the request body buffer.
+	// This is meant to be used when the entire request body is available, as it avoids
+	// the need for an extra copy into the request body buffer. Because of this, this method
+	// is expected to be called just once, further calls to UseRequestBody have to be avoided.
+	// If the body size exceeds the limit and the action is to reject, an interruption will be returned.
+	//
+	// Note: The new internal buffer takes ownership of the provided data, the caller should NOT use b slice
+	// after this call.
+	//
+	// It returns the relevant interruption, the final internal body buffer length and any error that occurs.
+	UseRequestBody(b []byte) (*types.Interruption, int, error)
+
+	// UseResponseBody directly sets the provided byte slice as the response body buffer.
+	// This is meant to be used when the entire response body is available, as it avoids
+	// the need for an extra copy into the response body buffer. Because of this, this method is expected to
+	// be called just once, further calls to UseResponseBody have to be avoided.
+	// If the body size exceeds the limit and the action is to reject, an interruption will be returned.
+	//
+	// Note: The new internal buffer takes ownership of the provided data, the caller should NOT use b slice
+	// after this call.
+	//
+	// It returns the relevant interruption, the final internal body buffer length and any error that occurs.
+	UseResponseBody(b []byte) (*types.Interruption, int, error)
+}

internal/corazawaf/body_buffer.go

@@ -94,6 +94,24 @@ func (br *BodyBuffer) Write(data []byte) (n int, err error) {
 	return br.buffer.Write(data)
 }
 
+// SetBuffer sets the buffer to the provided slice of bytes.
+func (br *BodyBuffer) SetBuffer(data []byte) error {
+	if len(data) == 0 {
+		return errors.New("provided data is empty")
+	}
+
+	// Check if the provided data exceeds the memory limit
+	if int64(len(data)) > br.options.MemoryLimit {
+		return errors.New("memoryLimit reached while writing")
+	}
+
+	// Set the buffer to the provided slice
+	br.buffer = bytes.NewBuffer(data)
+	br.length = int64(len(data))
+
+	return nil
+}
+
 type bodyBufferReader struct {
 	pos int
 	br  *BodyBuffer

internal/corazawaf/transaction.go

@@ -968,6 +968,60 @@ func (tx *Transaction) ReadRequestBodyFrom(r io.Reader) (*types.Interruption, in
 	return tx.interruption, int(w), err
 }
 
+// UseRequestBody directly sets the provided byte slice as the request body buffer.
+// This is meant to be used when the entire request body is available, as it avoids
+// the need for an extra copy into the request body buffer. Because of this, this method
+// is expected to be called just once, further calls to UseRequestBody have to be avoided.
+// If the body size exceeds the limit and the action is to reject, an interruption will be returned.
+// The caller should not use b slice after this call.
+//
+// It returns the relevant interruption, the final internal body buffer length and any error that occurs.
+func (tx *Transaction) UseRequestBody(b []byte) (*types.Interruption, int, error) {
+	if tx.RuleEngine == types.RuleEngineOff {
+		return nil, 0, nil
+	}
+
+	if !tx.RequestBodyAccess {
+		return nil, 0, nil
+	}
+
+	if tx.lastPhase >= types.PhaseRequestBody {
+		return nil, 0, fmt.Errorf("request body buffer set more than once, which has been already been processed")
+	}
+
+	bodySize := int64(len(b))
+	var runProcessRequestBody bool
+
+	if bodySize > tx.RequestBodyLimit {
+		tx.variables.inboundDataError.Set("1")
+		if tx.WAF.RequestBodyLimitAction == types.BodyLimitActionReject {
+			// We interrupt this transaction in case RequestBodyLimitAction is Reject
+			return setAndReturnBodyLimitInterruption(tx)
+		}
+
+		if tx.WAF.RequestBodyLimitAction == types.BodyLimitActionProcessPartial {
+			// Truncate the body slice to the configured limit
+			b = b[:tx.RequestBodyLimit]
+			bodySize = tx.RequestBodyLimit
+			runProcessRequestBody = true
+		}
+	}
+
+	// Point the internal buffer to the provided slice
+	err := tx.requestBodyBuffer.SetBuffer(b)
+	if err != nil {
+		return nil, 0, err
+	}
+
+	err = nil
+	if runProcessRequestBody {
+		tx.debugLogger.Warn().Msg("Processing request body whose size reached the configured limit (Action ProcessPartial)")
+		_, err = tx.ProcessRequestBody()
+	}
+
+	return tx.interruption, int(bodySize), err
+}
+
 // ProcessRequestBody Performs the analysis of the request body (if any)
 //
 // It is recommended to call this method even if it is not expected to have a body.
@@ -1218,6 +1272,59 @@ func (tx *Transaction) ReadResponseBodyFrom(r io.Reader) (*types.Interruption, i
 	return tx.interruption, int(w), err
 }
 
+// UseResponseBody directly sets the provided byte slice as the response body buffer.
+// This is meant to be used when the entire response body is available, as it avoids
+// the need for an extra copy into the response body buffer. Because of this, this method is expected to
+// be called just once, further calls to UseResponseBody have to be avoided.
+// If the body size exceeds the limit and the action is to reject, an interruption will be returned.
+// The caller should not use b slice after this call.
+//
+// It returns the relevant interruption, the final internal body buffer length and any error that occurs.
+func (tx *Transaction) UseResponseBody(b []byte) (*types.Interruption, int, error) {
+	if tx.RuleEngine == types.RuleEngineOff {
+		return nil, 0, nil
+	}
+
+	if !tx.ResponseBodyAccess {
+		return nil, 0, nil
+	}
+
+	if tx.lastPhase >= types.PhaseResponseBody {
+		return nil, 0, fmt.Errorf("response body buffer set more than once, which has been already been processed")
+	}
+
+	var (
+		bodySize               = int64(len(b))
+		runProcessResponseBody = false
+	)
+	if bodySize >= tx.ResponseBodyLimit {
+		tx.variables.outboundDataError.Set("1")
+		if tx.WAF.ResponseBodyLimitAction == types.BodyLimitActionReject {
+			// We interrupt this transaction in case ResponseBodyLimitAction is Reject
+			return setAndReturnBodyLimitInterruption(tx)
+		}
+
+		if tx.WAF.ResponseBodyLimitAction == types.BodyLimitActionProcessPartial {
+			// Truncate the body slice to the configured limit
+			b = b[:tx.ResponseBodyLimit]
+			bodySize = tx.ResponseBodyLimit
+			runProcessResponseBody = true
+		}
+	}
+	// Point the internal buffer to the provided slice
+	err := tx.responseBodyBuffer.SetBuffer(b)
+	if err != nil {
+		return nil, 0, err
+	}
+
+	err = nil
+	if runProcessResponseBody {
+		tx.debugLogger.Debug().Msg("Processing response body whose size reached the configured limit (Action ProcessPartial)")
+		_, err = tx.ProcessResponseBody()
+	}
+	return tx.interruption, int(bodySize), err
+}
+
 // ProcessResponseBody Perform the analysis of the the response body (if any)
 //
 // It is recommended to call this method even if it is not expected to have a body.