github.com/xmidt-org/webpa-common@v1.11.9/xhttp/xcontext/contextaware.go (about)

     1  package xcontext
     2  
     3  import (
     4  	"bufio"
     5  	"context"
     6  	"errors"
     7  	"net"
     8  	"net/http"
     9  )
    10  
    11  // ContextAware is an optional mixin implemented by anything with can hold a context
    12  type ContextAware interface {
    13  	// Context *never* returns a nil context
    14  	Context() context.Context
    15  	SetContext(context.Context)
    16  }
    17  
    18  type contextAwareResponseWriter struct {
    19  	http.ResponseWriter
    20  	ctx context.Context
    21  }
    22  
    23  var _ ContextAware = &contextAwareResponseWriter{}
    24  var _ http.Hijacker = &contextAwareResponseWriter{}
    25  var _ http.Flusher = &contextAwareResponseWriter{}
    26  var _ http.Pusher = &contextAwareResponseWriter{}
    27  
    28  func (carw contextAwareResponseWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
    29  	if h, ok := carw.ResponseWriter.(http.Hijacker); ok {
    30  		return h.Hijack()
    31  	}
    32  
    33  	return nil, nil, errors.New("Hijacker not supported")
    34  }
    35  
    36  func (carw contextAwareResponseWriter) Flush() {
    37  	if f, ok := carw.ResponseWriter.(http.Flusher); ok {
    38  		f.Flush()
    39  	}
    40  }
    41  
    42  func (carw contextAwareResponseWriter) Push(target string, opts *http.PushOptions) error {
    43  	if p, ok := carw.ResponseWriter.(http.Pusher); ok {
    44  		return p.Push(target, opts)
    45  	}
    46  
    47  	return errors.New("Pusher not supported")
    48  }
    49  
    50  func (carw *contextAwareResponseWriter) Context() context.Context {
    51  	if carw.ctx != nil {
    52  		return carw.ctx
    53  	}
    54  
    55  	return context.Background()
    56  }
    57  
    58  func (carw *contextAwareResponseWriter) SetContext(ctx context.Context) {
    59  	if ctx == nil {
    60  		// mimic the behavior of the net/http package
    61  		panic("nil context")
    62  	}
    63  
    64  	carw.ctx = ctx
    65  }
    66  
    67  func Context(response http.ResponseWriter, request *http.Request) context.Context {
    68  	if ca, ok := response.(ContextAware); ok {
    69  		return ca.Context()
    70  	}
    71  
    72  	// fallback to the request's context
    73  	return request.Context()
    74  }
    75  
    76  // SetContext associates a context with a response.  Useful for decorated code that needs to communicate
    77  // a context back up the call stack.
    78  //
    79  // Note that since ContextAware is an optional interface, it's possible that the supplied ResponseWriter does
    80  // not implement ContextAware.  This is tolerated, so as to be backward compatible.
    81  //
    82  // The returned ResponseWriter will always be ContextAware.  This writer can be used for subsequent handling code.
    83  func SetContext(response http.ResponseWriter, ctx context.Context) http.ResponseWriter {
    84  	if ca, ok := response.(ContextAware); ok {
    85  		ca.SetContext(ctx)
    86  		return response
    87  	}
    88  
    89  	if ctx == nil {
    90  		panic("nil context")
    91  	}
    92  
    93  	return &contextAwareResponseWriter{response, ctx}
    94  }
    95  
    96  // WithContext associates a context with the response/request pair that can later be accessed via the Context function.
    97  // If response is already ContextAware, it is used and returned as is.
    98  //
    99  // Useful for code that is decorating http handling code in order to establish a context.
   100  func WithContext(response http.ResponseWriter, request *http.Request, ctx context.Context) (http.ResponseWriter, *http.Request) {
   101  	if ca, ok := response.(ContextAware); ok {
   102  		ca.SetContext(ctx)
   103  		return response, request.WithContext(ctx)
   104  	}
   105  
   106  	if ctx == nil {
   107  		ctx = context.Background()
   108  	}
   109  
   110  	return &contextAwareResponseWriter{response, ctx}, request.WithContext(ctx)
   111  }