sdk/go/httpserver/responsewriter.go

   1 // Copyright (C) The Arvados Authors. All rights reserved.
   2 //
   3 // SPDX-License-Identifier: Apache-2.0
   4
   5 package httpserver
   6
   7 import (
   8         "net/http"
   9 )
  10
  11 type ResponseWriter interface {
  12         http.ResponseWriter
  13         WroteStatus() int
  14         WroteBodyBytes() int
  15 }
  16
  17 // responseWriter wraps http.ResponseWriter and exposes the status
  18 // sent, the number of bytes sent to the client, and the last write
  19 // error.
  20 type responseWriter struct {
  21         http.ResponseWriter
  22         wroteStatus    int   // Last status given to WriteHeader()
  23         wroteBodyBytes int   // Bytes successfully written
  24         err            error // Last error returned from Write()
  25 }
  26
  27 func WrapResponseWriter(orig http.ResponseWriter) ResponseWriter {
  28         return &responseWriter{ResponseWriter: orig}
  29 }
  30
  31 func (w *responseWriter) CloseNotify() <-chan bool {
  32         if cn, ok := w.ResponseWriter.(http.CloseNotifier); ok {
  33                 return cn.CloseNotify()
  34         }
  35         return nil
  36 }
  37
  38 func (w *responseWriter) WriteHeader(s int) {
  39         w.wroteStatus = s
  40         w.ResponseWriter.WriteHeader(s)
  41 }
  42
  43 func (w *responseWriter) Write(data []byte) (n int, err error) {
  44         if w.wroteStatus == 0 {
  45                 w.WriteHeader(http.StatusOK)
  46         }
  47         n, err = w.ResponseWriter.Write(data)
  48         w.wroteBodyBytes += n
  49         w.err = err
  50         return
  51 }
  52
  53 func (w *responseWriter) WroteStatus() int {
  54         return w.wroteStatus
  55 }
  56
  57 func (w *responseWriter) WroteBodyBytes() int {
  58         return w.wroteBodyBytes
  59 }
  60
  61 func (w *responseWriter) Err() error {
  62         return w.err
  63 }