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         "bufio"
   9         "net"
  10         "net/http"
  11         "time"
  12 )
  13
  14 const sniffBytes = 1024
  15
  16 type ResponseWriter interface {
  17         http.ResponseWriter
  18         WroteStatus() int
  19         WroteBodyBytes() int
  20         Sniffed() []byte
  21 }
  22
  23 // responseWriter wraps http.ResponseWriter and exposes the status
  24 // sent, the number of bytes sent to the client, and the last write
  25 // error.
  26 type responseWriter struct {
  27         http.ResponseWriter
  28         wroteStatus    int   // First status given to WriteHeader()
  29         wroteBodyBytes int   // Bytes successfully written
  30         err            error // Last error returned from Write()
  31         sniffed        []byte
  32 }
  33
  34 func WrapResponseWriter(orig http.ResponseWriter) ResponseWriter {
  35         return &responseWriter{ResponseWriter: orig}
  36 }
  37
  38 func (w *responseWriter) WriteHeader(s int) {
  39         if w.wroteStatus == 0 {
  40                 w.wroteStatus = s
  41         }
  42         // ...else it's too late to change the status seen by the
  43         // client -- but we call the wrapped WriteHeader() anyway so
  44         // it can log a warning.
  45         w.ResponseWriter.WriteHeader(s)
  46 }
  47
  48 func (w *responseWriter) Write(data []byte) (n int, err error) {
  49         if w.wroteStatus == 0 {
  50                 w.WriteHeader(http.StatusOK)
  51         } else if w.wroteStatus >= 400 {
  52                 w.sniff(data)
  53         }
  54         n, err = w.ResponseWriter.Write(data)
  55         w.wroteBodyBytes += n
  56         w.err = err
  57         return
  58 }
  59
  60 func (w *responseWriter) WroteStatus() int {
  61         return w.wroteStatus
  62 }
  63
  64 func (w *responseWriter) WroteBodyBytes() int {
  65         return w.wroteBodyBytes
  66 }
  67
  68 func (w *responseWriter) Err() error {
  69         return w.err
  70 }
  71
  72 func (w *responseWriter) sniff(data []byte) {
  73         max := sniffBytes - len(w.sniffed)
  74         if max <= 0 {
  75                 return
  76         } else if max < len(data) {
  77                 data = data[:max]
  78         }
  79         w.sniffed = append(w.sniffed, data...)
  80 }
  81
  82 func (w *responseWriter) Sniffed() []byte {
  83         return w.sniffed
  84 }
  85
  86 func (w *responseWriter) Unwrap() http.ResponseWriter {
  87         return w.ResponseWriter
  88 }
  89
  90 // ResponseControllerShim uses a ResponseController to re-add the
  91 // optional interface methods to a ResponseWriter that has lost them
  92 // via wrapping by middleware.
  93 //
  94 // This allows us to combine old code (like x/net/websocket) with
  95 // middleware that doesn't explicitly support the optional interfaces
  96 // (like responseTimer and responseWriter here).
  97 type ResponseControllerShim struct{ http.ResponseWriter }
  98
  99 func (s ResponseControllerShim) EnableFullDuplex() error {
 100         return http.NewResponseController(s.ResponseWriter).EnableFullDuplex()
 101 }
 102
 103 func (s ResponseControllerShim) Hijack() (net.Conn, *bufio.ReadWriter, error) {
 104         return http.NewResponseController(s.ResponseWriter).Hijack()
 105 }
 106
 107 func (s ResponseControllerShim) SetReadDeadline(d time.Time) error {
 108         return http.NewResponseController(s.ResponseWriter).SetReadDeadline(d)
 109 }
 110
 111 func (s ResponseControllerShim) SetWriteDeadline(d time.Time) error {
 112         return http.NewResponseController(s.ResponseWriter).SetWriteDeadline(d)
 113 }
 114
 115 func (s ResponseControllerShim) Flush() error {
 116         return http.NewResponseController(s.ResponseWriter).Flush()
 117 }