github.com/richardwilkes/toolbox@v1.121.0/xio/network/xhttp/status_response_writer.go (about)

     1  // Copyright (c) 2016-2024 by Richard A. Wilkes. All rights reserved.
     2  //
     3  // This Source Code Form is subject to the terms of the Mozilla Public
     4  // License, version 2.0. If a copy of the MPL was not distributed with
     5  // this file, You can obtain one at http://mozilla.org/MPL/2.0/.
     6  //
     7  // This Source Code Form is "Incompatible With Secondary Licenses", as
     8  // defined by the Mozilla Public License, version 2.0.
     9  
    10  package xhttp
    11  
    12  import "net/http"
    13  
    14  // StatusResponseWriter wraps an http.ResponseWriter and provides methods to retrieve the status code and number of
    15  // bytes written.
    16  type StatusResponseWriter struct {
    17  	Original http.ResponseWriter
    18  	Head     bool
    19  	status   int
    20  	written  int
    21  }
    22  
    23  // Status returns the status that was set, or http.StatusOK if no call to WriteHeader() was made.
    24  func (w *StatusResponseWriter) Status() int {
    25  	if w.status != 0 {
    26  		return w.status
    27  	}
    28  	return http.StatusOK
    29  }
    30  
    31  // BytesWritten returns the number of bytes written.
    32  func (w *StatusResponseWriter) BytesWritten() int {
    33  	return w.written
    34  }
    35  
    36  // Header implements http.ResponseWriter.
    37  func (w *StatusResponseWriter) Header() http.Header {
    38  	return w.Original.Header()
    39  }
    40  
    41  // Write implements http.ResponseWriter.
    42  func (w *StatusResponseWriter) Write(data []byte) (int, error) {
    43  	if w.Head {
    44  		return len(data), nil
    45  	}
    46  	n, err := w.Original.Write(data)
    47  	w.written += n
    48  	return n, err
    49  }
    50  
    51  // WriteHeader implements http.ResponseWriter.
    52  func (w *StatusResponseWriter) WriteHeader(status int) {
    53  	w.status = status
    54  	w.Original.WriteHeader(status)
    55  }
    56  
    57  // Flush implements http.Flusher.
    58  func (w *StatusResponseWriter) Flush() {
    59  	f, ok := w.Original.(http.Flusher)
    60  	if ok {
    61  		f.Flush()
    62  	}
    63  }