Merge pull request #2355 from onflow/bastian/2353-allow-account-linki…

…ng-pragma

runtime/account_test.go

@@ -2453,7 +2453,7 @@ func TestRuntimeAccountLink(t *testing.T) {
 
 	t.Parallel()
 
-	t.Run("disabled", func(t *testing.T) {
+	t.Run("disabled, missing pragma", func(t *testing.T) {
 
 		t.Parallel()
 
@@ -2524,7 +2524,78 @@ func TestRuntimeAccountLink(t *testing.T) {
 		assert.ErrorContains(t, err, "value of type `AuthAccount` has no member `linkAccount`")
 	})
 
-	t.Run("enabled", func(t *testing.T) {
+	t.Run("enabled, missing pragma", func(t *testing.T) {
+
+		t.Parallel()
+
+		runtime := NewInterpreterRuntime(Config{
+			AtreeValidationEnabled: true,
+			AccountLinkingEnabled:  true,
+		})
+
+		address := common.MustBytesToAddress([]byte{0x1})
+
+		accountCodes := map[Location][]byte{}
+		var logs []string
+
+		signerAccount := address
+
+		runtimeInterface := &testRuntimeInterface{
+			getCode: func(location Location) (bytes []byte, err error) {
+				return accountCodes[location], nil
+			},
+			storage: newTestLedger(nil, nil),
+			getSigningAccounts: func() ([]Address, error) {
+				return []Address{signerAccount}, nil
+			},
+			resolveLocation: singleIdentifierLocationResolver(t),
+			getAccountContractCode: func(address Address, name string) (code []byte, err error) {
+				location := common.AddressLocation{
+					Address: address,
+					Name:    name,
+				}
+				return accountCodes[location], nil
+			},
+			updateAccountContractCode: func(address Address, name string, code []byte) (err error) {
+				location := common.AddressLocation{
+					Address: address,
+					Name:    name,
+				}
+				accountCodes[location] = code
+				return nil
+			},
+			log: func(message string) {
+				logs = append(logs, message)
+			},
+		}
+
+		nextTransactionLocation := newTransactionLocationGenerator()
+
+		// Set up account
+
+		setupTransaction := []byte(`
+          transaction {
+              prepare(acct: AuthAccount) {
+                  acct.linkAccount(/public/foo)
+              }
+          }
+        `)
+
+		err := runtime.ExecuteTransaction(
+			Script{
+				Source: setupTransaction,
+			},
+			Context{
+				Interface: runtimeInterface,
+				Location:  nextTransactionLocation(),
+			},
+		)
+		require.Error(t, err)
+
+		assert.ErrorContains(t, err, "value of type `AuthAccount` has no member `linkAccount`")
+	})
+
+	t.Run("enabled, pragma", func(t *testing.T) {
 
 		t.Parallel()
 
@@ -2575,6 +2646,8 @@ func TestRuntimeAccountLink(t *testing.T) {
 		// Set up account
 
 		setupTransaction := []byte(`
+          #allowAccountLinking
+
           transaction {
               prepare(acct: AuthAccount) {
                   acct.linkAccount(/public/foo)
@@ -2683,6 +2756,8 @@ func TestRuntimeAccountLink(t *testing.T) {
 		// Set up account
 
 		setupTransaction := []byte(`
+          #allowAccountLinking
+
           transaction {
               prepare(acct: AuthAccount) {
                   let cap = acct.linkAccount(/private/foo)!

runtime/common/declarationkind.go

@@ -20,6 +20,7 @@ package common
 
 import (
 	"encoding/json"
+	"math"
 
 	"github.com/onflow/cadence/runtime/errors"
 )
@@ -190,3 +191,26 @@ func (k DeclarationKind) Keywords() string {
 func (k DeclarationKind) MarshalJSON() ([]byte, error) {
 	return json.Marshal(k.String())
 }
+
+type DeclarationKindSet uint64
+
+const (
+	EmptyDeclarationKindSet DeclarationKindSet = 0
+	AllDeclarationKindsSet  DeclarationKindSet = math.MaxUint64
+)
+
+func NewDeclarationKindSet(declarationKinds ...DeclarationKind) DeclarationKindSet {
+	var set DeclarationKindSet
+	for _, declarationKind := range declarationKinds {
+		set = set.With(declarationKind)
+	}
+	return set
+}
+
+func (s DeclarationKindSet) With(kind DeclarationKind) DeclarationKindSet {
+	return s | DeclarationKindSet(1<<kind)
+}
+
+func (s DeclarationKindSet) Has(kind DeclarationKind) bool {
+	return s&(1<<kind) != 0
+}

runtime/runtime.go

@@ -164,28 +164,29 @@ type Runtime interface {
 
 type ImportResolver = func(location Location) (program *ast.Program, e error)
 
-var validTopLevelDeclarationsInTransaction = []common.DeclarationKind{
+var validTopLevelDeclarationsInTransaction = common.NewDeclarationKindSet(
+	common.DeclarationKindPragma,
 	common.DeclarationKindImport,
 	common.DeclarationKindFunction,
 	common.DeclarationKindTransaction,
-}
+)
 
-var validTopLevelDeclarationsInAccountCode = []common.DeclarationKind{
+var validTopLevelDeclarationsInAccountCode = common.NewDeclarationKindSet(
 	common.DeclarationKindPragma,
 	common.DeclarationKindImport,
 	common.DeclarationKindContract,
 	common.DeclarationKindContractInterface,
-}
+)
 
-func validTopLevelDeclarations(location Location) []common.DeclarationKind {
+func validTopLevelDeclarations(location Location) common.DeclarationKindSet {
 	switch location.(type) {
 	case common.TransactionLocation:
 		return validTopLevelDeclarationsInTransaction
 	case common.AddressLocation:
 		return validTopLevelDeclarationsInAccountCode
 	}
 
-	return nil
+	return common.AllDeclarationKindsSet
 }
 
 func reportMetric(

runtime/sema/check_pragma.go

@@ -20,38 +20,73 @@ package sema
 
 import "github.com/onflow/cadence/runtime/ast"
 
-func (checker *Checker) VisitPragmaDeclaration(p *ast.PragmaDeclaration) (_ struct{}) {
+func (checker *Checker) VisitPragmaDeclaration(declaration *ast.PragmaDeclaration) (_ struct{}) {
 
-	switch e := p.Expression.(type) {
+	switch expression := declaration.Expression.(type) {
 	case *ast.InvocationExpression:
 		// Type arguments are not supported for pragmas
-		if len(e.TypeArguments) > 0 {
+		if len(expression.TypeArguments) > 0 {
 			checker.report(&InvalidPragmaError{
-				Message: "type arguments not supported",
-				Range:   ast.NewRangeFromPositioned(checker.memoryGauge, e),
+				Message: "type arguments are not supported",
+				Range: ast.NewRangeFromPositioned(
+					checker.memoryGauge,
+					expression.TypeArguments[0],
+				),
 			})
 		}
 
 		// Ensure arguments are string expressions
-		for _, arg := range e.Arguments {
+		for _, arg := range expression.Arguments {
 			_, ok := arg.Expression.(*ast.StringExpression)
 			if !ok {
 				checker.report(&InvalidPragmaError{
 					Message: "invalid non-string argument",
-					Range:   ast.NewRangeFromPositioned(checker.memoryGauge, e),
+					Range: ast.NewRangeFromPositioned(
+						checker.memoryGauge,
+						arg.Expression,
+					),
 				})
 			}
 		}
 
 	case *ast.IdentifierExpression:
-		// valid, NO-OP
+		if IsAllowAccountLinkingPragma(declaration) {
+			checker.reportInvalidNonHeaderPragma(declaration)
+		}
 
 	default:
 		checker.report(&InvalidPragmaError{
 			Message: "pragma must be identifier or invocation expression",
-			Range:   ast.NewRangeFromPositioned(checker.memoryGauge, p.Expression),
+			Range: ast.NewRangeFromPositioned(
+				checker.memoryGauge,
+				declaration,
+			),
 		})
 	}
 
 	return
 }
+
+func (checker *Checker) reportInvalidNonHeaderPragma(declaration *ast.PragmaDeclaration) {
+	checker.report(&InvalidPragmaError{
+		Message: "pragma must appear at top-level, before all other declarations",
+		Range: ast.NewRangeFromPositioned(
+			checker.memoryGauge,
+			declaration,
+		),
+	})
+}
+
+// allowAccountLinkingPragmaIdentifier is the identifier that needs to be used in a pragma to allow account linking.
+// This is a temporary feature.
+const allowAccountLinkingPragmaIdentifier = "allowAccountLinking"
+
+func IsAllowAccountLinkingPragma(declaration *ast.PragmaDeclaration) bool {
+	identifierExpression, ok := declaration.Expression.(*ast.IdentifierExpression)
+	if !ok {
+		return false
+	}
+
+	return identifierExpression.Identifier.Identifier ==
+		allowAccountLinkingPragmaIdentifier
+}