github.com/secoba/wails/v2@v2.6.4/pkg/assetserver/body_recorder.go (about)

     1  package assetserver
     2  
     3  import (
     4  	"bytes"
     5  	"net/http"
     6  )
     7  
     8  type bodyRecorder struct {
     9  	http.ResponseWriter
    10  	doRecord func(code int, header http.Header) bool
    11  
    12  	body        *bytes.Buffer
    13  	code        int
    14  	wroteHeader bool
    15  }
    16  
    17  func (rw *bodyRecorder) Write(buf []byte) (int, error) {
    18  	rw.writeHeader(buf, http.StatusOK)
    19  	if rw.body != nil {
    20  		return rw.body.Write(buf)
    21  	}
    22  	return rw.ResponseWriter.Write(buf)
    23  }
    24  
    25  func (rw *bodyRecorder) WriteHeader(code int) {
    26  	rw.writeHeader(nil, code)
    27  }
    28  
    29  func (rw *bodyRecorder) Code() int {
    30  	return rw.code
    31  }
    32  
    33  func (rw *bodyRecorder) Body() *bytes.Buffer {
    34  	return rw.body
    35  }
    36  
    37  func (rw *bodyRecorder) writeHeader(buf []byte, code int) {
    38  	if rw.wroteHeader {
    39  		return
    40  	}
    41  
    42  	if rw.doRecord != nil {
    43  		header := rw.Header()
    44  		if len(buf) != 0 {
    45  			if _, hasType := header[HeaderContentType]; !hasType {
    46  				header.Set(HeaderContentType, http.DetectContentType(buf))
    47  			}
    48  		}
    49  
    50  		if rw.doRecord(code, header) {
    51  			rw.body = bytes.NewBuffer(nil)
    52  		}
    53  	}
    54  
    55  	if rw.body == nil {
    56  		rw.ResponseWriter.WriteHeader(code)
    57  	}
    58  
    59  	rw.code = code
    60  	rw.wroteHeader = true
    61  }