fix(logic): ensure keepers in interpreter exec ctx

axone-protocol · Feb 4, 2023 · 901d7f2 · 901d7f2
1 parent 7f1d60f
commit 901d7f2
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 11 deletions.
diff --git a/x/logic/keeper/grpc_query_ask.go b/x/logic/keeper/grpc_query_ask.go
@@ -6,7 +6,6 @@ import (
 	sdkerrors "cosmossdk.io/errors"
 	sdk "github.com/cosmos/cosmos-sdk/types"
 	"github.com/okp4/okp4d/x/logic/types"
-	"golang.org/x/net/context"
 )
 
 func (k Keeper) Ask(ctx goctx.Context, req *types.QueryServiceAskRequest) (response *types.QueryServiceAskResponse, err error) {
@@ -39,7 +38,7 @@ func (k Keeper) Ask(ctx goctx.Context, req *types.QueryServiceAskRequest) (respo
 
 	//nolint:contextcheck
 	return k.execute(
-		context.WithValue(sdkCtx, sdk.SdkContextKey, sdkCtx),
+		sdkCtx,
 		req.Program,
 		req.Query)
 }

diff --git a/x/logic/keeper/interpreter.go b/x/logic/keeper/interpreter.go
@@ -18,18 +18,28 @@ func (k Keeper) limits(ctx goctx.Context) types.Limits {
 	return params.GetLimits()
 }
 
-func (k Keeper) execute(goctx goctx.Context, program, query string) (*types.QueryServiceAskResponse, error) {
-	sdkCtx := sdk.UnwrapSDKContext(goctx)
-	i, err := k.newInterpreter(goctx)
+func (k Keeper) enhanceContext(ctx goctx.Context) goctx.Context {
+	sdkCtx := sdk.UnwrapSDKContext(ctx)
+	sdkCtx = sdkCtx.WithValue(types.AuthKeeperContextKey, k.authKeeper)
+	sdkCtx = sdkCtx.WithValue(types.BankKeeperContextKey, k.bankKeeper)
+
+	return sdkCtx
+}
+
+func (k Keeper) execute(ctx goctx.Context, program, query string) (*types.QueryServiceAskResponse, error) {
+	ctx = k.enhanceContext(ctx)
+	sdkCtx := sdk.UnwrapSDKContext(ctx)
+
+	i, err := k.newInterpreter(ctx)
 	if err != nil {
 		return nil, sdkerrors.Wrapf(types.Internal, "error creating interpreter: %v", err.Error())
 	}
 
-	if err := i.ExecContext(goctx, program); err != nil {
+	if err := i.ExecContext(ctx, program); err != nil {
 		return nil, sdkerrors.Wrapf(types.InvalidArgument, "error compiling query: %v", err.Error())
 	}
 
-	sols, err := i.QueryContext(goctx, query)
+	sols, err := i.QueryContext(ctx, query)
 	if err != nil {
 		return nil, sdkerrors.Wrapf(types.InvalidArgument, "error executing query: %v", err.Error())
 	}
@@ -38,7 +48,7 @@ func (k Keeper) execute(goctx goctx.Context, program, query string) (*types.Quer
 	}()
 
 	success := false
-	limits := k.limits(goctx)
+	limits := k.limits(ctx)
 	var variables []string
 	results := make([]types.Result, 0)
 	for nb := sdkmath.ZeroUint(); nb.LT(*limits.MaxResultCount) && sols.Next(); nb = nb.Incr() {
@@ -81,9 +91,6 @@ func checkLimits(request *types.QueryServiceAskRequest, limits types.Limits) err
 // newInterpreter creates a new interpreter properly configured.
 func (k Keeper) newInterpreter(ctx goctx.Context) (*prolog.Interpreter, error) {
 	sdkctx := sdk.UnwrapSDKContext(ctx)
-	sdkctx = sdkctx.WithValue(types.AuthKeeperContextKey, k.authKeeper)
-	sdkctx = sdkctx.WithValue(types.BankKeeperContextKey, k.bankKeeper)
-
 	params := k.GetParams(sdkctx)
 
 	interpreterParams := params.GetInterpreter()