github.com/gocaveman/caveman@v0.0.0-20191211162744-0ddf99dbdf6e/webutil/response-context.go (about)

     1  package webutil
     2  
     3  import (
     4  	"bufio"
     5  	"context"
     6  	"net"
     7  	"net/http"
     8  )
     9  
    10  // TODO: a struct (WrapResponseWriter) that cleanly wraps a ResponseWriter
    11  // (can be embedded easily by another struct to add functionality)
    12  
    13  // TODO: a struct that embeds WrapResponseWriter to implement context cancellation
    14  // when WriteHeader is called.  Name: CancelingResponseWriter
    15  
    16  // TODO: we could also add one here that dumps everything to the log, lower priority
    17  // but probably useful (it should be smart enough to ungzip what GzipResponseWriter
    18  // has done in order to make it human readable).
    19  
    20  // DummyResponseWriterCloser implements only the context cancellation and otherwise does nothing.
    21  
    22  type ContextCancelHandler struct{}
    23  
    24  func NewContextCancelHandler() ContextCancelHandler {
    25  	return ContextCancelHandler{}
    26  }
    27  
    28  func (h ContextCancelHandler) ServeHTTPChain(w http.ResponseWriter, r *http.Request) (http.ResponseWriter, *http.Request) {
    29  	ctx, cancelFunc := context.WithCancel(r.Context())
    30  	retw := &ContextCancelResponseWriter{
    31  		ResponseWriter: w,
    32  		cancelFunc:     cancelFunc,
    33  	}
    34  	return retw, r.WithContext(ctx)
    35  }
    36  
    37  type ContextCancelResponseWriter struct {
    38  	http.ResponseWriter
    39  	cancelFunc context.CancelFunc
    40  }
    41  
    42  func (w *ContextCancelResponseWriter) Write(p []byte) (int, error) {
    43  	// TODO: Either as an option here or in a separate handler,
    44  	// we should have something that can dump out the stack trace; as
    45  	// a debug tool so we can see what handler ended up writing.
    46  	w.cancelFunc()
    47  	return w.ResponseWriter.Write(p)
    48  }
    49  
    50  func (w *ContextCancelResponseWriter) WriteHeader(c int) {
    51  	w.cancelFunc()
    52  	w.ResponseWriter.WriteHeader(c)
    53  }
    54  
    55  // func (w *ContextCancelResponseWriter) Close() (err error) {
    56  // 	log.Printf("FIXME: Close() is probably not a good idea - using Flush() instead of Close keeps to the existing API and can serve the same purpose")
    57  // 	w.cancelFunc()
    58  // 	if c, ok := w.ResponseWriter.(io.Closer); ok {
    59  // 		err = c.Close()
    60  // 	}
    61  // 	return err
    62  // }
    63  
    64  // Flush cancels the context and calls Flush on the underlying ResponseWriter.
    65  func (w *ContextCancelResponseWriter) Flush() {
    66  	w.cancelFunc()
    67  	w.ResponseWriter.(http.Flusher).Flush()
    68  }
    69  
    70  // Hijack cancels the context and calls Hijack on the underlying ResponseWriter.
    71  func (w *ContextCancelResponseWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
    72  	w.cancelFunc()
    73  	return w.ResponseWriter.(http.Hijacker).Hijack()
    74  }
    75  
    76  func (w *ContextCancelResponseWriter) CloseNotify() <-chan bool {
    77  	return w.ResponseWriter.(http.CloseNotifier).CloseNotify()
    78  }
    79  
    80  func (w *ContextCancelResponseWriter) Push(target string, opts *http.PushOptions) error {
    81  	return w.ResponseWriter.(http.Pusher).Push(target, opts)
    82  }