From 21fd15fa3624ac538e28b567d19ccc8c60ff027c Mon Sep 17 00:00:00 2001
From: wz <mars@bytom.io>
Date: Mon, 20 May 2019 15:48:40 +0800
Subject: [PATCH] fix votetx for validation (#75)

---
 api/block_retrieve.go            | 10 +++++++---
 api/query.go                     |  2 ++
 protocol/validation/block.go     | 13 +++++++++----
 protocol/validation/vmcontext.go | 27 +++++++++++++++++++--------
 wallet/utxo.go                   | 11 ++++++++---
 5 files changed, 45 insertions(+), 18 deletions(-)

diff --git a/api/block_retrieve.go b/api/block_retrieve.go
index 4b85fde8..27897ffd 100644
--- a/api/block_retrieve.go
+++ b/api/block_retrieve.go
@@ -93,10 +93,14 @@ func (a *API) getBlock(ins BlockReq) Response {
 		}
 
 		resOutID := orig.ResultIds[0]
-		resOut, ok := orig.Entries[*resOutID].(*bc.IntraChainOutput)
-		if ok {
+		switch resOut := orig.Entries[*resOutID].(type) {
+		case *bc.IntraChainOutput:
 			tx.MuxID = *resOut.Source.Ref
-		} else {
+		case *bc.CrossChainOutput:
+			tx.MuxID = *resOut.Source.Ref
+		case *bc.VoteOutput:
+			tx.MuxID = *resOut.Source.Ref
+		case *bc.Retirement:
 			resRetire, _ := orig.Entries[*resOutID].(*bc.Retirement)
 			tx.MuxID = *resRetire.Source.Ref
 		}
diff --git a/api/query.go b/api/query.go
index 369443d1..a3e927bf 100644
--- a/api/query.go
+++ b/api/query.go
@@ -203,6 +203,8 @@ func (a *API) getUnconfirmedTx(ctx context.Context, filter struct {
 	switch out := resOut.(type) {
 	case *bc.IntraChainOutput:
 		tx.MuxID = *out.Source.Ref
+	case *bc.VoteOutput:
+		tx.MuxID = *out.Source.Ref
 	case *bc.Retirement:
 		tx.MuxID = *out.Source.Ref
 	}
diff --git a/protocol/validation/block.go b/protocol/validation/block.go
index 3ded430b..8b63df0e 100644
--- a/protocol/validation/block.go
+++ b/protocol/validation/block.go
@@ -46,12 +46,17 @@ func checkCoinbaseAmount(b *bc.Block, amount uint64) error {
 		return errors.Wrap(ErrWrongCoinbaseTransaction, "have more than 1 output")
 	}
 
-	output, err := tx.IntraChainOutput(*tx.TxHeader.ResultIds[0])
-	if err != nil {
-		return err
+	var SourceAmount uint64
+	switch output := tx.Entries[*tx.TxHeader.ResultIds[0]].(type) {
+	case *bc.IntraChainOutput:
+		SourceAmount = output.Source.Value.Amount
+	case *bc.VoteOutput:
+		SourceAmount = output.Source.Value.Amount
+	default:
+		return errors.Wrapf(bc.ErrEntryType, "entry %x has unexpected type %T", tx.TxHeader.ResultIds[0].Bytes(), output)
 	}
 
-	if output.Source.Value.Amount != amount {
+	if SourceAmount != amount {
 		return errors.Wrap(ErrWrongCoinbaseTransaction, "dismatch output amount")
 	}
 	return nil
diff --git a/protocol/validation/vmcontext.go b/protocol/validation/vmcontext.go
index f7b62bcf..610a8eb7 100644
--- a/protocol/validation/vmcontext.go
+++ b/protocol/validation/vmcontext.go
@@ -25,15 +25,23 @@ func NewTxVMContext(vs *validationState, entry bc.Entry, prog *bc.Program, args
 	)
 
 	switch e := entry.(type) {
-
 	case *bc.Spend:
-		spentOutput := tx.Entries[*e.SpentOutputId].(*bc.IntraChainOutput)
-		a1 := spentOutput.Source.Value.AssetId.Bytes()
-		assetID = &a1
-		amount = &spentOutput.Source.Value.Amount
-		destPos = &e.WitnessDestination.Position
-		s := e.SpentOutputId.Bytes()
-		spentOutputID = &s
+		switch spentOutput := tx.Entries[*e.SpentOutputId].(type) {
+		case *bc.IntraChainOutput:
+			a1 := spentOutput.Source.Value.AssetId.Bytes()
+			assetID = &a1
+			amount = &spentOutput.Source.Value.Amount
+			destPos = &e.WitnessDestination.Position
+			s := e.SpentOutputId.Bytes()
+			spentOutputID = &s
+		case *bc.VoteOutput:
+			a1 := spentOutput.Source.Value.AssetId.Bytes()
+			assetID = &a1
+			amount = &spentOutput.Source.Value.Amount
+			destPos = &e.WitnessDestination.Position
+			s := e.SpentOutputId.Bytes()
+			spentOutputID = &s
+		}
 	}
 
 	var txSigHash *[]byte
@@ -111,6 +119,9 @@ func (ec *entryContext) checkOutput(index uint64, amount uint64, assetID []byte,
 		case *bc.IntraChainOutput:
 			return check(e.ControlProgram, e.Source.Value), nil
 
+		case *bc.VoteOutput:
+			return check(e.ControlProgram, e.Source.Value), nil
+
 		case *bc.Retirement:
 			var prog bc.Program
 			if expansion {
diff --git a/wallet/utxo.go b/wallet/utxo.go
index 20775fdd..c5d106c3 100644
--- a/wallet/utxo.go
+++ b/wallet/utxo.go
@@ -79,12 +79,17 @@ func (w *Wallet) detachUtxos(batch dbm.Batch, b *types.Block, txStatus *bc.Trans
 	for txIndex := len(b.Transactions) - 1; txIndex >= 0; txIndex-- {
 		tx := b.Transactions[txIndex]
 		for j := range tx.Outputs {
-			resOut, err := tx.IntraChainOutput(*tx.ResultIds[j])
-			if err != nil {
+			code := []byte{}
+			switch resOut := tx.Entries[*tx.ResultIds[j]].(type) {
+			case *bc.IntraChainOutput:
+				code = resOut.ControlProgram.Code
+			case *bc.VoteOutput:
+				code = resOut.ControlProgram.Code
+			default:
 				continue
 			}
 
-			if segwit.IsP2WScript(resOut.ControlProgram.Code) {
+			if segwit.IsP2WScript(code) {
 				batch.Delete(account.StandardUTXOKey(*tx.ResultIds[j]))
 			} else {
 				batch.Delete(account.ContractUTXOKey(*tx.ResultIds[j]))
-- 
2.11.0