github.com/epsagon/epsagon-go@v1.39.0/wrappers/net/http/response_writer.go (about)

     1  package epsagonhttp
     2  
     3  import (
     4  	"bytes"
     5  	"fmt"
     6  	"net/http"
     7  
     8  	"github.com/epsagon/epsagon-go/epsagon"
     9  	"github.com/epsagon/epsagon-go/protocol"
    10  )
    11  
    12  // WrappedResponseWriter is wrapping Resposne writer with Epsagon
    13  // to enrich the trace with data from the response
    14  type WrappedResponseWriter struct {
    15  	http.ResponseWriter
    16  	resource *protocol.Resource
    17  	buf      bytes.Buffer
    18  }
    19  
    20  // CreateWrappedResponseWriter creates a newWrappedResponseWriter
    21  func CreateWrappedResponseWriter(rw http.ResponseWriter, resource *protocol.Resource) *WrappedResponseWriter {
    22  	return &WrappedResponseWriter{
    23  		ResponseWriter: rw,
    24  		resource:       resource,
    25  		buf:            bytes.Buffer{},
    26  	}
    27  }
    28  
    29  // Header wrapper
    30  func (w *WrappedResponseWriter) Header() http.Header {
    31  	return w.ResponseWriter.Header()
    32  }
    33  
    34  // WriteHeader wrapper, will set status_code immediately
    35  func (w *WrappedResponseWriter) WriteHeader(statusCode int) {
    36  	w.resource.Metadata["status_code"] = fmt.Sprint(statusCode)
    37  	w.ResponseWriter.WriteHeader(statusCode)
    38  }
    39  
    40  // Write wrapper
    41  func (w *WrappedResponseWriter) Write(data []byte) (int, error) {
    42  	w.buf.Write(data)
    43  	return w.ResponseWriter.Write(data)
    44  }
    45  
    46  // UpdateResource updates the connected resource with the response headers and body
    47  func (w *WrappedResponseWriter) UpdateResource() {
    48  	w.resource.Metadata["response_headers"], _ = epsagon.FormatHeaders(w.ResponseWriter.Header())
    49  	w.resource.Metadata["response_body"] = w.buf.String()
    50  }