go.undefinedlabs.com/scopeagent@v0.4.2/instrumentation/nethttp/responseTracker.go (about)

     1  package nethttp
     2  
     3  import (
     4  	"io"
     5  	"net/http"
     6  )
     7  
     8  type responseTracker struct {
     9  	http.ResponseWriter
    10  	status                 int
    11  	wroteheader            bool
    12  	payloadInstrumentation bool
    13  	payloadBuffer          []byte
    14  }
    15  
    16  var payloadBufferSize = 512
    17  
    18  func (w *responseTracker) WriteHeader(status int) {
    19  	w.status = status
    20  	w.wroteheader = true
    21  	w.ResponseWriter.WriteHeader(status)
    22  }
    23  
    24  func (w *responseTracker) Write(b []byte) (int, error) {
    25  	if !w.wroteheader {
    26  		w.wroteheader = true
    27  		w.status = 200
    28  	}
    29  	if w.payloadInstrumentation {
    30  		payloadSize := len(w.payloadBuffer)
    31  		if payloadSize < payloadBufferSize {
    32  			bufferMissing := payloadBufferSize - payloadSize
    33  			if bufferMissing < len(b) {
    34  				w.payloadBuffer = append(w.payloadBuffer, b[:bufferMissing]...)
    35  			} else {
    36  				w.payloadBuffer = append(w.payloadBuffer, b...)
    37  			}
    38  		}
    39  	}
    40  	return w.ResponseWriter.Write(b)
    41  }
    42  
    43  // wrappedResponseWriter returns a wrapped version of the original
    44  // ResponseWriter and only implements the same combination of additional
    45  // interfaces as the original.  This implementation is based on
    46  // https://github.com/felixge/httpsnoop.
    47  func (w *responseTracker) wrappedResponseWriter() http.ResponseWriter {
    48  	var (
    49  		hj, i0 = w.ResponseWriter.(http.Hijacker)
    50  		cn, i1 = w.ResponseWriter.(http.CloseNotifier)
    51  		pu, i2 = w.ResponseWriter.(http.Pusher)
    52  		fl, i3 = w.ResponseWriter.(http.Flusher)
    53  		rf, i4 = w.ResponseWriter.(io.ReaderFrom)
    54  	)
    55  
    56  	switch {
    57  	case !i0 && !i1 && !i2 && !i3 && !i4:
    58  		return struct {
    59  			http.ResponseWriter
    60  		}{w}
    61  	case !i0 && !i1 && !i2 && !i3 && i4:
    62  		return struct {
    63  			http.ResponseWriter
    64  			io.ReaderFrom
    65  		}{w, rf}
    66  	case !i0 && !i1 && !i2 && i3 && !i4:
    67  		return struct {
    68  			http.ResponseWriter
    69  			http.Flusher
    70  		}{w, fl}
    71  	case !i0 && !i1 && !i2 && i3 && i4:
    72  		return struct {
    73  			http.ResponseWriter
    74  			http.Flusher
    75  			io.ReaderFrom
    76  		}{w, fl, rf}
    77  	case !i0 && !i1 && i2 && !i3 && !i4:
    78  		return struct {
    79  			http.ResponseWriter
    80  			http.Pusher
    81  		}{w, pu}
    82  	case !i0 && !i1 && i2 && !i3 && i4:
    83  		return struct {
    84  			http.ResponseWriter
    85  			http.Pusher
    86  			io.ReaderFrom
    87  		}{w, pu, rf}
    88  	case !i0 && !i1 && i2 && i3 && !i4:
    89  		return struct {
    90  			http.ResponseWriter
    91  			http.Pusher
    92  			http.Flusher
    93  		}{w, pu, fl}
    94  	case !i0 && !i1 && i2 && i3 && i4:
    95  		return struct {
    96  			http.ResponseWriter
    97  			http.Pusher
    98  			http.Flusher
    99  			io.ReaderFrom
   100  		}{w, pu, fl, rf}
   101  	case !i0 && i1 && !i2 && !i3 && !i4:
   102  		return struct {
   103  			http.ResponseWriter
   104  			http.CloseNotifier
   105  		}{w, cn}
   106  	case !i0 && i1 && !i2 && !i3 && i4:
   107  		return struct {
   108  			http.ResponseWriter
   109  			http.CloseNotifier
   110  			io.ReaderFrom
   111  		}{w, cn, rf}
   112  	case !i0 && i1 && !i2 && i3 && !i4:
   113  		return struct {
   114  			http.ResponseWriter
   115  			http.CloseNotifier
   116  			http.Flusher
   117  		}{w, cn, fl}
   118  	case !i0 && i1 && !i2 && i3 && i4:
   119  		return struct {
   120  			http.ResponseWriter
   121  			http.CloseNotifier
   122  			http.Flusher
   123  			io.ReaderFrom
   124  		}{w, cn, fl, rf}
   125  	case !i0 && i1 && i2 && !i3 && !i4:
   126  		return struct {
   127  			http.ResponseWriter
   128  			http.CloseNotifier
   129  			http.Pusher
   130  		}{w, cn, pu}
   131  	case !i0 && i1 && i2 && !i3 && i4:
   132  		return struct {
   133  			http.ResponseWriter
   134  			http.CloseNotifier
   135  			http.Pusher
   136  			io.ReaderFrom
   137  		}{w, cn, pu, rf}
   138  	case !i0 && i1 && i2 && i3 && !i4:
   139  		return struct {
   140  			http.ResponseWriter
   141  			http.CloseNotifier
   142  			http.Pusher
   143  			http.Flusher
   144  		}{w, cn, pu, fl}
   145  	case !i0 && i1 && i2 && i3 && i4:
   146  		return struct {
   147  			http.ResponseWriter
   148  			http.CloseNotifier
   149  			http.Pusher
   150  			http.Flusher
   151  			io.ReaderFrom
   152  		}{w, cn, pu, fl, rf}
   153  	case i0 && !i1 && !i2 && !i3 && !i4:
   154  		return struct {
   155  			http.ResponseWriter
   156  			http.Hijacker
   157  		}{w, hj}
   158  	case i0 && !i1 && !i2 && !i3 && i4:
   159  		return struct {
   160  			http.ResponseWriter
   161  			http.Hijacker
   162  			io.ReaderFrom
   163  		}{w, hj, rf}
   164  	case i0 && !i1 && !i2 && i3 && !i4:
   165  		return struct {
   166  			http.ResponseWriter
   167  			http.Hijacker
   168  			http.Flusher
   169  		}{w, hj, fl}
   170  	case i0 && !i1 && !i2 && i3 && i4:
   171  		return struct {
   172  			http.ResponseWriter
   173  			http.Hijacker
   174  			http.Flusher
   175  			io.ReaderFrom
   176  		}{w, hj, fl, rf}
   177  	case i0 && !i1 && i2 && !i3 && !i4:
   178  		return struct {
   179  			http.ResponseWriter
   180  			http.Hijacker
   181  			http.Pusher
   182  		}{w, hj, pu}
   183  	case i0 && !i1 && i2 && !i3 && i4:
   184  		return struct {
   185  			http.ResponseWriter
   186  			http.Hijacker
   187  			http.Pusher
   188  			io.ReaderFrom
   189  		}{w, hj, pu, rf}
   190  	case i0 && !i1 && i2 && i3 && !i4:
   191  		return struct {
   192  			http.ResponseWriter
   193  			http.Hijacker
   194  			http.Pusher
   195  			http.Flusher
   196  		}{w, hj, pu, fl}
   197  	case i0 && !i1 && i2 && i3 && i4:
   198  		return struct {
   199  			http.ResponseWriter
   200  			http.Hijacker
   201  			http.Pusher
   202  			http.Flusher
   203  			io.ReaderFrom
   204  		}{w, hj, pu, fl, rf}
   205  	case i0 && i1 && !i2 && !i3 && !i4:
   206  		return struct {
   207  			http.ResponseWriter
   208  			http.Hijacker
   209  			http.CloseNotifier
   210  		}{w, hj, cn}
   211  	case i0 && i1 && !i2 && !i3 && i4:
   212  		return struct {
   213  			http.ResponseWriter
   214  			http.Hijacker
   215  			http.CloseNotifier
   216  			io.ReaderFrom
   217  		}{w, hj, cn, rf}
   218  	case i0 && i1 && !i2 && i3 && !i4:
   219  		return struct {
   220  			http.ResponseWriter
   221  			http.Hijacker
   222  			http.CloseNotifier
   223  			http.Flusher
   224  		}{w, hj, cn, fl}
   225  	case i0 && i1 && !i2 && i3 && i4:
   226  		return struct {
   227  			http.ResponseWriter
   228  			http.Hijacker
   229  			http.CloseNotifier
   230  			http.Flusher
   231  			io.ReaderFrom
   232  		}{w, hj, cn, fl, rf}
   233  	case i0 && i1 && i2 && !i3 && !i4:
   234  		return struct {
   235  			http.ResponseWriter
   236  			http.Hijacker
   237  			http.CloseNotifier
   238  			http.Pusher
   239  		}{w, hj, cn, pu}
   240  	case i0 && i1 && i2 && !i3 && i4:
   241  		return struct {
   242  			http.ResponseWriter
   243  			http.Hijacker
   244  			http.CloseNotifier
   245  			http.Pusher
   246  			io.ReaderFrom
   247  		}{w, hj, cn, pu, rf}
   248  	case i0 && i1 && i2 && i3 && !i4:
   249  		return struct {
   250  			http.ResponseWriter
   251  			http.Hijacker
   252  			http.CloseNotifier
   253  			http.Pusher
   254  			http.Flusher
   255  		}{w, hj, cn, pu, fl}
   256  	case i0 && i1 && i2 && i3 && i4:
   257  		return struct {
   258  			http.ResponseWriter
   259  			http.Hijacker
   260  			http.CloseNotifier
   261  			http.Pusher
   262  			http.Flusher
   263  			io.ReaderFrom
   264  		}{w, hj, cn, pu, fl, rf}
   265  	default:
   266  		return struct {
   267  			http.ResponseWriter
   268  		}{w}
   269  	}
   270  }