fix db getting corrupt on reorgs because of token data getting lost

Author: OPReturnCode

blockchain/chainio.go

@@ -666,16 +666,23 @@ func serializeUtxoEntry(entry *UtxoEntry) ([]byte, error) {
 		return nil, err
 	}
 
+	pkScript := entry.PkScript()
+	if !entry.tokenData.IsEmpty() {
+		buf := entry.tokenData.TokenDataBuffer()
+		buf.Write(pkScript)
+		pkScript = buf.Bytes()
+	}
+
 	// Calculate the size needed to serialize the entry.
 	size := serializeSizeVLQ(headerCode) +
-		compressedTxOutSize(uint64(entry.Amount()), entry.PkScript())
+		compressedTxOutSize(uint64(entry.Amount()), pkScript)
 
 	// Serialize the header code followed by the compressed unspent
 	// transaction output.
 	serialized := make([]byte, size)
 	offset := putVLQ(serialized, headerCode)
 	putCompressedTxOut(serialized[offset:], uint64(entry.Amount()),
-		entry.PkScript())
+		pkScript)
 
 	return serialized, nil
 }
@@ -710,6 +717,8 @@ func DeserializeUtxoEntry(serialized []byte) (*UtxoEntry, error) {
 		blockHeight: blockHeight,
 		packedFlags: 0,
 	}
+	entry.pkScript, _ = entry.tokenData.SeparateTokenDataFromPKScriptIfExists(entry.pkScript, 0)
+
 	if isCoinBase {
 		entry.packedFlags |= tfCoinBase
 	}
@@ -744,6 +753,7 @@ func deserializeUtxoCommitmentFormat(serialized []byte) (*wire.OutPoint, *UtxoEn
 		pkScript:    pkScript,
 		packedFlags: 0,
 	}
+	entry.pkScript, _ = entry.tokenData.SeparateTokenDataFromPKScriptIfExists(entry.pkScript, 0)
 	if coinbaseFlag > 0 {
 		entry.packedFlags |= tfCoinBase
 	}

blockchain/compress.go

@@ -253,9 +253,25 @@ func isPubKey(script []byte) (bool, []byte) {
 	return false, nil
 }
 
