github.com/gofunct/common@v0.0.0-20190131174352-fd058c7fbf22/pkg/transport/middleware/passing_header_middleware.go (about)

     1  package middleware
     2  
     3  import (
     4  	"net/http"
     5  	"sync"
     6  
     7  	"github.com/grpc-ecosystem/grpc-gateway/runtime"
     8  )
     9  
    10  // PassedHeaderDeciderFunc returns true if given header should be passed to gRPC server metadata.
    11  type PassedHeaderDeciderFunc func(string) bool
    12  
    13  func CreatePassingHeaderMiddleware(decide PassedHeaderDeciderFunc) HTTPServerMiddleware {
    14  	return func(next http.Handler) http.Handler {
    15  		cache := new(sync.Map)
    16  		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
    17  			newHeader := make(http.Header, 2*len(r.Header))
    18  
    19  			for k := range r.Header {
    20  				v := r.Header.Get(k)
    21  				if newKey, ok := cache.Load(k); ok {
    22  					newHeader.Set(newKey.(string), v)
    23  				} else if decide(k) {
    24  					newKey := runtime.MetadataHeaderPrefix + k
    25  					cache.Store(k, newKey)
    26  					newHeader.Set(newKey, v)
    27  				}
    28  				newHeader.Set(k, v)
    29  			}
    30  
    31  			r.Header = newHeader
    32  
    33  			next.ServeHTTP(w, r)
    34  		})
    35  	}
    36  }