github.com/hellobchain/third_party@v0.0.0-20230331131523-deb0478a2e52/go-chi/chi/middleware/closenotify17.go (about)

     1  //go:build go1.7 && !go1.8
     2  // +build go1.7,!go1.8
     3  
     4  package middleware
     5  
     6  import (
     7  	"context"
     8  	"github.com/hellobchain/newcryptosm/http"
     9  )
    10  
    11  // CloseNotify is a middleware that cancels ctx when the underlying
    12  // connection has gone away. It can be used to cancel long operations
    13  // on the server when the client disconnects before the response is ready.
    14  //
    15  // Note: this behaviour is standard in Go 1.8+, so the middleware does nothing
    16  // on 1.8+ and exists just for backwards compatibility.
    17  func CloseNotify(next http.Handler) http.Handler {
    18  	fn := func(w http.ResponseWriter, r *http.Request) {
    19  		cn, ok := w.(http.CloseNotifier)
    20  		if !ok {
    21  			panic("chi/middleware: CloseNotify expects http.ResponseWriter to implement http.CloseNotifier interface")
    22  		}
    23  		closeNotifyCh := cn.CloseNotify()
    24  
    25  		ctx, cancel := context.WithCancel(r.Context())
    26  		defer cancel()
    27  
    28  		go func() {
    29  			select {
    30  			case <-ctx.Done():
    31  				return
    32  			case <-closeNotifyCh:
    33  				cancel()
    34  				return
    35  			}
    36  		}()
    37  
    38  		r = r.WithContext(ctx)
    39  		next.ServeHTTP(w, r)
    40  	}
    41  
    42  	return http.HandlerFunc(fn)
    43  }