response_writer.go

package coraza

import (
	"io"

	"github.com/corazawaf/coraza/v2"
	"github.com/corazawaf/coraza/v2/types"
	"github.com/gin-gonic/gin"
)

type responseWriter struct {
	gin.ResponseWriter
	tx *coraza.Transaction

	headersProcessed bool
	size             int
}

func (w responseWriter) Write(b []byte) (n int, err error) {
	if it := w.processResponseHeaders(); it != nil {
		// transaction was interrupted :(
		return
	}
	w.WriteHeaderNow()
	n, err = w.tx.ResponseBodyBuffer.Write(b)
	w.size += n
	return
}

func (w *responseWriter) WriteString(s string) (n int, err error) {
	if it := w.processResponseHeaders(); it != nil {
		// transaction was interrupted :(
		return
	}
	w.WriteHeaderNow()
	n, err = io.WriteString(w.tx.ResponseBodyBuffer, s)
	w.size += n
	return
}

func (w *responseWriter) processResponseHeaders() *types.Interruption {
	if w.headersProcessed || w.tx.Interruption != nil {
		return w.tx.Interruption
	}
	w.headersProcessed = true
	for k, vv := range w.ResponseWriter.Header() {
		for _, v := range vv {
			w.tx.AddResponseHeader(k, v)
		}
	}
	return w.tx.ProcessResponseHeaders(w.ResponseWriter.Status(), "http/1.1")
}

func (w *responseWriter) Status() int {
	if w.tx.Interruption != nil {
		return w.tx.Interruption.Status
	}
	return w.ResponseWriter.Status()
}

func (w *responseWriter) Size() int {
	return w.size
}