From 9b8db6b37cbeb41b314ee172975e0ea60fca63e3 Mon Sep 17 00:00:00 2001
From: Ethan Buchman <ethan@coinculture.info>
Date: Sun, 4 Mar 2018 02:32:39 -0500
Subject: [PATCH] x/auth: clean up ante handler

---
 x/auth/ante.go | 121 ++++++++++++++++++++++++++-----------------------
 1 file changed, 65 insertions(+), 56 deletions(-)

diff --git a/x/auth/ante.go b/x/auth/ante.go
index 11aa03c0d7..17e9ccdd94 100644
--- a/x/auth/ante.go
+++ b/x/auth/ante.go
@@ -9,82 +9,91 @@ func NewAnteHandler(accountMapper sdk.AccountMapper) sdk.AnteHandler {
 		ctx sdk.Context, tx sdk.Tx,
 	) (_ sdk.Context, _ sdk.Result, abort bool) {
 
-		// Deduct the fee from the fee payer.
-		// This is done first because it only
-		// requires fetching 1 account.
-		payerAddr := tx.GetFeePayer()
-		if payerAddr != nil {
-			payerAcc := accountMapper.GetAccount(ctx, payerAddr)
-			if payerAcc == nil {
-				return ctx,
-					sdk.ErrUnrecognizedAddress(payerAddr).Result(),
-					true
-			}
-			// TODO: Charge fee from payerAcc.
-			// TODO: accountMapper.SetAccount(ctx, payerAddr)
-		} else {
-			// TODO: Ensure that some other spam prevention is used.
-		}
-
-		var sigs = tx.GetSignatures()
-
 		// Assert that there are signatures.
+		var sigs = tx.GetSignatures()
 		if len(sigs) == 0 {
 			return ctx,
 				sdk.ErrUnauthorized("no signers").Result(),
 				true
 		}
 
-		// Ensure that sigs are correct.
-		var msg = tx.GetMsg()
-		var signerAddrs = msg.GetSigners()
-		var signerAccs = make([]sdk.Account, len(signerAddrs))
+		// TODO: can tx just implement message?
+		msg := tx.GetMsg()
 
 		// Assert that number of signatures is correct.
+		var signerAddrs = msg.GetSigners()
 		if len(sigs) != len(signerAddrs) {
 			return ctx,
 				sdk.ErrUnauthorized("wrong number of signers").Result(),
 				true
 		}
 
-		// Check each nonce and sig.
-		// TODO Refactor out.
-		for i, sig := range sigs {
+		// Collect accounts to set in the context
+		var signerAccs = make([]sdk.Account, len(signerAddrs))
 
-			var signerAcc = accountMapper.GetAccount(ctx, signerAddrs[i])
+		signBytes := msg.GetSignBytes()
+
+		// First sig is the fee payer.
+		// Check sig and nonce, deduct fee.
+		// This is done first because it only
+		// requires fetching 1 account.
+		payerAcc, res := processSig(ctx, accountMapper, signerAddrs[0], sigs[0], signBytes)
+		if !res.IsOK() {
+			return ctx, res, true
+		}
+		signerAccs[0] = payerAcc
+		// TODO: Charge fee from payerAcc.
+		// TODO: accountMapper.SetAccount(ctx, payerAddr)
+
+		// Check sig and nonce for the rest.
+		for i := 1; i < len(sigs); i++ {
+			signerAcc, res := processSig(ctx, accountMapper, signerAddrs[i], sigs[i], signBytes)
+			if !res.IsOK() {
+				return ctx, res, true
+			}
 			signerAccs[i] = signerAcc
-
-			// If no pubkey, set pubkey.
-			if signerAcc.GetPubKey().Empty() {
-				err := signerAcc.SetPubKey(sig.PubKey)
-				if err != nil {
-					return ctx,
-						sdk.ErrInternal("setting PubKey on signer").Result(),
-						true
-				}
-			}
-
-			// Check and increment sequence number.
-			seq := signerAcc.GetSequence()
-			if seq != sig.Sequence {
-				return ctx,
-					sdk.ErrInvalidSequence("").Result(),
-					true
-			}
-			signerAcc.SetSequence(seq + 1)
-
-			// Check sig.
-			if !sig.PubKey.VerifyBytes(msg.GetSignBytes(), sig.Signature) {
-				return ctx,
-					sdk.ErrUnauthorized("").Result(),
-					true
-			}
-
-			// Save the account.
-			accountMapper.SetAccount(ctx, signerAcc)
 		}
 
 		ctx = WithSigners(ctx, signerAccs)
 		return ctx, sdk.Result{}, false // continue...
 	}
 }
+
+// verify the signature and increment the sequence.
+// if the account doesn't have a pubkey, set it as well.
+func processSig(ctx sdk.Context, am sdk.AccountMapper, addr sdk.Address, sig sdk.StdSignature, signBytes []byte) (acc sdk.Account, res sdk.Result) {
+
+	// Get the account
+	acc = am.GetAccount(ctx, addr)
+	if acc == nil {
+		return nil, sdk.ErrUnrecognizedAddress(addr).Result()
+	}
+
+	// Check and increment sequence number.
+	seq := acc.GetSequence()
+	if seq != sig.Sequence {
+		return nil, sdk.ErrInvalidSequence("").Result()
+	}
+	acc.SetSequence(seq + 1)
+
+	// Check and possibly set pubkey.
+	pubKey := acc.GetPubKey()
+	if pubKey.Empty() {
+		pubKey = sig.PubKey
+		err := acc.SetPubKey(pubKey)
+		if err != nil {
+			return nil, sdk.ErrInternal("setting PubKey on signer").Result()
+		}
+	}
+	// TODO: should we enforce pubKey == sig.PubKey ?
+	// If not, ppl can send useless PubKeys after first tx
+
+	// Check sig.
+	if !sig.PubKey.VerifyBytes(signBytes, sig.Signature) {
+		return nil, sdk.ErrUnauthorized("").Result()
+	}
+
+	// Save the account.
+	am.SetAccount(ctx, acc)
+	return
+}