support Grpc-Timeout header for passing to the backend

Author: Michal Witkowski

README.md

@@ -198,6 +198,7 @@ More examples are available under `examples` directory.
 * Mapping streaming APIs to JSON streams
 * Mapping HTTP headers with `Grpc-Metadata-` prefix to gRPC metadata
 * Optionally emitting API definition for [Swagger](http://swagger.io).
+* Setting [gRPC timeouts](http://www.grpc.io/docs/guides/wire.html) through inbound HTTP `Grpc-Timeout` header.
 
 ### Want to support
 But not yet.

runtime/context.go

@@ -1,20 +1,30 @@
 package runtime
 
 import (
+	"fmt"
+	"net"
 	"net/http"
+	"strconv"
 	"strings"
-
-	"net"
+	"time"
 
 	"golang.org/x/net/context"
 	"google.golang.org/grpc/metadata"
 )
 
 const metadataHeaderPrefix = "Grpc-Metadata-"
 const metadataTrailerPrefix = "Grpc-Trailer-"
+const metadataGrpcTimeout = "Grpc-Timeout"
+
 const xForwardedFor = "X-Forwarded-For"
 const xForwardedHost = "X-Forwarded-Host"
 
+var (
+	// DefaultContextTimeout is used for gRPC call context.WithTimeout whenever a Grpc-Timeout inbound
+	// header isn't present. If the value is 0 the sent `context` will not have a timeout.
+	DefaultContextTimeout = 0 * time.Second
+)
+
 /*
 AnnotateContext adds context information such as metadata from the request.
 
@@ -23,6 +33,10 @@ will be the same context.
 */
 func AnnotateContext(ctx context.Context, req *http.Request) context.Context {
 	var pairs []string
+	timeout := DefaultContextTimeout
+	if tm := req.Header.Get(metadataGrpcTimeout); tm != "" {
+		timeout, _ = timeoutDecode(tm)
+	}
 	for key, vals := range req.Header {
 		for _, val := range vals {
 			if key == "Authorization" {
@@ -47,7 +61,9 @@ func AnnotateContext(ctx context.Context, req *http.Request) context.Context {
 			pairs = append(pairs, strings.ToLower(xForwardedFor), req.Header.Get(xForwardedFor)+", "+remoteIp)
 		}
 	}
-
+	if timeout != 0 {
+		ctx, _ = context.WithTimeout(ctx, timeout)
+	}
 	if len(pairs) == 0 {
 		return ctx
 	}
@@ -72,3 +88,38 @@ func ServerMetadataFromContext(ctx context.Context) (md ServerMetadata, ok bool)
 	md, ok = ctx.Value(serverMetadataKey{}).(ServerMetadata)
 	return
 }
+
+func timeoutDecode(s string) (time.Duration, error) {
+	size := len(s)
+	if size < 2 {
+		return 0, fmt.Errorf("timeout string is too short: %q", s)
+	}
+	d, ok := timeoutUnitToDuration(s[size-1])
+	if !ok {
+		return 0, fmt.Errorf("timeout unit is not recognized: %q", s)
+	}
+	t, err := strconv.ParseInt(s[:size-1], 10, 64)
+	if err != nil {
+		return 0, err
+	}
+	return d * time.Duration(t), nil
+}
+
+func timeoutUnitToDuration(u uint8) (d time.Duration, ok bool) {
+	switch u {
+	case 'H':
+		return time.Hour, true
+	case 'M':
+		return time.Minute, true
+	case 'S':
+		return time.Second, true
+	case 'm':
+		return time.Millisecond, true
+	case 'u':
+		return time.Microsecond, true
+	case 'n':
+		return time.Nanosecond, true
+	default:
+	}
+	return
+}

runtime/context_test.go

@@ -4,6 +4,7 @@ import (
 	"net/http"
 	"reflect"
 	"testing"
+	"time"
 
 	"github.com/gengo/grpc-gateway/runtime"
 	"golang.org/x/net/context"
@@ -83,3 +84,22 @@ func TestAnnotateContext_XForwardedFor(t *testing.T) {
 		t.Errorf("md[\"x-forwarded-for\"] = %v want %v", got, want)
 	}
 }
+
+func TestAnnotateContext_SupportsTimeouts(t *testing.T) {
+	ctx := context.Background()
+	request, _ := http.NewRequest("GET", "http://bar.foo.example.com", nil)
+	annotated := runtime.AnnotateContext(ctx, request)
+	if _, exists := annotated.Deadline(); exists {
+		t.Errorf("by default Timeouts should not be set")
+	}
+	runtime.DefaultContextTimeout = 10 * time.Second
+	annotated = runtime.AnnotateContext(ctx, request)
+	if _, exists := annotated.Deadline(); !exists {
+		t.Errorf("if DefaultContextTimeout is set, the annotated context should have a timeout")
+	}
+	request.Header.Add("Grpc-Timeout", "50S")
+	annotated = runtime.AnnotateContext(ctx, request)
+	if deadline, _ := annotated.Deadline(); deadline.Sub(time.Now()) < 2*time.Second { // in case time.Now() drifts
+		t.Errorf("if Grpc-Timeout=50s is present, the timeout should be 50s")
+	}
+}