rpc/v2: optionally passing context argument to rpc v2 api methods

1 files changed, 17 insertions, 8 deletions

diff --git a/rpc/v2/utils.go b/rpc/v2/utils.go
index a564b2473..ca37924a3 100644
--- a/rpc/v2/utils.go
+++ b/rpc/v2/utils.go
@@ -24,6 +24,8 @@ import (
 	"reflect"
 	"unicode"
 	"unicode/utf8"
+
+	"golang.org/x/net/context"
 )
 
 // Is this an exported - upper case - name?
@@ -107,6 +109,8 @@ func isBlockNumber(t reflect.Type) bool {
 	return t == blockNumberType
 }
 
+var contextType = reflect.TypeOf(new(context.Context)).Elem()
+
 // suitableCallbacks iterates over the methods of the given type. It will determine if a method satisfies the criteria
 // for a RPC callback or a subscription callback and adds it to the collection of callbacks or subscriptions. See server
 // documentation for a summary of these criteria.
@@ -129,12 +133,19 @@ METHODS:
 		h.method = method
 		h.errPos = -1
 
+		firstArg := 1
+		numIn := mtype.NumIn()
+		if numIn >= 2 && mtype.In(1) == contextType {
+			h.hasCtx = true
+			firstArg = 2
+		}
+
 		if h.isSubscribe {
-			h.argTypes = make([]reflect.Type, mtype.NumIn()-1) // skip rcvr type
-			for i := 1; i < mtype.NumIn(); i++ {
+			h.argTypes = make([]reflect.Type, numIn-firstArg) // skip rcvr type
+			for i := firstArg; i < numIn; i++ {
 				argType := mtype.In(i)
 				if isExportedOrBuiltinType(argType) {
-					h.argTypes[i-1] = argType
+					h.argTypes[i-firstArg] = argType
 				} else {
 					continue METHODS
 				}
@@ -144,17 +155,15 @@ METHODS:
 			continue METHODS
 		}
 
-		numIn := mtype.NumIn()
-
 		// determine method arguments, ignore first arg since it's the receiver type
 		// Arguments must be exported or builtin types
-		h.argTypes = make([]reflect.Type, numIn-1)
-		for i := 1; i < numIn; i++ {
+		h.argTypes = make([]reflect.Type, numIn-firstArg)
+		for i := firstArg; i < numIn; i++ {
 			argType := mtype.In(i)
 			if !isExportedOrBuiltinType(argType) {
 				continue METHODS
 			}
-			h.argTypes[i-1] = argType
+			h.argTypes[i-firstArg] = argType
 		}
 
 		// check that all returned values are exported or builtin types