code.gitea.io/gitea@v1.22.3/services/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  	web_types "code.gitea.io/gitea/modules/web/types"
    10  )
    11  
    12  // ResponseWriter represents a response writer for HTTP
    13  type ResponseWriter interface {
    14  	http.ResponseWriter
    15  	http.Flusher
    16  	web_types.ResponseStatusProvider
    17  
    18  	Before(func(ResponseWriter))
    19  
    20  	Status() int // used by access logger template
    21  	Size() int   // used by access logger template
    22  }
    23  
    24  var _ ResponseWriter = &Response{}
    25  
    26  // Response represents a response
    27  type Response struct {
    28  	http.ResponseWriter
    29  	written        int
    30  	status         int
    31  	befores        []func(ResponseWriter)
    32  	beforeExecuted bool
    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  func (r *Response) Status() int {
    55  	return r.status
    56  }
    57  
    58  func (r *Response) Size() int {
    59  	return r.written
    60  }
    61  
    62  // WriteHeader write status code
    63  func (r *Response) WriteHeader(statusCode int) {
    64  	if !r.beforeExecuted {
    65  		for _, before := range r.befores {
    66  			before(r)
    67  		}
    68  		r.beforeExecuted = true
    69  	}
    70  	if r.status == 0 {
    71  		r.status = statusCode
    72  		r.ResponseWriter.WriteHeader(statusCode)
    73  	}
    74  }
    75  
    76  // Flush flushes cached data
    77  func (r *Response) Flush() {
    78  	if f, ok := r.ResponseWriter.(http.Flusher); ok {
    79  		f.Flush()
    80  	}
    81  }
    82  
    83  // WrittenStatus returned status code written
    84  func (r *Response) WrittenStatus() int {
    85  	return r.status
    86  }
    87  
    88  // Before allows for a function to be called before the ResponseWriter has been written to. This is
    89  // useful for setting headers or any other operations that must happen before a response has been written.
    90  func (r *Response) Before(f func(ResponseWriter)) {
    91  	r.befores = append(r.befores, f)
    92  }
    93  
    94  func WrapResponseWriter(resp http.ResponseWriter) *Response {
    95  	if v, ok := resp.(*Response); ok {
    96  		return v
    97  	}
    98  	return &Response{
    99  		ResponseWriter: resp,
   100  		status:         0,
   101  		befores:        make([]func(ResponseWriter), 0),
   102  	}
   103  }