diff options
author | GitLab Bot <gitlab-bot@gitlab.com> | 2020-12-17 14:59:07 +0300 |
---|---|---|
committer | GitLab Bot <gitlab-bot@gitlab.com> | 2020-12-17 14:59:07 +0300 |
commit | 8b573c94895dc0ac0e1d9d59cf3e8745e8b539ca (patch) | |
tree | 544930fb309b30317ae9797a9683768705d664c4 /workhorse/internal/helper/countingresponsewriter.go | |
parent | 4b1de649d0168371549608993deac953eb692019 (diff) |
Add latest changes from gitlab-org/gitlab@13-7-stable-eev13.7.0-rc42
Diffstat (limited to 'workhorse/internal/helper/countingresponsewriter.go')
-rw-r--r-- | workhorse/internal/helper/countingresponsewriter.go | 56 |
1 files changed, 56 insertions, 0 deletions
diff --git a/workhorse/internal/helper/countingresponsewriter.go b/workhorse/internal/helper/countingresponsewriter.go new file mode 100644 index 00000000000..a79d51d4c6a --- /dev/null +++ b/workhorse/internal/helper/countingresponsewriter.go @@ -0,0 +1,56 @@ +package helper + +import ( + "net/http" +) + +type CountingResponseWriter interface { + http.ResponseWriter + Count() int64 + Status() int +} + +type countingResponseWriter struct { + rw http.ResponseWriter + status int + count int64 +} + +func NewCountingResponseWriter(rw http.ResponseWriter) CountingResponseWriter { + return &countingResponseWriter{rw: rw} +} + +func (c *countingResponseWriter) Header() http.Header { + return c.rw.Header() +} + +func (c *countingResponseWriter) Write(data []byte) (int, error) { + if c.status == 0 { + c.WriteHeader(http.StatusOK) + } + + n, err := c.rw.Write(data) + c.count += int64(n) + return n, err +} + +func (c *countingResponseWriter) WriteHeader(status int) { + if c.status != 0 { + return + } + + c.status = status + c.rw.WriteHeader(status) +} + +// Count returns the number of bytes written to the ResponseWriter. This +// function is not thread-safe. +func (c *countingResponseWriter) Count() int64 { + return c.count +} + +// Status returns the first HTTP status value that was written to the +// ResponseWriter. This function is not thread-safe. +func (c *countingResponseWriter) Status() int { + return c.status +} |