github.com/go-kivik/kivik/v4@v4.3.2/couchdb/chttp/trace.go (about)

     1  // Licensed under the Apache License, Version 2.0 (the "License"); you may not
     2  // use this file except in compliance with the License. You may obtain a copy of
     3  // the License at
     4  //
     5  //  http://www.apache.org/licenses/LICENSE-2.0
     6  //
     7  // Unless required by applicable law or agreed to in writing, software
     8  // distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
     9  // WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
    10  // License for the specific language governing permissions and limitations under
    11  // the License.
    12  
    13  package chttp
    14  
    15  import (
    16  	"bytes"
    17  	"context"
    18  	"io"
    19  	"net/http"
    20  )
    21  
    22  var clientTraceContextKey = &struct{ name string }{"client trace"}
    23  
    24  // ContextClientTrace returns the ClientTrace associated with the
    25  // provided context. If none, it returns nil.
    26  func ContextClientTrace(ctx context.Context) *ClientTrace {
    27  	trace, _ := ctx.Value(clientTraceContextKey).(*ClientTrace)
    28  	return trace
    29  }
    30  
    31  // ClientTrace is a set of hooks to run at various stages of an outgoing
    32  // HTTP request. Any particular hook may be nil. Functions may be
    33  // called concurrently from different goroutines and some may be called
    34  // after the request has completed or failed.
    35  type ClientTrace struct {
    36  	// HTTPResponse returns a cloe of the *http.Response received from the
    37  	// server, with the body set to nil. If you need the body, use the more
    38  	// expensive HTTPResponseBody.
    39  	HTTPResponse func(*http.Response)
    40  
    41  	// HTTPResponseBody returns a clone of the *http.Response received from the
    42  	// server, with the body cloned. This can be expensive for responses
    43  	// with large bodies.
    44  	HTTPResponseBody func(*http.Response)
    45  
    46  	// HTTPRequest returns a clone of the *http.Request sent to the server, with
    47  	// the body set to nil. If you need the body, use the more expensive
    48  	// HTTPRequestBody.
    49  	HTTPRequest func(*http.Request)
    50  
    51  	// HTTPRequestBody returns a clone of the *http.Request sent to the server,
    52  	// with the body cloned, if it is set. This can be expensive for requests
    53  	// with large bodies.
    54  	HTTPRequestBody func(*http.Request)
    55  }
    56  
    57  // WithClientTrace returns a new context based on the provided parent
    58  // ctx. HTTP client requests made with the returned context will use
    59  // the provided trace hooks, in addition to any previous hooks
    60  // registered with ctx. Any hooks defined in the provided trace will
    61  // be called first.
    62  func WithClientTrace(ctx context.Context, trace *ClientTrace) context.Context {
    63  	if trace == nil {
    64  		panic("nil trace")
    65  	}
    66  	return context.WithValue(ctx, clientTraceContextKey, trace)
    67  }
    68  
    69  func (t *ClientTrace) httpResponse(r *http.Response) {
    70  	if t.HTTPResponse == nil || r == nil {
    71  		return
    72  	}
    73  	clone := new(http.Response)
    74  	*clone = *r
    75  	clone.Body = nil
    76  	t.HTTPResponse(clone)
    77  }
    78  
    79  func (t *ClientTrace) httpResponseBody(r *http.Response) {
    80  	if t.HTTPResponseBody == nil || r == nil {
    81  		return
    82  	}
    83  	clone := new(http.Response)
    84  	*clone = *r
    85  	rBody := r.Body
    86  	body, readErr := io.ReadAll(rBody)
    87  	closeErr := rBody.Close()
    88  	r.Body = newReplay(body, readErr, closeErr)
    89  	clone.Body = newReplay(body, readErr, closeErr)
    90  	t.HTTPResponseBody(clone)
    91  }
    92  
    93  func (t *ClientTrace) httpRequest(r *http.Request) {
    94  	if t.HTTPRequest == nil {
    95  		return
    96  	}
    97  	clone := new(http.Request)
    98  	*clone = *r
    99  	clone.Body = nil
   100  	t.HTTPRequest(clone)
   101  }
   102  
   103  func (t *ClientTrace) httpRequestBody(r *http.Request) {
   104  	if t.HTTPRequestBody == nil {
   105  		return
   106  	}
   107  	clone := new(http.Request)
   108  	*clone = *r
   109  	if r.Body != nil {
   110  		rBody := r.Body
   111  		body, readErr := io.ReadAll(rBody)
   112  		closeErr := rBody.Close()
   113  		r.Body = newReplay(body, readErr, closeErr)
   114  		clone.Body = newReplay(body, readErr, closeErr)
   115  	}
   116  	t.HTTPRequestBody(clone)
   117  }
   118  
   119  func newReplay(body []byte, readErr, closeErr error) io.ReadCloser {
   120  	if readErr == nil && closeErr == nil {
   121  		return io.NopCloser(bytes.NewReader(body))
   122  	}
   123  	return &replayReadCloser{
   124  		Reader:   io.NopCloser(bytes.NewReader(body)),
   125  		readErr:  readErr,
   126  		closeErr: closeErr,
   127  	}
   128  }
   129  
   130  // replayReadCloser replays read and close errors
   131  type replayReadCloser struct {
   132  	io.Reader
   133  	readErr  error
   134  	closeErr error
   135  }
   136  
   137  func (r *replayReadCloser) Read(p []byte) (int, error) {
   138  	c, err := r.Reader.Read(p)
   139  	if err == io.EOF && r.readErr != nil {
   140  		err = r.readErr
   141  	}
   142  	return c, err
   143  }
   144  
   145  func (r *replayReadCloser) Close() error {
   146  	return r.closeErr
   147  }