+// isCashTokensScript returns whether or not the script is a
+// cashtokens script along with the script itself if it is.
+// We can do better compression on this.
+func isCashTokensScript(script []byte) (bool, []byte) {
+	if len(script) > 0 && script[0] == txscript.SPECIAL_TOKEN_PREFIX {
+		return true, script
+	}
+	return false, nil
+}
+
 // compressedScriptSize returns the number of bytes the passed script would take
 // when encoded with the domain specific compression algorithm described above.
 func compressedScriptSize(pkScript []byte) int {
+
+	if valid, _ := isCashTokensScript(pkScript); valid {
+		return serializeSizeVLQ(uint64(len(pkScript)+numSpecialScripts)) +
+			len(pkScript)
+	}
+
 	// Pay-to-pubkey-hash script.
 	if valid, _ := isPubKeyHash(pkScript); valid {
 		return 21
@@ -311,34 +327,38 @@ func decodeCompressedScriptSize(serialized []byte) int {
 // handle the number of bytes returned by the compressedScriptSize function or
 // it will panic.
 func putCompressedScript(target, pkScript []byte) int {
-	// Pay-to-pubkey-hash script.
-	if valid, hash := isPubKeyHash(pkScript); valid {
-		target[0] = cstPayToPubKeyHash
-		copy(target[1:21], hash)
-		return 21
-	}
 
-	// Pay-to-script-hash script.
-	if valid, hash := isScriptHash(pkScript); valid {
-		target[0] = cstPayToScriptHash
-		copy(target[1:21], hash)
-		return 21
-	}
+	if valid, _ := isCashTokensScript(pkScript); !valid {
 
-	// Pay-to-pubkey (compressed or uncompressed) script.
-	if valid, serializedPubKey := isPubKey(pkScript); valid {
-		pubKeyFormat := serializedPubKey[0]
-		switch pubKeyFormat {
-		case 0x02, 0x03:
-			target[0] = pubKeyFormat
-			copy(target[1:33], serializedPubKey[1:33])
-			return 33
-		case 0x04:
-			// Encode the oddness of the serialized pubkey into the
-			// compressed script type.
-			target[0] = pubKeyFormat | (serializedPubKey[64] & 0x01)
-			copy(target[1:33], serializedPubKey[1:33])
-			return 33
+		// Pay-to-pubkey-hash script.
+		if valid, hash := isPubKeyHash(pkScript); valid {
+			target[0] = cstPayToPubKeyHash
+			copy(target[1:21], hash)
+			return 21
+		}
+
+		// Pay-to-script-hash script.
+		if valid, hash := isScriptHash(pkScript); valid {
+			target[0] = cstPayToScriptHash
+			copy(target[1:21], hash)
+			return 21
+		}
+
+		// Pay-to-pubkey (compressed or uncompressed) script.
+		if valid, serializedPubKey := isPubKey(pkScript); valid {
+			pubKeyFormat := serializedPubKey[0]
+			switch pubKeyFormat {
+			case 0x02, 0x03:
+				target[0] = pubKeyFormat
+				copy(target[1:33], serializedPubKey[1:33])
+				return 33
+			case 0x04:
+				// Encode the oddness of the serialized pubkey into the
+				// compressed script type.
+				target[0] = pubKeyFormat | (serializedPubKey[64] & 0x01)
+				copy(target[1:33], serializedPubKey[1:33])
+				return 33
+			}
 		}
 	}
 

blockchain/upgrade.go

@@ -428,6 +428,7 @@ func deserializeUtxoEntryV0(serialized []byte) (map[uint32]*UtxoEntry, error) {
 			blockHeight: int32(blockHeight),
 			packedFlags: packedFlags,
 		}
+		entries[outputIndex].pkScript, _ = entries[outputIndex].tokenData.SeparateTokenDataFromPKScriptIfExists(entries[outputIndex].pkScript, 0)
 	}
 
 	return entries, nil

blockchain/utxocache.go

@@ -88,7 +88,7 @@ func NewUtxoEntry(txOut *wire.TxOut, blockHeight int32, isCoinbase bool) *UtxoEn
 	if isCoinbase {
 		cbFlag |= tfCoinBase
 	}
-
+	txOut.PkScript, _ = txOut.TokenData.SeparateTokenDataFromPKScriptIfExists(txOut.PkScript, 0)
 	return &UtxoEntry{
 		amount:      txOut.Value,
 		pkScript:    txOut.PkScript,

blockchain/utxoviewpoint.go

@@ -85,6 +85,7 @@ func (view *UtxoViewpoint) addTxOut(outpoint wire.OutPoint, txOut *wire.TxOut, i
 	entry.amount = txOut.Value
 	entry.pkScript = pkScript
 	entry.blockHeight = blockHeight
+	entry.pkScript, _ = entry.tokenData.SeparateTokenDataFromPKScriptIfExists(pkScript, 0)
 	entry.packedFlags = tfModified
 	if isCoinBase {
 		entry.packedFlags |= tfCoinBase
@@ -185,6 +186,7 @@ func addTxOuts(view utxoView, tx *bchutil.Tx, blockHeight int32, overwrite bool)
 		// Create a new entry from the output.
 		pkScript := make([]byte, len(txOut.PkScript))
 		copy(pkScript, txOut.PkScript)
+		pkScript, _ = txOut.TokenData.SeparateTokenDataFromPKScriptIfExists(pkScript, 0)
 
 		entry := &UtxoEntry{
 			amount:      txOut.Value,
@@ -234,6 +236,12 @@ func spendTransactionInputs(view utxoView, tx *bchutil.Tx, stxos *[]SpentTxOut)
 			pkScript := make([]byte, len(entry.PkScript()))
 			copy(pkScript, entry.PkScript())
 
+			if !entry.tokenData.IsEmpty() {
+				buf := entry.tokenData.TokenDataBuffer()
+				buf.Write(pkScript)
+				pkScript = buf.Bytes()
+			}
+
 			// Populate the stxo details using the utxo entry.
 			var stxo = SpentTxOut{
 				Amount:     entry.Amount(),
@@ -347,6 +355,7 @@ func disconnectTransactions(view utxoView, block *bchutil.Block, stxos []SpentTx
 				blockHeight: stxo.Height,
 				packedFlags: tfModified,
 			}
+			entry.pkScript, _ = entry.tokenData.SeparateTokenDataFromPKScriptIfExists(entry.pkScript, 0)
 			if stxo.IsCoinBase {
 				entry.packedFlags |= tfCoinBase
 			}
@@ -388,7 +397,7 @@ func disconnectTransactions(view utxoView, block *bchutil.Block, stxos []SpentTx
 
 			pkScript := make([]byte, len(txOut.PkScript))
 			copy(pkScript, txOut.PkScript)
-
+			pkScript, _ = txOut.TokenData.SeparateTokenDataFromPKScriptIfExists(pkScript, 0)
 			// Mark the entry as spent.  To make sure the view has the entry,
 			// create one to pass along.
 			entry := &UtxoEntry{