github.com/zak-blake/goa@v1.4.1/middleware/request_id.go (about)

     1  package middleware
     2  
     3  import (
     4  	"crypto/rand"
     5  	"encoding/base64"
     6  	"fmt"
     7  	"net/http"
     8  	"strings"
     9  	"sync/atomic"
    10  
    11  	"github.com/goadesign/goa"
    12  
    13  	"context"
    14  )
    15  
    16  const (
    17  	// RequestIDHeader is the name of the header used to transmit the request ID.
    18  	RequestIDHeader = "X-Request-Id"
    19  
    20  	// DefaultRequestIDLengthLimit is the default maximum length for the request ID header value.
    21  	DefaultRequestIDLengthLimit = 128
    22  )
    23  
    24  // Counter used to create new request ids.
    25  var reqID int64
    26  
    27  // Common prefix to all newly created request ids for this process.
    28  var reqPrefix string
    29  
    30  // Initialize common prefix on process startup.
    31  func init() {
    32  	// algorithm taken from https://github.com/zenazn/goji/blob/master/web/middleware/request_id.go#L44-L50
    33  	var buf [12]byte
    34  	var b64 string
    35  	for len(b64) < 10 {
    36  		rand.Read(buf[:])
    37  		b64 = base64.StdEncoding.EncodeToString(buf[:])
    38  		b64 = strings.NewReplacer("+", "", "/", "").Replace(b64)
    39  	}
    40  	reqPrefix = string(b64[0:10])
    41  }
    42  
    43  // RequestIDWithHeader behaves like the middleware RequestID, but it takes the request id header
    44  // as the (first) argument.
    45  func RequestIDWithHeader(requestIDHeader string) goa.Middleware {
    46  	return RequestIDWithHeaderAndLengthLimit(requestIDHeader, DefaultRequestIDLengthLimit)
    47  }
    48  
    49  // RequestIDWithHeaderAndLengthLimit behaves like the middleware RequestID, but it takes the
    50  // request id header as the (first) argument and a length limit for truncation of the request
    51  // header value if it exceeds a reasonable length. The limit can be negative for unlimited.
    52  func RequestIDWithHeaderAndLengthLimit(requestIDHeader string, lengthLimit int) goa.Middleware {
    53  	return func(h goa.Handler) goa.Handler {
    54  		return func(ctx context.Context, rw http.ResponseWriter, req *http.Request) error {
    55  			id := req.Header.Get(requestIDHeader)
    56  			if id == "" {
    57  				id = fmt.Sprintf("%s-%d", reqPrefix, atomic.AddInt64(&reqID, 1))
    58  			} else if lengthLimit >= 0 && len(id) > lengthLimit {
    59  				id = id[:lengthLimit]
    60  			}
    61  			ctx = context.WithValue(ctx, reqIDKey, id)
    62  
    63  			return h(ctx, rw, req)
    64  		}
    65  	}
    66  }
    67  
    68  // RequestID is a middleware that injects a request ID into the context of each request.
    69  // Retrieve it using ctx.Value(ReqIDKey). If the incoming request has a RequestIDHeader header then
    70  // that value is used else a random value is generated.
    71  func RequestID() goa.Middleware {
    72  	return RequestIDWithHeader(RequestIDHeader)
    73  }
    74  
    75  // ContextRequestID extracts the Request ID from the context.
    76  func ContextRequestID(ctx context.Context) (reqID string) {
    77  	id := ctx.Value(reqIDKey)
    78  	if id != nil {
    79  		reqID = id.(string)
    80  	}
    81  	return
    82  }