Add tests for setRawTransactionPrivateFrom, simplify code (Consensys#…

…1245)
danielporto · Sep 7, 2021 · 1517050 · 1517050
1 parent 213e705
commit 1517050
Show file tree

Hide file tree

Showing 2 changed files with 135 additions and 12 deletions.
diff --git a/internal/ethapi/api.go b/internal/ethapi/api.go
@@ -1798,7 +1798,8 @@ func (args *PrivateTxArgs) SetDefaultPrivateFrom(ctx context.Context, b Backend)
 
 func (args *PrivateTxArgs) SetRawTransactionPrivateFrom(ctx context.Context, b Backend, tx *types.Transaction) error {
 	if args.PrivateFor != nil && b.ChainConfig().IsMPS {
-		retrievedPrivateFrom, err := getRawTransactionPrivateFrom(tx)
+		hash := common.BytesToEncryptedPayloadHash(tx.Data())
+		_, retrievedPrivateFrom, _, err := private.P.ReceiveRaw(hash)
 		if err != nil {
 			return err
 		}
@@ -1813,7 +1814,7 @@ func (args *PrivateTxArgs) SetRawTransactionPrivateFrom(ctx context.Context, b B
 			return err
 		}
 		if psm.NotIncludeAny(args.PrivateFrom) {
-			return fmt.Errorf("The PrivateFrom address does not match the specified private state (%s) ", psm.ID)
+			return fmt.Errorf("The PrivateFrom address does not match the specified private state (%s)", psm.ID)
 		}
 	}
 	return nil
@@ -2773,16 +2774,6 @@ func handlePrivateTransaction(ctx context.Context, b Backend, tx *types.Transact
 	return
 }
 
-func getRawTransactionPrivateFrom(tx *types.Transaction) (string, error) {
-	data := tx.Data()
-	hash := common.BytesToEncryptedPayloadHash(data)
-	_, privateFrom, _, revErr := private.P.ReceiveRaw(hash)
-	if revErr != nil {
-		return "", revErr
-	}
-	return privateFrom, revErr
-}
-
 // Quorum
 func handleRawPrivateTransaction(ctx context.Context, b Backend, tx *types.Transaction, privateTxArgs *PrivateTxArgs, from common.Address) (hash common.EncryptedPayloadHash, err error) {
 	data := tx.Data()

diff --git a/internal/ethapi/api_test.go b/internal/ethapi/api_test.go
@@ -37,6 +37,7 @@ import (
 	"github.com/golang/mock/gomock"
 	"github.com/jpmorganchase/quorum-security-plugin-sdk-go/proto"
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 )
 
 var (
@@ -593,6 +594,137 @@ func TestSubmitPrivateTransactionWithPrivacyMarkerEnabled(t *testing.T) {
 	assert.NotEqual(hexutil.Uint64(stbBackend.txThatWasSent.Gas()), gas, "privacy marker transaction should not have same gas value as internal private tx")
 }
 
+func TestSetRawTransactionPrivateFrom(t *testing.T) {
+	somePTMAddr := "some-ptm-addr"
+	psiID := types.PrivateStateIdentifier("myPSI")
+	mpsPTMAddrs := []string{somePTMAddr}
+
+	tests := []struct {
+		name                  string
+		receiveRawPrivateFrom string
+		argsPrivateFrom       string
+		wantPrivateFrom       string
+	}{
+		{
+			name:                  "receiveRawPrivateFromIfNoArgPrivateFrom",
+			receiveRawPrivateFrom: somePTMAddr,
+			argsPrivateFrom:       "",
+			wantPrivateFrom:       somePTMAddr,
+		},
+		{
+			name:                  "argPrivateFromOnly",
+			receiveRawPrivateFrom: "",
+			argsPrivateFrom:       somePTMAddr,
+			wantPrivateFrom:       somePTMAddr,
+		},
+		{
+			name:                  "equalArgAndReceiveRawPrivateFrom",
+			receiveRawPrivateFrom: somePTMAddr,
+			argsPrivateFrom:       somePTMAddr,
+			wantPrivateFrom:       somePTMAddr,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			ctrl := gomock.NewController(t)
+			defer ctrl.Finish()
+
+			savedPTM := private.P
+			defer func() { private.P = savedPTM }()
+
+			mockPTM := private.NewMockPrivateTransactionManager(ctrl)
+			mockPTM.EXPECT().ReceiveRaw(gomock.Any()).Return(nil, somePTMAddr, nil, nil).Times(1)
+			private.P = mockPTM
+
+			psm := mps.NewPrivateStateMetadata(psiID, "", "", 0, mpsPTMAddrs)
+
+			mockPSMR := mps.NewMockPrivateStateMetadataResolver(ctrl)
+			mockPSMR.EXPECT().ResolveForUserContext(gomock.Any()).Return(psm, nil).Times(1)
+
+			b := &MPSStubBackend{
+				psmr: mockPSMR,
+			}
+
+			tx := types.NewTransaction(0, common.Address{}, nil, 0, nil, []byte("ptm-hash"))
+
+			args := &PrivateTxArgs{
+				PrivateFor:  []string{"some-ptm-recipient"},
+				PrivateFrom: tt.argsPrivateFrom,
+			}
+
+			err := args.SetRawTransactionPrivateFrom(context.Background(), b, tx)
+
+			require.NoError(t, err)
+			require.Equal(t, tt.wantPrivateFrom, args.PrivateFrom)
+		})
+	}
+}
+
+func TestSetRawTransactionPrivateFrom_DifferentArgPrivateFromAndReceiveRawPrivateFrom(t *testing.T) {
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+
+	savedPTM := private.P
+	defer func() { private.P = savedPTM }()
+
+	receiveRawPrivateFrom := "some-ptm-addr"
+	argsPrivateFrom := "other-ptm-addr"
+
+	mockPTM := private.NewMockPrivateTransactionManager(ctrl)
+	mockPTM.EXPECT().ReceiveRaw(gomock.Any()).Return(nil, receiveRawPrivateFrom, nil, nil).Times(1)
+	private.P = mockPTM
+
+	b := &MPSStubBackend{}
+
+	tx := types.NewTransaction(0, common.Address{}, nil, 0, nil, []byte("ptm-hash"))
+
+	args := &PrivateTxArgs{
+		PrivateFor:  []string{"some-ptm-recipient"},
+		PrivateFrom: argsPrivateFrom,
+	}
+
+	err := args.SetRawTransactionPrivateFrom(context.Background(), b, tx)
+
+	require.EqualError(t, err, "The PrivateFrom address retrieved from the privacy manager does not match private PrivateFrom (other-ptm-addr) specified in transaction arguments.")
+}
+
+func TestSetRawTransactionPrivateFrom_ResolvePrivateFromIsNotMPSTenantAddr(t *testing.T) {
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+
+	savedPTM := private.P
+	defer func() { private.P = savedPTM }()
+
+	receiveRawPrivateFrom := "some-ptm-addr"
+	psiID := types.PrivateStateIdentifier("myPSI")
+
+	mpsPTMAddrs := []string{"other-ptm-addr"}
+
+	mockPTM := private.NewMockPrivateTransactionManager(ctrl)
+	mockPTM.EXPECT().ReceiveRaw(gomock.Any()).Return(nil, receiveRawPrivateFrom, nil, nil).Times(1)
+	private.P = mockPTM
+
+	psm := mps.NewPrivateStateMetadata(psiID, "", "", 0, mpsPTMAddrs)
+
+	mockPSMR := mps.NewMockPrivateStateMetadataResolver(ctrl)
+	mockPSMR.EXPECT().ResolveForUserContext(gomock.Any()).Return(psm, nil).Times(1)
+
+	b := &MPSStubBackend{
+		psmr: mockPSMR,
+	}
+
+	tx := types.NewTransaction(0, common.Address{}, nil, 0, nil, []byte("ptm-hash"))
+
+	args := &PrivateTxArgs{
+		PrivateFor: []string{"some-ptm-recipient"},
+	}
+
+	err := args.SetRawTransactionPrivateFrom(context.Background(), b, tx)
+
+	require.EqualError(t, err, "The PrivateFrom address does not match the specified private state (myPSI)")
+}
+
 func createKeystore(t *testing.T) (*keystore.KeyStore, accounts.Account, accounts.Account) {
 	assert := assert.New(t)