code.gitea.io/gitea@v1.19.3/modules/context/response.go (about)

     1  // Copyright 2021 The Gitea Authors. All rights reserved.
     2  // SPDX-License-Identifier: MIT
     3  
     4  package context
     5  
     6  import (
     7  	"net/http"
     8  )
     9  
    10  // ResponseWriter represents a response writer for HTTP
    11  type ResponseWriter interface {
    12  	http.ResponseWriter
    13  	Flush()
    14  	Status() int
    15  	Before(func(ResponseWriter))
    16  	Size() int
    17  }
    18  
    19  var _ ResponseWriter = &Response{}
    20  
    21  // Response represents a response
    22  type Response struct {
    23  	http.ResponseWriter
    24  	written        int
    25  	status         int
    26  	befores        []func(ResponseWriter)
    27  	beforeExecuted bool
    28  }
    29  
    30  // Size return written size
    31  func (r *Response) Size() int {
    32  	return r.written
    33  }
    34  
    35  // Write writes bytes to HTTP endpoint
    36  func (r *Response) Write(bs []byte) (int, error) {
    37  	if !r.beforeExecuted {
    38  		for _, before := range r.befores {
    39  			before(r)
    40  		}
    41  		r.beforeExecuted = true
    42  	}
    43  	size, err := r.ResponseWriter.Write(bs)
    44  	r.written += size
    45  	if err != nil {
    46  		return size, err
    47  	}
    48  	if r.status == 0 {
    49  		r.status = http.StatusOK
    50  	}
    51  	return size, nil
    52  }
    53  
    54  // WriteHeader write status code
    55  func (r *Response) WriteHeader(statusCode int) {
    56  	if !r.beforeExecuted {
    57  		for _, before := range r.befores {
    58  			before(r)
    59  		}
    60  		r.beforeExecuted = true
    61  	}
    62  	if r.status == 0 {
    63  		r.status = statusCode
    64  		r.ResponseWriter.WriteHeader(statusCode)
    65  	}
    66  }
    67  
    68  // Flush flush cached data
    69  func (r *Response) Flush() {
    70  	if f, ok := r.ResponseWriter.(http.Flusher); ok {
    71  		f.Flush()
    72  	}
    73  }
    74  
    75  // Status returned status code written
    76  func (r *Response) Status() int {
    77  	return r.status
    78  }
    79  
    80  // Before allows for a function to be called before the ResponseWriter has been written to. This is
    81  // useful for setting headers or any other operations that must happen before a response has been written.
    82  func (r *Response) Before(f func(ResponseWriter)) {
    83  	r.befores = append(r.befores, f)
    84  }
    85  
    86  // NewResponse creates a response
    87  func NewResponse(resp http.ResponseWriter) *Response {
    88  	if v, ok := resp.(*Response); ok {
    89  		return v
    90  	}
    91  	return &Response{
    92  		ResponseWriter: resp,
    93  		status:         0,
    94  		befores:        make([]func(ResponseWriter), 0),
    95  	}
    96  }