<commit_before>package wallet

import (
	"math"

	"github.com/NebulousLabs/Sia/build"
	"github.com/NebulousLabs/Sia/modules"
	"github.com/NebulousLabs/Sia/types"

	"github.com/NebulousLabs/bolt"
)

// updateConfirmedSet uses a consensus change to update the confirmed set of
// outputs as understood by the wallet.
func (w *Wallet) updateConfirmedSet(cc modules.ConsensusChange) {
	for _, diff := range cc.SiacoinOutputDiffs {
		// Verify that the diff is relevant to the wallet.
		_, exists := w.keys[diff.SiacoinOutput.UnlockHash]
		if !exists {
			continue
		}

		_, exists = w.siacoinOutputs[diff.ID]
		if diff.Direction == modules.DiffApply {
			if build.DEBUG && exists {
				panic("adding an existing output to wallet")
			}
			w.siacoinOutputs[diff.ID] = diff.SiacoinOutput
		} else {
			if build.DEBUG && !exists {
				panic("deleting nonexisting output from wallet")
			}
			delete(w.siacoinOutputs, diff.ID)
		}
	}
	for _, diff := range cc.SiafundOutputDiffs {
		// Verify that the diff is relevant to the wallet.
		_, exists := w.keys[diff.SiafundOutput.UnlockHash]
		if !exists {
			continue
		}

		_, exists = w.siafundOutputs[diff.ID]
		if diff.Direction == modules.DiffApply {
			if build.DEBUG && exists {
				panic("adding an existing output to wallet")
			}
			w.siafundOutputs[diff.ID] = diff.SiafundOutput
		} else {
			if build.DEBUG && !exists {
				panic("deleting nonexisting output from wallet")
			}
			delete(w.siafundOutputs, diff.ID)
		}
	}
	for _, diff := range cc.SiafundPoolDiffs {
		if diff.Direction == modules.DiffApply {
			w.siafundPool = diff.Adjusted
		} else {
			w.siafundPool = diff.Previous
		}
	}
}

// revertHistory reverts any transaction history that was destroyed by reverted
// blocks in the consensus change.
func (w *Wallet) revertHistory(cc modules.ConsensusChange) {
	for _, block := range cc.RevertedBlocks {
		// Remove any transactions that have been reverted.
		for i := len(block.Transactions) - 1; i >= 0; i-- {
			// If the transaction is relevant to the wallet, it will be the
			// most recent transaction appended to w.processedTransactions.
			// Relevance can be determined just by looking at the last element
			// of w.processedTransactions.
			txn := block.Transactions[i]
			txid := txn.ID()
			if len(w.processedTransactions) > 0 && txid == w.processedTransactions[len(w.processedTransactions)-1].TransactionID {
				w.processedTransactions = w.processedTransactions[:len(w.processedTransactions)-1]
				delete(w.processedTransactionMap, txid)
			}
		}

		// Remove the miner payout transaction if applicable.
		for _, mp := range block.MinerPayouts {
			_, exists := w.keys[mp.UnlockHash]
			if exists {
				w.processedTransactions = w.processedTransactions[:len(w.processedTransactions)-1]
				delete(w.processedTransactionMap, types.TransactionID(block.ID()))
				break
			}
		}
		w.consensusSetHeight--
	}
}

// applyHistory applies any transaction history that was introduced by the
// applied blocks.
func (w *Wallet) applyHistory(cc modules.ConsensusChange) {
	err := w.db.Update(func(tx *bolt.Tx) error {
		for _, block := range cc.AppliedBlocks {
			w.consensusSetHeight++
			// Apply the miner payout transaction if applicable.
			minerPT := modules.ProcessedTransaction{
				Transaction:           types.Transaction{},
				TransactionID:         types.TransactionID(block.ID()),
				ConfirmationHeight:    w.consensusSetHeight,
				ConfirmationTimestamp: block.Timestamp,
			}
			relevant := false
			for i, mp := range block.MinerPayouts {
				_, exists := w.keys[mp.UnlockHash]
				if exists {
					relevant = true
				}
				minerPT.Outputs = append(minerPT.Outputs, modules.ProcessedOutput{
					FundType:       types.SpecifierMinerPayout,
					MaturityHeight: w.consensusSetHeight + types.MaturityDelay,
					WalletAddress:  exists,
					RelatedAddress: mp.UnlockHash,
					Value:          mp.Value,
				})
				err := dbPutHistoricOutput(tx, types.OutputID(block.MinerPayoutID(uint64(i))), mp.Value)
				if err != nil {
					return err
				}
			}
			if relevant {
				w.processedTransactions = append(w.processedTransactions, minerPT)
				w.processedTransactionMap[minerPT.TransactionID] = &w.processedTransactions[len(w.processedTransactions)-1]
			}
			for _, txn := range block.Transactions {
				relevant := false
				pt := modules.ProcessedTransaction{
					Transaction:           txn,
					TransactionID:         txn.ID(),
					ConfirmationHeight:    w.consensusSetHeight,
					ConfirmationTimestamp: block.Timestamp,
				}

				for _, sci := range txn.SiacoinInputs {
					_, exists := w.keys[sci.UnlockConditions.UnlockHash()]
					if exists {
						relevant = true
					}
					val, err := dbGetHistoricOutput(tx, types.OutputID(sci.ParentID))
					if err != nil {
						return err
					}
					pt.Inputs = append(pt.Inputs, modules.ProcessedInput{
						FundType:       types.SpecifierSiacoinInput,
						WalletAddress:  exists,
						RelatedAddress: sci.UnlockConditions.UnlockHash(),
						Value:          val,
					})
				}

				for i, sco := range txn.SiacoinOutputs {
					_, exists := w.keys[sco.UnlockHash]
					if exists {
						relevant = true
					}
					pt.Outputs = append(pt.Outputs, modules.ProcessedOutput{
						FundType:       types.SpecifierSiacoinOutput,
						MaturityHeight: w.consensusSetHeight,
						WalletAddress:  exists,
						RelatedAddress: sco.UnlockHash,
						Value:          sco.Value,
					})
					err := dbPutHistoricOutput(tx, types.OutputID(txn.SiacoinOutputID(uint64(i))), sco.Value)
					if err != nil {
						return err
					}
				}

				for _, sfi := range txn.SiafundInputs {
					_, exists := w.keys[sfi.UnlockConditions.UnlockHash()]
					if exists {
						relevant = true
					}
					sfiValue, err := dbGetHistoricOutput(tx, types.OutputID(sfi.ParentID))
					if err != nil {
						return err
					}
					pt.Inputs = append(pt.Inputs, modules.ProcessedInput{
						FundType:       types.SpecifierSiafundInput,
						WalletAddress:  exists,
						RelatedAddress: sfi.UnlockConditions.UnlockHash(),
						Value:          sfiValue,
					})
					startVal, err := dbGetHistoricClaimStart(tx, sfi.ParentID)
					if err != nil {
						return err
					}
					claimValue := w.siafundPool.Sub(startVal).Mul(sfiValue)
					pt.Outputs = append(pt.Outputs, modules.ProcessedOutput{
						FundType:       types.SpecifierClaimOutput,
						MaturityHeight: w.consensusSetHeight + types.MaturityDelay,
						WalletAddress:  exists,
						RelatedAddress: sfi.ClaimUnlockHash,
						Value:          claimValue,
					})
				}

				for i, sfo := range txn.SiafundOutputs {
					_, exists := w.keys[sfo.UnlockHash]
					if exists {
						relevant = true
					}
					pt.Outputs = append(pt.Outputs, modules.ProcessedOutput{
						FundType:       types.SpecifierSiafundOutput,
						MaturityHeight: w.consensusSetHeight,
						WalletAddress:  exists,
						RelatedAddress: sfo.UnlockHash,
						Value:          sfo.Value,
					})
					id := txn.SiafundOutputID(uint64(i))
					err := dbPutHistoricOutput(tx, types.OutputID(id), sfo.Value)
					if err != nil {
						return err
					}
					err = dbPutHistoricClaimStart(tx, id, sfo.ClaimStart)
					if err != nil {
						return err
					}
				}

				for _, fee := range txn.MinerFees {
					pt.Outputs = append(pt.Outputs, modules.ProcessedOutput{
						FundType: types.SpecifierMinerFee,
						Value:    fee,
					})
				}
				if relevant {
					w.processedTransactions = append(w.processedTransactions, pt)
					w.processedTransactionMap[pt.TransactionID] = &w.processedTransactions[len(w.processedTransactions)-1]
				}
			}
		}

		return nil
	})
	if err != nil {
		w.log.Println("ERROR: failed to add consensus change:", err)
	}
}

// ProcessConsensusChange parses a consensus change to update the set of
// confirmed outputs known to the wallet.
func (w *Wallet) ProcessConsensusChange(cc modules.ConsensusChange) {
	if err := w.tg.Add(); err != nil {
		// The wallet should gracefully reject updates from the consensus set
		// or transaction pool that are sent after the wallet's Close method
		// has closed the wallet's ThreadGroup.
		return
	}
	defer w.tg.Done()
	w.mu.Lock()
	defer w.mu.Unlock()
	w.updateConfirmedSet(cc)
	w.revertHistory(cc)
	w.applyHistory(cc)
}

// ReceiveUpdatedUnconfirmedTransactions updates the wallet's unconfirmed
// transaction set.
func (w *Wallet) ReceiveUpdatedUnconfirmedTransactions(txns []types.Transaction, _ modules.ConsensusChange) {
	if err := w.tg.Add(); err != nil {
		// Gracefully reject transactions if the wallet's Close method has
		// closed the wallet's ThreadGroup already.
		return
	}
	defer w.tg.Done()
	w.mu.Lock()
	defer w.mu.Unlock()

	w.unconfirmedProcessedTransactions = nil
	for _, txn := range txns {
		// To save on code complexity, relevancy is determined while building
		// up the wallet transaction.
		relevant := false
		pt := modules.ProcessedTransaction{
			Transaction:           txn,
			TransactionID:         txn.ID(),
			ConfirmationHeight:    types.BlockHeight(math.MaxUint64),
			ConfirmationTimestamp: types.Timestamp(math.MaxUint64),
		}
		w.db.Update(func(tx *bolt.Tx) error {
			for _, sci := range txn.SiacoinInputs {
				_, exists := w.keys[sci.UnlockConditions.UnlockHash()]
				if exists {
					relevant = true
				}
				val, err := dbGetHistoricOutput(tx, types.OutputID(sci.ParentID))
				if err != nil {
					return err
				}
				pt.Inputs = append(pt.Inputs, modules.ProcessedInput{
					FundType:       types.SpecifierSiacoinInput,
					WalletAddress:  exists,
					RelatedAddress: sci.UnlockConditions.UnlockHash(),
					Value:          val,
				})
			}

			for i, sco := range txn.SiacoinOutputs {
				_, exists := w.keys[sco.UnlockHash]
				if exists {
					relevant = true
				}
				pt.Outputs = append(pt.Outputs, modules.ProcessedOutput{
					FundType:       types.SpecifierSiacoinOutput,
					MaturityHeight: types.BlockHeight(math.MaxUint64),
					WalletAddress:  exists,
					RelatedAddress: sco.UnlockHash,
					Value:          sco.Value,
				})
				err := dbPutHistoricOutput(tx, types.OutputID(txn.SiacoinOutputID(uint64(i))), sco.Value)
				if err != nil {
					return err
				}
			}

			return nil
		})
		for _, fee := range txn.MinerFees {
			pt.Outputs = append(pt.Outputs, modules.ProcessedOutput{
				FundType: types.SpecifierMinerFee,
				Value:    fee,
			})
		}
		if relevant {
			w.unconfirmedProcessedTransactions = append(w.unconfirmedProcessedTransactions, pt)
		}
	}
}
<commit_msg>move Update outside applyHistory<commit_after>package wallet

import (
	"math"

	"github.com/NebulousLabs/Sia/build"
	"github.com/NebulousLabs/Sia/modules"
	"github.com/NebulousLabs/Sia/types"

	"github.com/NebulousLabs/bolt"
)

// updateConfirmedSet uses a consensus change to update the confirmed set of
// outputs as understood by the wallet.
func (w *Wallet) updateConfirmedSet(cc modules.ConsensusChange) {
	for _, diff := range cc.SiacoinOutputDiffs {
		// Verify that the diff is relevant to the wallet.
		_, exists := w.keys[diff.SiacoinOutput.UnlockHash]
		if !exists {
			continue
		}

		_, exists = w.siacoinOutputs[diff.ID]
		if diff.Direction == modules.DiffApply {
			if build.DEBUG && exists {
				panic("adding an existing output to wallet")
			}
			w.siacoinOutputs[diff.ID] = diff.SiacoinOutput
		} else {
			if build.DEBUG && !exists {
				panic("deleting nonexisting output from wallet")
			}
			delete(w.siacoinOutputs, diff.ID)
		}
	}
	for _, diff := range cc.SiafundOutputDiffs {
		// Verify that the diff is relevant to the wallet.
		_, exists := w.keys[diff.SiafundOutput.UnlockHash]
		if !exists {
			continue
		}

		_, exists = w.siafundOutputs[diff.ID]
		if diff.Direction == modules.DiffApply {
			if build.DEBUG && exists {
				panic("adding an existing output to wallet")
			}
			w.siafundOutputs[diff.ID] = diff.SiafundOutput
		} else {
			if build.DEBUG && !exists {
				panic("deleting nonexisting output from wallet")
			}
			delete(w.siafundOutputs, diff.ID)
		}
	}
	for _, diff := range cc.SiafundPoolDiffs {
		if diff.Direction == modules.DiffApply {
			w.siafundPool = diff.Adjusted
		} else {
			w.siafundPool = diff.Previous
		}
	}
}

// revertHistory reverts any transaction history that was destroyed by reverted
// blocks in the consensus change.
func (w *Wallet) revertHistory(cc modules.ConsensusChange) {
	for _, block := range cc.RevertedBlocks {
		// Remove any transactions that have been reverted.
		for i := len(block.Transactions) - 1; i >= 0; i-- {
			// If the transaction is relevant to the wallet, it will be the
			// most recent transaction appended to w.processedTransactions.
			// Relevance can be determined just by looking at the last element
			// of w.processedTransactions.
			txn := block.Transactions[i]
			txid := txn.ID()
			if len(w.processedTransactions) > 0 && txid == w.processedTransactions[len(w.processedTransactions)-1].TransactionID {
				w.processedTransactions = w.processedTransactions[:len(w.processedTransactions)-1]
				delete(w.processedTransactionMap, txid)
			}
		}

		// Remove the miner payout transaction if applicable.
		for _, mp := range block.MinerPayouts {
			_, exists := w.keys[mp.UnlockHash]
			if exists {
				w.processedTransactions = w.processedTransactions[:len(w.processedTransactions)-1]
				delete(w.processedTransactionMap, types.TransactionID(block.ID()))
				break
			}
		}
		w.consensusSetHeight--
	}
}

// applyHistory applies any transaction history that was introduced by the
// applied blocks.
func (w *Wallet) applyHistory(tx *bolt.Tx, applied []types.Block) error {
	for _, block := range applied {
		w.consensusSetHeight++
		// Apply the miner payout transaction if applicable.
		minerPT := modules.ProcessedTransaction{
			Transaction:           types.Transaction{},
			TransactionID:         types.TransactionID(block.ID()),
			ConfirmationHeight:    w.consensusSetHeight,
			ConfirmationTimestamp: block.Timestamp,
		}
		relevant := false
		for i, mp := range block.MinerPayouts {
			_, exists := w.keys[mp.UnlockHash]
			if exists {
				relevant = true
			}
			minerPT.Outputs = append(minerPT.Outputs, modules.ProcessedOutput{
				FundType:       types.SpecifierMinerPayout,
				MaturityHeight: w.consensusSetHeight + types.MaturityDelay,
				WalletAddress:  exists,
				RelatedAddress: mp.UnlockHash,
				Value:          mp.Value,
			})
			err := dbPutHistoricOutput(tx, types.OutputID(block.MinerPayoutID(uint64(i))), mp.Value)
			if err != nil {
				return err
			}
		}
		if relevant {
			w.processedTransactions = append(w.processedTransactions, minerPT)
			w.processedTransactionMap[minerPT.TransactionID] = &w.processedTransactions[len(w.processedTransactions)-1]
		}
		for _, txn := range block.Transactions {
			relevant := false
			pt := modules.ProcessedTransaction{
				Transaction:           txn,
				TransactionID:         txn.ID(),
				ConfirmationHeight:    w.consensusSetHeight,
				ConfirmationTimestamp: block.Timestamp,
			}

			for _, sci := range txn.SiacoinInputs {
				_, exists := w.keys[sci.UnlockConditions.UnlockHash()]
				if exists {
					relevant = true
				}
				val, err := dbGetHistoricOutput(tx, types.OutputID(sci.ParentID))
				if err != nil {
					return err
				}
				pt.Inputs = append(pt.Inputs, modules.ProcessedInput{
					FundType:       types.SpecifierSiacoinInput,
					WalletAddress:  exists,
					RelatedAddress: sci.UnlockConditions.UnlockHash(),
					Value:          val,
				})
			}

			for i, sco := range txn.SiacoinOutputs {
				_, exists := w.keys[sco.UnlockHash]
				if exists {
					relevant = true
				}
				pt.Outputs = append(pt.Outputs, modules.ProcessedOutput{
					FundType:       types.SpecifierSiacoinOutput,
					MaturityHeight: w.consensusSetHeight,
					WalletAddress:  exists,
					RelatedAddress: sco.UnlockHash,
					Value:          sco.Value,
				})
				err := dbPutHistoricOutput(tx, types.OutputID(txn.SiacoinOutputID(uint64(i))), sco.Value)
				if err != nil {
					return err
				}
			}

			for _, sfi := range txn.SiafundInputs {
				_, exists := w.keys[sfi.UnlockConditions.UnlockHash()]
				if exists {
					relevant = true
				}
				sfiValue, err := dbGetHistoricOutput(tx, types.OutputID(sfi.ParentID))
				if err != nil {
					return err
				}
				pt.Inputs = append(pt.Inputs, modules.ProcessedInput{
					FundType:       types.SpecifierSiafundInput,
					WalletAddress:  exists,
					RelatedAddress: sfi.UnlockConditions.UnlockHash(),
					Value:          sfiValue,
				})
				startVal, err := dbGetHistoricClaimStart(tx, sfi.ParentID)
				if err != nil {
					return err
				}
				claimValue := w.siafundPool.Sub(startVal).Mul(sfiValue)
				pt.Outputs = append(pt.Outputs, modules.ProcessedOutput{
					FundType:       types.SpecifierClaimOutput,
					MaturityHeight: w.consensusSetHeight + types.MaturityDelay,
					WalletAddress:  exists,
					RelatedAddress: sfi.ClaimUnlockHash,
					Value:          claimValue,
				})
			}

			for i, sfo := range txn.SiafundOutputs {
				_, exists := w.keys[sfo.UnlockHash]
				if exists {
					relevant = true
				}
				pt.Outputs = append(pt.Outputs, modules.ProcessedOutput{
					FundType:       types.SpecifierSiafundOutput,
					MaturityHeight: w.consensusSetHeight,
					WalletAddress:  exists,
					RelatedAddress: sfo.UnlockHash,
					Value:          sfo.Value,
				})
				id := txn.SiafundOutputID(uint64(i))
				err := dbPutHistoricOutput(tx, types.OutputID(id), sfo.Value)
				if err != nil {
					return err
				}
				err = dbPutHistoricClaimStart(tx, id, sfo.ClaimStart)
				if err != nil {
					return err
				}
			}

			for _, fee := range txn.MinerFees {
				pt.Outputs = append(pt.Outputs, modules.ProcessedOutput{
					FundType: types.SpecifierMinerFee,
					Value:    fee,
				})
			}
			if relevant {
				w.processedTransactions = append(w.processedTransactions, pt)
				w.processedTransactionMap[pt.TransactionID] = &w.processedTransactions[len(w.processedTransactions)-1]
			}
		}
	}

	return nil
}

// ProcessConsensusChange parses a consensus change to update the set of
// confirmed outputs known to the wallet.
func (w *Wallet) ProcessConsensusChange(cc modules.ConsensusChange) {
	if err := w.tg.Add(); err != nil {
		// The wallet should gracefully reject updates from the consensus set
		// or transaction pool that are sent after the wallet's Close method
		// has closed the wallet's ThreadGroup.
		return
	}
	defer w.tg.Done()
	w.mu.Lock()
	defer w.mu.Unlock()
	err := w.db.Update(func(tx *bolt.Tx) error {
		w.updateConfirmedSet(cc)
		w.revertHistory(cc)
		return w.applyHistory(tx, cc.AppliedBlocks)
	})
	if err != nil {
		w.log.Println("ERROR: failed to add consensus change:", err)
	}
}

// ReceiveUpdatedUnconfirmedTransactions updates the wallet's unconfirmed
// transaction set.
func (w *Wallet) ReceiveUpdatedUnconfirmedTransactions(txns []types.Transaction, _ modules.ConsensusChange) {
	if err := w.tg.Add(); err != nil {
		// Gracefully reject transactions if the wallet's Close method has
		// closed the wallet's ThreadGroup already.
		return
	}
	defer w.tg.Done()
	w.mu.Lock()
	defer w.mu.Unlock()
	err := w.db.Update(func(tx *bolt.Tx) error {
		w.unconfirmedProcessedTransactions = nil
		for _, txn := range txns {
			// To save on code complexity, relevancy is determined while building
			// up the wallet transaction.
			relevant := false
			pt := modules.ProcessedTransaction{
				Transaction:           txn,
				TransactionID:         txn.ID(),
				ConfirmationHeight:    types.BlockHeight(math.MaxUint64),
				ConfirmationTimestamp: types.Timestamp(math.MaxUint64),
			}
			for _, sci := range txn.SiacoinInputs {
				_, exists := w.keys[sci.UnlockConditions.UnlockHash()]
				if exists {
					relevant = true
				}
				val, err := dbGetHistoricOutput(tx, types.OutputID(sci.ParentID))
				if err != nil {
					return err
				}
				pt.Inputs = append(pt.Inputs, modules.ProcessedInput{
					FundType:       types.SpecifierSiacoinInput,
					WalletAddress:  exists,
					RelatedAddress: sci.UnlockConditions.UnlockHash(),
					Value:          val,
				})
			}
			for i, sco := range txn.SiacoinOutputs {
				_, exists := w.keys[sco.UnlockHash]
				if exists {
					relevant = true
				}
				pt.Outputs = append(pt.Outputs, modules.ProcessedOutput{
					FundType:       types.SpecifierSiacoinOutput,
					MaturityHeight: types.BlockHeight(math.MaxUint64),
					WalletAddress:  exists,
					RelatedAddress: sco.UnlockHash,
					Value:          sco.Value,
				})
				err := dbPutHistoricOutput(tx, types.OutputID(txn.SiacoinOutputID(uint64(i))), sco.Value)
				if err != nil {
					return err
				}
			}
			for _, fee := range txn.MinerFees {
				pt.Outputs = append(pt.Outputs, modules.ProcessedOutput{
					FundType: types.SpecifierMinerFee,
					Value:    fee,
				})
			}
			if relevant {
				w.unconfirmedProcessedTransactions = append(w.unconfirmedProcessedTransactions, pt)
			}
		}
		return nil
	})
	if err != nil {
		w.log.Println("ERROR: failed to add unconfirmed transactions:", err)
	}
}
<|endoftext|>


<commit_before>package main_test

import (
	"io/ioutil"
	"net/url"

	. "github.com/alphagov/govuk_crawler_worker"
	. "github.com/alphagov/govuk_crawler_worker/http_crawler"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"

	"os"

	"github.com/streadway/amqp"
)

var _ = Describe("CrawlerMessageItem", func() {
	var mirrorRoot, testUrl, urlPath string
	var delivery amqp.Delivery
	var err error
	var html []byte
	var item *CrawlerMessageItem
	var rootURL *url.URL

	BeforeEach(func() {
		mirrorRoot = os.Getenv("MIRROR_ROOT")
		if mirrorRoot == "" {
			mirrorRoot, err = ioutil.TempDir("", "crawler_message_item_test")
			Expect(err).To(BeNil())
		}

		rootURL, _ = url.Parse("https://www.gov.uk")
		testUrl = rootURL.String() + urlPath
		urlPath = "/government/organisations"

		delivery = amqp.Delivery{Body: []byte(testUrl)}
		item = NewCrawlerMessageItem(delivery, rootURL, []string{})

		html = []byte(`
<title>test</title

TEST

`)
		item.Response = &CrawlerResponse{Body: html}
	})

	AfterEach(func() {
		DeleteMirrorFilesFromDisk(mirrorRoot)
	})

	It("generates a CrawlerMessageItem object", func() {
		Expect(NewCrawlerMessageItem(delivery, rootURL, []string{})).ToNot(BeNil())
	})

	Describe("getting and setting the Response.Body", func() {
		It("can get the Response.Body of the crawled URL", func() {
			Expect(item.Response.Body).To(Equal(html))
		})

		It("can set the Response.Body of the crawled URL", func() {
			item := NewCrawlerMessageItem(delivery, rootURL, []string{})
			item.Response = &CrawlerResponse{Body: []byte("foo")}

			Expect(item.Response.Body).To(Equal([]byte("foo")))
		})
	})

	It("detects when a URL is blacklisted", func() {
		delivery = amqp.Delivery{Body: []byte("https://www.example.com/blacklisted")}
		item := NewCrawlerMessageItem(delivery, rootURL, []string{"/blacklisted"})
		Expect(item.IsBlacklisted()).To(BeTrue())
	})

	It("returns its URL", func() {
		Expect(item.URL()).To(Equal(testUrl))
	})

	Describe("generating a sane filename", func() {
		It("strips out the domain, protocol, auth and ports", func() {
			testUrl = "https://user:pass@example.com:8080/test/url"
			delivery = amqp.Delivery{Body: []byte(testUrl)}

			item = NewCrawlerMessageItem(delivery, rootURL, []string{})
			item.Response = &CrawlerResponse{Body: html, ContentType: HTML}

			Expect(item.RelativeFilePath()).To(Equal("test/url.html"))
		})

		It("strips preceeding path traversals and resolves the remaining path", func() {
			testUrl = rootURL.String() + "/../../one/./two/../three"
			delivery = amqp.Delivery{Body: []byte(testUrl)}

			item = NewCrawlerMessageItem(delivery, rootURL, []string{})
			item.Response = &CrawlerResponse{Body: html, ContentType: HTML}

			Expect(item.RelativeFilePath()).To(Equal("one/three.html"))
		})

		It("preserves case sensitivity", func() {
			testUrl = rootURL.String() + "/test/UPPER/MiXeD"
			delivery = amqp.Delivery{Body: []byte(testUrl)}

			item = NewCrawlerMessageItem(delivery, rootURL, []string{})
			item.Response = &CrawlerResponse{Body: html, ContentType: HTML}

			Expect(item.RelativeFilePath()).To(Equal("test/UPPER/MiXeD.html"))
		})

		It("preserves non-alphanumeric characters", func() {
			testUrl = rootURL.String() + "/test/!T@e£s$t/U^R*L(){}"
			delivery = amqp.Delivery{Body: []byte(testUrl)}

			item = NewCrawlerMessageItem(delivery, rootURL, []string{})
			item.Response = &CrawlerResponse{Body: html, ContentType: HTML}

			Expect(item.RelativeFilePath()).To(Equal("test/!T@e£s$t/U^R*L(){}.html"))
		})

		It("preserves multiple dashes", func() {
			testUrl = rootURL.String() + "/test/one-two--three---"
			delivery = amqp.Delivery{Body: []byte(testUrl)}

			item = NewCrawlerMessageItem(delivery, rootURL, []string{})
			item.Response = &CrawlerResponse{Body: html, ContentType: HTML}

			Expect(item.RelativeFilePath()).To(Equal("test/one-two--three---.html"))
		})

		It("preserves non-latin chars and not URL encode them", func() {
			testUrl = rootURL.String() + `/test/如何在香港申請英國簽證`
			delivery = amqp.Delivery{Body: []byte(testUrl)}

			item = NewCrawlerMessageItem(delivery, rootURL, []string{})
			item.Response = &CrawlerResponse{Body: html, ContentType: HTML}

			Expect(item.RelativeFilePath()).To(Equal(`test/如何在香港申請英國簽證.html`))
		})

		It("adds an index.html suffix when URL references a directory", func() {
			testUrl = rootURL.String() + "/this/url/has/a/trailing/slash/"
			delivery = amqp.Delivery{Body: []byte(testUrl)}

			item = NewCrawlerMessageItem(delivery, rootURL, []string{})
			item.Response = &CrawlerResponse{Body: html, ContentType: HTML}

			Expect(item.RelativeFilePath()).To(Equal("this/url/has/a/trailing/slash/index.html"))
		})

		It("adds an index.html suffix when URL has no path and no trailing slash", func() {
			testUrl = rootURL.String() + "/"
			delivery = amqp.Delivery{Body: []byte(testUrl)}

			item = NewCrawlerMessageItem(delivery, rootURL, []string{})
			item.Response = &CrawlerResponse{Body: html, ContentType: HTML}

			Expect(item.RelativeFilePath()).To(Equal("index.html"))
		})

		It("omits URL query parameters", func() {
			delivery := amqp.Delivery{Body: []byte(testUrl + "?foo=bar")}
			item = NewCrawlerMessageItem(delivery, rootURL, []string{})

			item.Response = &CrawlerResponse{Body: html, ContentType: HTML}

			Expect(item.RelativeFilePath()).To(Equal("government/organisations.html"))
		})

		It("omits URL fragments", func() {
			delivery := amqp.Delivery{Body: []byte(testUrl + "#foo")}

			item = NewCrawlerMessageItem(delivery, rootURL, []string{})
			item.Response = &CrawlerResponse{Body: html, ContentType: HTML}

			Expect(item.RelativeFilePath()).To(Equal("government/organisations.html"))
		})

		It("supports ATOM URLs", func() {
			testUrl = rootURL.String() + "/things.atom"
			delivery = amqp.Delivery{Body: []byte(testUrl)}

			item = NewCrawlerMessageItem(delivery, rootURL, []string{})
			item.Response = &CrawlerResponse{Body: []byte(""), ContentType: ATOM}

			Expect(item.RelativeFilePath()).To(Equal("things.atom"))
		})

		It("supports JSON URLs", func() {
			testUrl = rootURL.String() + "/api.json"
			delivery = amqp.Delivery{Body: []byte(testUrl)}

			item = NewCrawlerMessageItem(delivery, rootURL, []string{})
			item.Response = &CrawlerResponse{Body: []byte(""), ContentType: JSON}

			Expect(item.RelativeFilePath()).To(Equal("api.json"))
		})
	})

	Describe("ExtractURLs", func() {
		It("should return an empty array if it can't find any matching URLs", func() {
			item.Response.Body = []byte("")
			urls, err := item.ExtractURLs()

			Expect(err).To(BeNil())
			Expect(urls).To(Equal([]*url.URL{}))
		})

		It("should extract all a[@href] URLs from a given HTML document", func() {
			item.Response.Body = []byte(`<div><a href="https://www.gov.uk/"></a></div>`)
			urls, err := item.ExtractURLs()
			expectedURL, _ := url.Parse("https://www.gov.uk/")

			Expect(err).To(BeNil())
			Expect(urls).To(ContainElement(expectedURL))
		})

		It("should extract all img[@src] URLs from a given HTML document", func() {
			item.Response.Body = []byte(`<div><img src="https://www.gov.uk/image.png" /></div>`)
			urls, err := item.ExtractURLs()
			expectedURL, _ := url.Parse("https://www.gov.uk/image.png")

			Expect(err).To(BeNil())
			Expect(urls).To(ContainElement(expectedURL))
		})

		It("should extract all link[@href] URLs from a given HTML document", func() {
			item.Response.Body = []byte(`<head><link rel="icon" href="https://www.gov.uk/favicon.ico"></head>`)
			expectedURL, _ := url.Parse("https://www.gov.uk/favicon.ico")
			urls, err := item.ExtractURLs()

			Expect(err).To(BeNil())
			Expect(urls).To(ContainElement(expectedURL))
		})

		It("should extract all script[@src] URLs from a given HTML document", func() {
			item.Response.Body = []byte(
				`<head><script type="text/javascript" src="https://www.gov.uk/jq.js"></script></head>`)
			urls, err := item.ExtractURLs()
			expectedURL, _ := url.Parse("https://www.gov.uk/jq.js")

			Expect(err).To(BeNil())
			Expect(urls).To(ContainElement(expectedURL))
		})

		It("successfully extracts multiple matching URLs from the provided DOM", func() {
			item.Response.Body = []byte(
				`

<link rel="icon" href="https://www.gov.uk/favicon.ico">
`)
			urls, err := item.ExtractURLs()
			expectedURL1, _ := url.Parse("https://www.gov.uk/jq.js")
			expectedURL2, _ := url.Parse("https://www.gov.uk/favicon.ico")

			Expect(err).To(BeNil())
			Expect(urls).To(ContainElement(expectedURL1))
			Expect(urls).To(ContainElement(expectedURL2))
		})

		It("will not provide URLs that don't match the provided prefix rootURL", func() {
			item.Response.Body = []byte(
				`<head><script type="text/javascript" src="https://www.foobar.com/jq.js"></script></head>`)
			urls, err := item.ExtractURLs()

			Expect(err).To(BeNil())
			Expect(urls).To(BeEmpty())
		})

		It("will unescape URLs", func() {
			item.Response.Body = []byte(`<div><a href="http://www.gov.uk/bar%20"></a></div>`)
			expectedURL, _ := url.Parse("http://www.gov.uk/bar")
			urls, err := item.ExtractURLs()

			Expect(err).To(BeNil())
			Expect(urls).To(ContainElement(expectedURL))
		})

		It("should extract relative URLs", func() {
			item.Response.Body = []byte(`<div><a href="/foo/bar">a</a><a href="mailto:c@d.com">b</a></div>`)
			expectedURL, _ := url.Parse("https://www.gov.uk/foo/bar")
			urls, err := item.ExtractURLs()

			Expect(err).To(BeNil())
			Expect(len(urls)).To(Equal(1))
			Expect(urls).To(ContainElement(expectedURL))
		})

		It("should remove the #fragment when extracting URLs", func() {
			item.Response.Body = []byte(`<div><a href="http://www.gov.uk/#germany"></a></div>`)
			expectedURL, _ := url.Parse("http://www.gov.uk/")
			urls, err := item.ExtractURLs()

			Expect(err).To(BeNil())
			Expect(urls).To(ContainElement(expectedURL))
		})

		It("removes paths that are blacklisted", func() {
			item := NewCrawlerMessageItem(delivery, rootURL, []string{"/trade-tariff"})
			item.Response = &CrawlerResponse{
				Body: []byte(`<div><a href="/foo/bar">a</a><a href="/trade-tariff">b</a></div>`),
			}

			urls, err := item.ExtractURLs()

			Expect(err).To(BeNil())
			Expect(len(urls)).To(Equal(1))
		})

		It("should only return unique URLs", func() {
			item.Response.Body = []byte(`<a href="http://www.gov.uk/foo">a</a><a href="http://www.gov.uk/foo">b</a>`)
			urls, err := item.ExtractURLs()

			Expect(err).To(BeNil())
			Expect(urls).To(HaveLen(1))
		})
	})
})
<commit_msg>Rename testUrl variable to testURL<commit_after>package main_test

import (
	"io/ioutil"
	"net/url"

	. "github.com/alphagov/govuk_crawler_worker"
	. "github.com/alphagov/govuk_crawler_worker/http_crawler"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"

	"os"

	"github.com/streadway/amqp"
)

var _ = Describe("CrawlerMessageItem", func() {
	var mirrorRoot, testURL, urlPath string
	var delivery amqp.Delivery
	var err error
	var html []byte
	var item *CrawlerMessageItem
	var rootURL *url.URL

	BeforeEach(func() {
		mirrorRoot = os.Getenv("MIRROR_ROOT")
		if mirrorRoot == "" {
			mirrorRoot, err = ioutil.TempDir("", "crawler_message_item_test")
			Expect(err).To(BeNil())
		}

		rootURL, _ = url.Parse("https://www.gov.uk")
		testURL = rootURL.String() + urlPath
		urlPath = "/government/organisations"

		delivery = amqp.Delivery{Body: []byte(testURL)}
		item = NewCrawlerMessageItem(delivery, rootURL, []string{})

		html = []byte(`
<title>test</title

TEST

`)
		item.Response = &CrawlerResponse{Body: html}
	})

	AfterEach(func() {
		DeleteMirrorFilesFromDisk(mirrorRoot)
	})

	It("generates a CrawlerMessageItem object", func() {
		Expect(NewCrawlerMessageItem(delivery, rootURL, []string{})).ToNot(BeNil())
	})

	Describe("getting and setting the Response.Body", func() {
		It("can get the Response.Body of the crawled URL", func() {
			Expect(item.Response.Body).To(Equal(html))
		})

		It("can set the Response.Body of the crawled URL", func() {
			item := NewCrawlerMessageItem(delivery, rootURL, []string{})
			item.Response = &CrawlerResponse{Body: []byte("foo")}

			Expect(item.Response.Body).To(Equal([]byte("foo")))
		})
	})

	It("detects when a URL is blacklisted", func() {
		delivery = amqp.Delivery{Body: []byte("https://www.example.com/blacklisted")}
		item := NewCrawlerMessageItem(delivery, rootURL, []string{"/blacklisted"})
		Expect(item.IsBlacklisted()).To(BeTrue())
	})

	It("returns its URL", func() {
		Expect(item.URL()).To(Equal(testURL))
	})

	Describe("generating a sane filename", func() {
		It("strips out the domain, protocol, auth and ports", func() {
			testURL = "https://user:pass@example.com:8080/test/url"
			delivery = amqp.Delivery{Body: []byte(testURL)}

			item = NewCrawlerMessageItem(delivery, rootURL, []string{})
			item.Response = &CrawlerResponse{Body: html, ContentType: HTML}

			Expect(item.RelativeFilePath()).To(Equal("test/url.html"))
		})

		It("strips preceeding path traversals and resolves the remaining path", func() {
			testURL = rootURL.String() + "/../../one/./two/../three"
			delivery = amqp.Delivery{Body: []byte(testURL)}

			item = NewCrawlerMessageItem(delivery, rootURL, []string{})
			item.Response = &CrawlerResponse{Body: html, ContentType: HTML}

			Expect(item.RelativeFilePath()).To(Equal("one/three.html"))
		})

		It("preserves case sensitivity", func() {
			testURL = rootURL.String() + "/test/UPPER/MiXeD"
			delivery = amqp.Delivery{Body: []byte(testURL)}

			item = NewCrawlerMessageItem(delivery, rootURL, []string{})
			item.Response = &CrawlerResponse{Body: html, ContentType: HTML}

			Expect(item.RelativeFilePath()).To(Equal("test/UPPER/MiXeD.html"))
		})

		It("preserves non-alphanumeric characters", func() {
			testURL = rootURL.String() + "/test/!T@e£s$t/U^R*L(){}"
			delivery = amqp.Delivery{Body: []byte(testURL)}

			item = NewCrawlerMessageItem(delivery, rootURL, []string{})
			item.Response = &CrawlerResponse{Body: html, ContentType: HTML}

			Expect(item.RelativeFilePath()).To(Equal("test/!T@e£s$t/U^R*L(){}.html"))
		})

		It("preserves multiple dashes", func() {
			testURL = rootURL.String() + "/test/one-two--three---"
			delivery = amqp.Delivery{Body: []byte(testURL)}

			item = NewCrawlerMessageItem(delivery, rootURL, []string{})
			item.Response = &CrawlerResponse{Body: html, ContentType: HTML}

			Expect(item.RelativeFilePath()).To(Equal("test/one-two--three---.html"))
		})

		It("preserves non-latin chars and not URL encode them", func() {
			testURL = rootURL.String() + `/test/如何在香港申請英國簽證`
			delivery = amqp.Delivery{Body: []byte(testURL)}

			item = NewCrawlerMessageItem(delivery, rootURL, []string{})
			item.Response = &CrawlerResponse{Body: html, ContentType: HTML}

			Expect(item.RelativeFilePath()).To(Equal(`test/如何在香港申請英國簽證.html`))
		})

		It("adds an index.html suffix when URL references a directory", func() {
			testURL = rootURL.String() + "/this/url/has/a/trailing/slash/"
			delivery = amqp.Delivery{Body: []byte(testURL)}

			item = NewCrawlerMessageItem(delivery, rootURL, []string{})
			item.Response = &CrawlerResponse{Body: html, ContentType: HTML}

			Expect(item.RelativeFilePath()).To(Equal("this/url/has/a/trailing/slash/index.html"))
		})

		It("adds an index.html suffix when URL has no path and no trailing slash", func() {
			testURL = rootURL.String() + "/"
			delivery = amqp.Delivery{Body: []byte(testURL)}

			item = NewCrawlerMessageItem(delivery, rootURL, []string{})
			item.Response = &CrawlerResponse{Body: html, ContentType: HTML}

			Expect(item.RelativeFilePath()).To(Equal("index.html"))
		})

		It("omits URL query parameters", func() {
			delivery := amqp.Delivery{Body: []byte(testURL + "?foo=bar")}
			item = NewCrawlerMessageItem(delivery, rootURL, []string{})

			item.Response = &CrawlerResponse{Body: html, ContentType: HTML}

			Expect(item.RelativeFilePath()).To(Equal("government/organisations.html"))
		})

		It("omits URL fragments", func() {
			delivery := amqp.Delivery{Body: []byte(testURL + "#foo")}

			item = NewCrawlerMessageItem(delivery, rootURL, []string{})
			item.Response = &CrawlerResponse{Body: html, ContentType: HTML}

			Expect(item.RelativeFilePath()).To(Equal("government/organisations.html"))
		})

		It("supports ATOM URLs", func() {
			testURL = rootURL.String() + "/things.atom"
			delivery = amqp.Delivery{Body: []byte(testURL)}

			item = NewCrawlerMessageItem(delivery, rootURL, []string{})
			item.Response = &CrawlerResponse{Body: []byte(""), ContentType: ATOM}

			Expect(item.RelativeFilePath()).To(Equal("things.atom"))
		})

		It("supports JSON URLs", func() {
			testURL = rootURL.String() + "/api.json"
			delivery = amqp.Delivery{Body: []byte(testURL)}

			item = NewCrawlerMessageItem(delivery, rootURL, []string{})
			item.Response = &CrawlerResponse{Body: []byte(""), ContentType: JSON}

			Expect(item.RelativeFilePath()).To(Equal("api.json"))
		})
	})

	Describe("ExtractURLs", func() {
		It("should return an empty array if it can't find any matching URLs", func() {
			item.Response.Body = []byte("")
			urls, err := item.ExtractURLs()

			Expect(err).To(BeNil())
			Expect(urls).To(Equal([]*url.URL{}))
		})

		It("should extract all a[@href] URLs from a given HTML document", func() {
			item.Response.Body = []byte(`<div><a href="https://www.gov.uk/"></a></div>`)
			urls, err := item.ExtractURLs()
			expectedURL, _ := url.Parse("https://www.gov.uk/")

			Expect(err).To(BeNil())
			Expect(urls).To(ContainElement(expectedURL))
		})

		It("should extract all img[@src] URLs from a given HTML document", func() {
			item.Response.Body = []byte(`<div><img src="https://www.gov.uk/image.png" /></div>`)
			urls, err := item.ExtractURLs()
			expectedURL, _ := url.Parse("https://www.gov.uk/image.png")

			Expect(err).To(BeNil())
			Expect(urls).To(ContainElement(expectedURL))
		})

		It("should extract all link[@href] URLs from a given HTML document", func() {
			item.Response.Body = []byte(`<head><link rel="icon" href="https://www.gov.uk/favicon.ico"></head>`)
			expectedURL, _ := url.Parse("https://www.gov.uk/favicon.ico")
			urls, err := item.ExtractURLs()

			Expect(err).To(BeNil())
			Expect(urls).To(ContainElement(expectedURL))
		})

		It("should extract all script[@src] URLs from a given HTML document", func() {
			item.Response.Body = []byte(
				`<head><script type="text/javascript" src="https://www.gov.uk/jq.js"></script></head>`)
			urls, err := item.ExtractURLs()
			expectedURL, _ := url.Parse("https://www.gov.uk/jq.js")

			Expect(err).To(BeNil())
			Expect(urls).To(ContainElement(expectedURL))
		})

		It("successfully extracts multiple matching URLs from the provided DOM", func() {
			item.Response.Body = []byte(
				`

<link rel="icon" href="https://www.gov.uk/favicon.ico">
`)
			urls, err := item.ExtractURLs()
			expectedURL1, _ := url.Parse("https://www.gov.uk/jq.js")
			expectedURL2, _ := url.Parse("https://www.gov.uk/favicon.ico")

			Expect(err).To(BeNil())
			Expect(urls).To(ContainElement(expectedURL1))
			Expect(urls).To(ContainElement(expectedURL2))
		})

		It("will not provide URLs that don't match the provided prefix rootURL", func() {
			item.Response.Body = []byte(
				`<head><script type="text/javascript" src="https://www.foobar.com/jq.js"></script></head>`)
			urls, err := item.ExtractURLs()

			Expect(err).To(BeNil())
			Expect(urls).To(BeEmpty())
		})

		It("will unescape URLs", func() {
			item.Response.Body = []byte(`<div><a href="http://www.gov.uk/bar%20"></a></div>`)
			expectedURL, _ := url.Parse("http://www.gov.uk/bar")
			urls, err := item.ExtractURLs()

			Expect(err).To(BeNil())
			Expect(urls).To(ContainElement(expectedURL))
		})

		It("should extract relative URLs", func() {
			item.Response.Body = []byte(`<div><a href="/foo/bar">a</a><a href="mailto:c@d.com">b</a></div>`)
			expectedURL, _ := url.Parse("https://www.gov.uk/foo/bar")
			urls, err := item.ExtractURLs()

			Expect(err).To(BeNil())
			Expect(len(urls)).To(Equal(1))
			Expect(urls).To(ContainElement(expectedURL))
		})

		It("should remove the #fragment when extracting URLs", func() {
			item.Response.Body = []byte(`<div><a href="http://www.gov.uk/#germany"></a></div>`)
			expectedURL, _ := url.Parse("http://www.gov.uk/")
			urls, err := item.ExtractURLs()

			Expect(err).To(BeNil())
			Expect(urls).To(ContainElement(expectedURL))
		})

		It("removes paths that are blacklisted", func() {
			item := NewCrawlerMessageItem(delivery, rootURL, []string{"/trade-tariff"})
			item.Response = &CrawlerResponse{
				Body: []byte(`<div><a href="/foo/bar">a</a><a href="/trade-tariff">b</a></div>`),
			}

			urls, err := item.ExtractURLs()

			Expect(err).To(BeNil())
			Expect(len(urls)).To(Equal(1))
		})

		It("should only return unique URLs", func() {
			item.Response.Body = []byte(`<a href="http://www.gov.uk/foo">a</a><a href="http://www.gov.uk/foo">b</a>`)
			urls, err := item.ExtractURLs()

			Expect(err).To(BeNil())
			Expect(urls).To(HaveLen(1))
		})
	})
})
<|endoftext|>


<commit_before>// +build !nopkcs11

package pkcs11key

import (
	"bytes"
	"crypto"
	"crypto/rand"
	"crypto/rsa"
	"crypto/ecdsa"
	"github.com/miekg/pkcs11"
	"testing"
)

type mockCtx struct{}

const rsaPrivateKeyHandle = pkcs11.ObjectHandle(23)
const ecPrivateKeyHandle = pkcs11.ObjectHandle(32)
const ecPublicKeyHandle = pkcs11.ObjectHandle(33)
// EC private key with no matching public key
const ecPrivNoPubHandle = pkcs11.ObjectHandle(34)
// EC private and public key with invalid EC point
const ecInvEcPointPrivHandle = pkcs11.ObjectHandle(35)
const ecInvEcPointPubHandle = pkcs11.ObjectHandle(36)
const sessionHandle = pkcs11.SessionHandle(17)

var slots = []uint{7, 8, 9}
var tokenInfo = pkcs11.TokenInfo{
	Label: "token label",
}
var currentSearch []*pkcs11.Attribute

func (c mockCtx) CloseSession(sh pkcs11.SessionHandle) error {
	return nil
}

func (c mockCtx) FindObjectsFinal(sh pkcs11.SessionHandle) error {
	return nil
}

func (c mockCtx) FindObjectsInit(sh pkcs11.SessionHandle, temp []*pkcs11.Attribute) error {
	currentSearch = temp
	return nil
}

func (c mockCtx) FindObjects(sh pkcs11.SessionHandle, max int) ([]pkcs11.ObjectHandle, bool, error) {
	for _, a := range currentSearch {
		// We search private keys using CKA_LABEL
		if a.Type == pkcs11.CKA_LABEL {
			if a.Value[0] == 0x72 { // Label start with 'r'
				return []pkcs11.ObjectHandle{rsaPrivateKeyHandle}, true, nil
			} else if a.Value[0] == 0x65 { // Label start with 'e'
				return []pkcs11.ObjectHandle{ecPrivateKeyHandle}, true, nil
			} else if a.Value[0] == 0x6E { // Label start with 'n'
				return []pkcs11.ObjectHandle{ecPrivNoPubHandle}, true, nil
			} else if a.Value[0] == 0x69 { // Label start with 'i'
				return []pkcs11.ObjectHandle{ecInvEcPointPrivHandle}, true, nil
			}
		}
		// We search th EC public key using CKA_ID
		if a.Type == pkcs11.CKA_ID {
			if a.Value[0] == 0x3 {
				return []pkcs11.ObjectHandle{ecPublicKeyHandle}, true, nil
			} else if a.Value[0] == 0x5 {
				return []pkcs11.ObjectHandle{ecInvEcPointPubHandle}, true, nil
			} else {
				return []pkcs11.ObjectHandle{}, true, nil
			}
		}
	}
	return nil, false, nil
}

func rsaPrivateAttributes(template []*pkcs11.Attribute) ([]*pkcs11.Attribute, error) {
	var output []*pkcs11.Attribute
	for _, a := range template {
		// Return simple values for these attributes. Note that a value of `1` for
		// both modulus an public exponent would be a very bad public key, but it's
		// sufficient to satisfy the current code.
		if a.Type == pkcs11.CKA_MODULUS ||
			a.Type == pkcs11.CKA_PUBLIC_EXPONENT ||
			a.Type == pkcs11.CKA_ALWAYS_AUTHENTICATE {
			output = append(output, &pkcs11.Attribute{
				Type:  a.Type,
				Value: []byte{byte(1)},
			})
		}
		if a.Type == pkcs11.CKA_KEY_TYPE {
			output = append(output, &pkcs11.Attribute{
				Type:  a.Type,
				Value: []byte{byte(0)},
			})
		}
	}
	return output, nil
}

var ecOid = []byte{0x06, 0x08, 0x2A, 0x86, 0x48, 0xCE, 0x3D, 0x03, 0x01, 0x07}

func ecPrivateAttributes(template []*pkcs11.Attribute) ([]*pkcs11.Attribute, error) {
	var output []*pkcs11.Attribute
	for _, a := range template {
		if a.Type == pkcs11.CKA_EC_PARAMS {
			output = append(output, &pkcs11.Attribute{
				Type:  a.Type,
				Value: ecOid,
			})
		}
		if a.Type == pkcs11.CKA_KEY_TYPE ||
			a.Type == pkcs11.CKA_ID {
			output = append(output, &pkcs11.Attribute{
				Type:  a.Type,
				Value: []byte{byte(3)},
			})
		}
	}
	return output, nil
}

var ecPoint = []byte{0x04, 0x41, 0x04, 0x4C, 0xD7, 0x7B, 0x7B, 0x2E,
	0x3D, 0x57, 0x98, 0xB8, 0x2F, 0x99, 0xB4, 0x83,
	0x99, 0xE6, 0xD4, 0x4C, 0x4F, 0xBC, 0x2D, 0x60,
	0xCD, 0x08, 0x8E, 0x93, 0x65, 0x6F, 0x20, 0x51,
	0x1C, 0xE7, 0xFD, 0x59, 0x34, 0xAA, 0xA9, 0x36,
	0x26, 0xCE, 0x4A, 0xC5, 0xA2, 0x4A, 0x85, 0x6C,
	0xB3, 0x95, 0xFF, 0x92, 0x0F, 0x56, 0x76, 0x34,
	0x1F, 0x69, 0x52, 0x5F, 0x20, 0x83, 0x13, 0x50,
	0xA3, 0xDE, 0xBE}

func ecPublicAttributes(template []*pkcs11.Attribute) ([]*pkcs11.Attribute, error) {
	var output []*pkcs11.Attribute
	for _, a := range template {
		if a.Type == pkcs11.CKA_EC_PARAMS {
			output = append(output, &pkcs11.Attribute{
				Type:  a.Type,
				Value: ecOid,
			})
		}
		if a.Type == pkcs11.CKA_EC_POINT {
			output = append(output, &pkcs11.Attribute{
				Type:  a.Type,
				Value: ecPoint,
			})
		}
		if a.Type == pkcs11.CKA_KEY_TYPE ||
			a.Type == pkcs11.CKA_ID {
			output = append(output, &pkcs11.Attribute{
				Type:  a.Type,
				Value: []byte{byte(3)},
			})
		}
	}
	return output, nil
}

func ecPrivNoPubAttributes(template []*pkcs11.Attribute) ([]*pkcs11.Attribute, error) {
	var output []*pkcs11.Attribute
	for _, a := range template {
		if a.Type == pkcs11.CKA_EC_PARAMS {
			output = append(output, &pkcs11.Attribute{
				Type:  a.Type,
				Value: ecOid,
			})
		}
		if a.Type == pkcs11.CKA_KEY_TYPE {
			output = append(output, &pkcs11.Attribute{
				Type:  a.Type,
				Value: []byte{byte(3)},
			})
		}
		if a.Type == pkcs11.CKA_ID {
			output = append(output, &pkcs11.Attribute{
				Type:  a.Type,
				Value: []byte{byte(4)},
			})
		}
	}
	return output, nil
}

func ecInvalidEcPointAttributes(template []*pkcs11.Attribute) ([]*pkcs11.Attribute, error) {
	var output []*pkcs11.Attribute
	for _, a := range template {
		if a.Type == pkcs11.CKA_EC_PARAMS {
			output = append(output, &pkcs11.Attribute{
				Type:  a.Type,
				Value: ecOid,
			})
		}
		if a.Type == pkcs11.CKA_EC_POINT {
			output = append(output, &pkcs11.Attribute{
				Type:  a.Type,
				Value: []byte{byte(0)},
			})
		}
		if a.Type == pkcs11.CKA_KEY_TYPE {
			output = append(output, &pkcs11.Attribute{
				Type:  a.Type,
				Value: []byte{byte(3)},
			})
		}
		if a.Type == pkcs11.CKA_ID {
			output = append(output, &pkcs11.Attribute{
				Type:  a.Type,
				Value: []byte{byte(5)},
			})
		}
	}
	return output, nil
}

func (c mockCtx) GetAttributeValue(sh pkcs11.SessionHandle, o pkcs11.ObjectHandle, template []*pkcs11.Attribute) ([]*pkcs11.Attribute, error) {
	if o == rsaPrivateKeyHandle {
		return rsaPrivateAttributes(template)
	} else if o == ecPrivateKeyHandle {
		return ecPrivateAttributes(template)
	} else if o == ecPublicKeyHandle {
		return ecPublicAttributes(template)
	} else if o == ecPrivNoPubHandle {
		return ecPrivNoPubAttributes(template)
	} else if o == ecInvEcPointPrivHandle {
		return ecInvalidEcPointAttributes(template)
	} else if o == ecInvEcPointPubHandle {
		return ecInvalidEcPointAttributes(template)
	}
	return nil, nil
}

func (c mockCtx) GetSlotList(tokenPresent bool) ([]uint, error) {
	return slots, nil
}

func (c mockCtx) GetTokenInfo(slotID uint) (pkcs11.TokenInfo, error) {
	return tokenInfo, nil
}

func (c mockCtx) Initialize() error {
	return nil
}

func (c mockCtx) Login(sh pkcs11.SessionHandle, userType uint, pin string) error {
	return nil
}

func (c mockCtx) Logout(sh pkcs11.SessionHandle) error {
	return nil
}

func (c mockCtx) OpenSession(slotID uint, flags uint) (pkcs11.SessionHandle, error) {
	return sessionHandle, nil
}

func (c mockCtx) SignInit(sh pkcs11.SessionHandle, m []*pkcs11.Mechanism, o pkcs11.ObjectHandle) error {
	return nil
}

func (c mockCtx) Sign(sh pkcs11.SessionHandle, message []byte) ([]byte, error) {
	return message, nil
}

func setup(t *testing.T, label string) *Key {
	ps := Key{
		module:     mockCtx{},
		tokenLabel: "token label",
		pin:        "unused",
	}
	err := ps.setup(label)
	if err != nil {
		t.Fatalf("Failed to set up Key: %s", err)
	}
	return &ps
}

var signInput = []byte("1234567890 1234567890 1234567890")
func sign(t *testing.T, ps *Key) ([]byte) {
	// Sign input must be exactly 32 bytes to match SHA256 size. In normally
	// usage, Sign would be called by e.g. x509.CreateCertificate, which would
	// handle padding to the necessary size.
	output, err := ps.Sign(rand.Reader, signInput, crypto.SHA256)
	if err != nil {
		t.Fatalf("Failed to sign: %s", err)
	}

	if len(output) < len(signInput) {
		t.Fatalf("Invalid signature size")
	}

	i := len(output) - len(signInput)
	if !bytes.Equal(output[i:], signInput) {
		t.Fatal("Incorrect sign output")
	}
	return output
}

func TestSign(t *testing.T) {
	ps := setup(t, "rsa")
	sig := sign(t, ps)

	// Check that the RSA signature starts with the SHA256 hash prefix
	var sha256_pre = []byte{0x30, 0x31, 0x30, 0x0d, 0x06, 0x09, 0x60, 0x86,
		0x48, 0x01, 0x65, 0x03, 0x04, 0x02, 0x01, 0x05, 0x00, 0x04, 0x20}
	if !(bytes.Equal(sha256_pre, sig[0:19])) {
		t.Fatal("RSA signature doesn't start with prefix")
	}

	pub := ps.Public()
	// Check public key is of right type
	_ = pub.(*rsa.PublicKey)

	ps = setup(t, "ec")
	sig = sign(t, ps)

	if !(bytes.Equal(signInput, sig)) {
		t.Fatal("ECDSA signature error")
	}

	pub = ps.Public()
	// Check public key is of right type
	ecPub := pub.(*ecdsa.PublicKey)
	if !(bytes.Equal(ecPub.X.Bytes(), ecPoint[3:35]) &&
		bytes.Equal(ecPub.Y.Bytes(), ecPoint[35:])) {
		t.Fatal("Incorrect decoding of EC Point")
	}

	k := Key{
		module:     mockCtx{},
		tokenLabel: "token label",
		pin:        "unused",
	}

	// Trying to load private EC key with no public key
	err := k.setup("no_pub_ec")
	if err == nil {
		t.Fatalf("Unexpected succes")
	}

	// Trying to load private EC key with invalid EC point
	err = k.setup("invalid_ec_point")
	if err == nil {
		t.Fatalf("Unexpected succes")
	}
}

// This is a version of the mock that gives CKR_ATTRIBUTE_TYPE_INVALID when
// asked about the CKA_ALWAYS_AUTHENTICATE attribute.
type mockCtxFailsAlwaysAuthenticate struct {
	mockCtx
}

func (c mockCtxFailsAlwaysAuthenticate) GetAttributeValue(sh pkcs11.SessionHandle, o pkcs11.ObjectHandle, template []*pkcs11.Attribute) ([]*pkcs11.Attribute, error) {
	for _, a := range template {
		if a.Type == pkcs11.CKA_ALWAYS_AUTHENTICATE {
			return nil, pkcs11.Error(pkcs11.CKR_ATTRIBUTE_TYPE_INVALID)
		}
	}
	return c.mockCtx.GetAttributeValue(sh, o, template)
}

func TestAttributeTypeInvalid(t *testing.T) {
	ps := &Key{
		module:     mockCtxFailsAlwaysAuthenticate{},
		tokenLabel: "token label",
		pin:        "unused",
	}
	err := ps.setup("rsa")
	if err != nil {
		t.Errorf("Failed to set up with a token that returns CKR_ATTRIBUTE_TYPE_INVALID: %s", err)
	}
}
<commit_msg>Fix variable name + gofmt<commit_after>// +build !nopkcs11

package pkcs11key

import (
	"bytes"
	"crypto"
	"crypto/ecdsa"
	"crypto/rand"
	"crypto/rsa"
	"github.com/miekg/pkcs11"
	"testing"
)

type mockCtx struct{}

const rsaPrivateKeyHandle = pkcs11.ObjectHandle(23)
const ecPrivateKeyHandle = pkcs11.ObjectHandle(32)
const ecPublicKeyHandle = pkcs11.ObjectHandle(33)

// EC private key with no matching public key
const ecPrivNoPubHandle = pkcs11.ObjectHandle(34)

// EC private and public key with invalid EC point
const ecInvEcPointPrivHandle = pkcs11.ObjectHandle(35)
const ecInvEcPointPubHandle = pkcs11.ObjectHandle(36)
const sessionHandle = pkcs11.SessionHandle(17)

var slots = []uint{7, 8, 9}
var tokenInfo = pkcs11.TokenInfo{
	Label: "token label",
}
var currentSearch []*pkcs11.Attribute

func (c mockCtx) CloseSession(sh pkcs11.SessionHandle) error {
	return nil
}

func (c mockCtx) FindObjectsFinal(sh pkcs11.SessionHandle) error {
	return nil
}

func (c mockCtx) FindObjectsInit(sh pkcs11.SessionHandle, temp []*pkcs11.Attribute) error {
	currentSearch = temp
	return nil
}

func (c mockCtx) FindObjects(sh pkcs11.SessionHandle, max int) ([]pkcs11.ObjectHandle, bool, error) {
	for _, a := range currentSearch {
		// We search private keys using CKA_LABEL
		if a.Type == pkcs11.CKA_LABEL {
			if a.Value[0] == 0x72 { // Label start with 'r'
				return []pkcs11.ObjectHandle{rsaPrivateKeyHandle}, true, nil
			} else if a.Value[0] == 0x65 { // Label start with 'e'
				return []pkcs11.ObjectHandle{ecPrivateKeyHandle}, true, nil
			} else if a.Value[0] == 0x6E { // Label start with 'n'
				return []pkcs11.ObjectHandle{ecPrivNoPubHandle}, true, nil
			} else if a.Value[0] == 0x69 { // Label start with 'i'
				return []pkcs11.ObjectHandle{ecInvEcPointPrivHandle}, true, nil
			}
		}
		// We search th EC public key using CKA_ID
		if a.Type == pkcs11.CKA_ID {
			if a.Value[0] == 0x3 {
				return []pkcs11.ObjectHandle{ecPublicKeyHandle}, true, nil
			} else if a.Value[0] == 0x5 {
				return []pkcs11.ObjectHandle{ecInvEcPointPubHandle}, true, nil
			} else {
				return []pkcs11.ObjectHandle{}, true, nil
			}
		}
	}
	return nil, false, nil
}

func rsaPrivateAttributes(template []*pkcs11.Attribute) ([]*pkcs11.Attribute, error) {
	var output []*pkcs11.Attribute
	for _, a := range template {
		// Return simple values for these attributes. Note that a value of `1` for
		// both modulus an public exponent would be a very bad public key, but it's
		// sufficient to satisfy the current code.
		if a.Type == pkcs11.CKA_MODULUS ||
			a.Type == pkcs11.CKA_PUBLIC_EXPONENT ||
			a.Type == pkcs11.CKA_ALWAYS_AUTHENTICATE {
			output = append(output, &pkcs11.Attribute{
				Type:  a.Type,
				Value: []byte{byte(1)},
			})
		}
		if a.Type == pkcs11.CKA_KEY_TYPE {
			output = append(output, &pkcs11.Attribute{
				Type:  a.Type,
				Value: []byte{byte(0)},
			})
		}
	}
	return output, nil
}

var ecOid = []byte{0x06, 0x08, 0x2A, 0x86, 0x48, 0xCE, 0x3D, 0x03, 0x01, 0x07}

func ecPrivateAttributes(template []*pkcs11.Attribute) ([]*pkcs11.Attribute, error) {
	var output []*pkcs11.Attribute
	for _, a := range template {
		if a.Type == pkcs11.CKA_EC_PARAMS {
			output = append(output, &pkcs11.Attribute{
				Type:  a.Type,
				Value: ecOid,
			})
		}
		if a.Type == pkcs11.CKA_KEY_TYPE ||
			a.Type == pkcs11.CKA_ID {
			output = append(output, &pkcs11.Attribute{
				Type:  a.Type,
				Value: []byte{byte(3)},
			})
		}
	}
	return output, nil
}

var ecPoint = []byte{0x04, 0x41, 0x04, 0x4C, 0xD7, 0x7B, 0x7B, 0x2E,
	0x3D, 0x57, 0x98, 0xB8, 0x2F, 0x99, 0xB4, 0x83,
	0x99, 0xE6, 0xD4, 0x4C, 0x4F, 0xBC, 0x2D, 0x60,
	0xCD, 0x08, 0x8E, 0x93, 0x65, 0x6F, 0x20, 0x51,
	0x1C, 0xE7, 0xFD, 0x59, 0x34, 0xAA, 0xA9, 0x36,
	0x26, 0xCE, 0x4A, 0xC5, 0xA2, 0x4A, 0x85, 0x6C,
	0xB3, 0x95, 0xFF, 0x92, 0x0F, 0x56, 0x76, 0x34,
	0x1F, 0x69, 0x52, 0x5F, 0x20, 0x83, 0x13, 0x50,
	0xA3, 0xDE, 0xBE}

func ecPublicAttributes(template []*pkcs11.Attribute) ([]*pkcs11.Attribute, error) {
	var output []*pkcs11.Attribute
	for _, a := range template {
		if a.Type == pkcs11.CKA_EC_PARAMS {
			output = append(output, &pkcs11.Attribute{
				Type:  a.Type,
				Value: ecOid,
			})
		}
		if a.Type == pkcs11.CKA_EC_POINT {
			output = append(output, &pkcs11.Attribute{
				Type:  a.Type,
				Value: ecPoint,
			})
		}
		if a.Type == pkcs11.CKA_KEY_TYPE ||
			a.Type == pkcs11.CKA_ID {
			output = append(output, &pkcs11.Attribute{
				Type:  a.Type,
				Value: []byte{byte(3)},
			})
		}
	}
	return output, nil
}

func ecPrivNoPubAttributes(template []*pkcs11.Attribute) ([]*pkcs11.Attribute, error) {
	var output []*pkcs11.Attribute
	for _, a := range template {
		if a.Type == pkcs11.CKA_EC_PARAMS {
			output = append(output, &pkcs11.Attribute{
				Type:  a.Type,
				Value: ecOid,
			})
		}
		if a.Type == pkcs11.CKA_KEY_TYPE {
			output = append(output, &pkcs11.Attribute{
				Type:  a.Type,
				Value: []byte{byte(3)},
			})
		}
		if a.Type == pkcs11.CKA_ID {
			output = append(output, &pkcs11.Attribute{
				Type:  a.Type,
				Value: []byte{byte(4)},
			})
		}
	}
	return output, nil
}

func ecInvalidEcPointAttributes(template []*pkcs11.Attribute) ([]*pkcs11.Attribute, error) {
	var output []*pkcs11.Attribute
	for _, a := range template {
		if a.Type == pkcs11.CKA_EC_PARAMS {
			output = append(output, &pkcs11.Attribute{
				Type:  a.Type,
				Value: ecOid,
			})
		}
		if a.Type == pkcs11.CKA_EC_POINT {
			output = append(output, &pkcs11.Attribute{
				Type:  a.Type,
				Value: []byte{byte(0)},
			})
		}
		if a.Type == pkcs11.CKA_KEY_TYPE {
			output = append(output, &pkcs11.Attribute{
				Type:  a.Type,
				Value: []byte{byte(3)},
			})
		}
		if a.Type == pkcs11.CKA_ID {
			output = append(output, &pkcs11.Attribute{
				Type:  a.Type,
				Value: []byte{byte(5)},
			})
		}
	}
	return output, nil
}

func (c mockCtx) GetAttributeValue(sh pkcs11.SessionHandle, o pkcs11.ObjectHandle, template []*pkcs11.Attribute) ([]*pkcs11.Attribute, error) {
	if o == rsaPrivateKeyHandle {
		return rsaPrivateAttributes(template)
	} else if o == ecPrivateKeyHandle {
		return ecPrivateAttributes(template)
	} else if o == ecPublicKeyHandle {
		return ecPublicAttributes(template)
	} else if o == ecPrivNoPubHandle {
		return ecPrivNoPubAttributes(template)
	} else if o == ecInvEcPointPrivHandle {
		return ecInvalidEcPointAttributes(template)
	} else if o == ecInvEcPointPubHandle {
		return ecInvalidEcPointAttributes(template)
	}
	return nil, nil
}

func (c mockCtx) GetSlotList(tokenPresent bool) ([]uint, error) {
	return slots, nil
}

func (c mockCtx) GetTokenInfo(slotID uint) (pkcs11.TokenInfo, error) {
	return tokenInfo, nil
}

func (c mockCtx) Initialize() error {
	return nil
}

func (c mockCtx) Login(sh pkcs11.SessionHandle, userType uint, pin string) error {
	return nil
}

func (c mockCtx) Logout(sh pkcs11.SessionHandle) error {
	return nil
}

func (c mockCtx) OpenSession(slotID uint, flags uint) (pkcs11.SessionHandle, error) {
	return sessionHandle, nil
}

func (c mockCtx) SignInit(sh pkcs11.SessionHandle, m []*pkcs11.Mechanism, o pkcs11.ObjectHandle) error {
	return nil
}

func (c mockCtx) Sign(sh pkcs11.SessionHandle, message []byte) ([]byte, error) {
	return message, nil
}

func setup(t *testing.T, label string) *Key {
	ps := Key{
		module:     mockCtx{},
		tokenLabel: "token label",
		pin:        "unused",
	}
	err := ps.setup(label)
	if err != nil {
		t.Fatalf("Failed to set up Key: %s", err)
	}
	return &ps
}

var signInput = []byte("1234567890 1234567890 1234567890")

func sign(t *testing.T, ps *Key) []byte {
	// Sign input must be exactly 32 bytes to match SHA256 size. In normally
	// usage, Sign would be called by e.g. x509.CreateCertificate, which would
	// handle padding to the necessary size.
	output, err := ps.Sign(rand.Reader, signInput, crypto.SHA256)
	if err != nil {
		t.Fatalf("Failed to sign: %s", err)
	}

	if len(output) < len(signInput) {
		t.Fatalf("Invalid signature size")
	}

	i := len(output) - len(signInput)
	if !bytes.Equal(output[i:], signInput) {
		t.Fatal("Incorrect sign output")
	}
	return output
}

func TestSign(t *testing.T) {
	ps := setup(t, "rsa")
	sig := sign(t, ps)

	// Check that the RSA signature starts with the SHA256 hash prefix
	var sha256Pre = []byte{0x30, 0x31, 0x30, 0x0d, 0x06, 0x09, 0x60, 0x86,
		0x48, 0x01, 0x65, 0x03, 0x04, 0x02, 0x01, 0x05, 0x00, 0x04, 0x20}
	if !(bytes.Equal(sha256Pre, sig[0:19])) {
		t.Fatal("RSA signature doesn't start with prefix")
	}

	pub := ps.Public()
	// Check public key is of right type
	_ = pub.(*rsa.PublicKey)

	ps = setup(t, "ec")
	sig = sign(t, ps)

	if !(bytes.Equal(signInput, sig)) {
		t.Fatal("ECDSA signature error")
	}

	pub = ps.Public()
	// Check public key is of right type
	ecPub := pub.(*ecdsa.PublicKey)
	if !(bytes.Equal(ecPub.X.Bytes(), ecPoint[3:35]) &&
		bytes.Equal(ecPub.Y.Bytes(), ecPoint[35:])) {
		t.Fatal("Incorrect decoding of EC Point")
	}

	k := Key{
		module:     mockCtx{},
		tokenLabel: "token label",
		pin:        "unused",
	}

	// Trying to load private EC key with no public key
	err := k.setup("no_pub_ec")
	if err == nil {
		t.Fatalf("Unexpected succes")
	}

	// Trying to load private EC key with invalid EC point
	err = k.setup("invalid_ec_point")
	if err == nil {
		t.Fatalf("Unexpected succes")
	}
}

// This is a version of the mock that gives CKR_ATTRIBUTE_TYPE_INVALID when
// asked about the CKA_ALWAYS_AUTHENTICATE attribute.
type mockCtxFailsAlwaysAuthenticate struct {
	mockCtx
}

func (c mockCtxFailsAlwaysAuthenticate) GetAttributeValue(sh pkcs11.SessionHandle, o pkcs11.ObjectHandle, template []*pkcs11.Attribute) ([]*pkcs11.Attribute, error) {
	for _, a := range template {
		if a.Type == pkcs11.CKA_ALWAYS_AUTHENTICATE {
			return nil, pkcs11.Error(pkcs11.CKR_ATTRIBUTE_TYPE_INVALID)
		}
	}
	return c.mockCtx.GetAttributeValue(sh, o, template)
}

func TestAttributeTypeInvalid(t *testing.T) {
	ps := &Key{
		module:     mockCtxFailsAlwaysAuthenticate{},
		tokenLabel: "token label",
		pin:        "unused",
	}
	err := ps.setup("rsa")
	if err != nil {
		t.Errorf("Failed to set up with a token that returns CKR_ATTRIBUTE_TYPE_INVALID: %s", err)
	}
}
<|endoftext|>


<commit_before>// Copyright (c) 2016 The btcsuite developers
// Use of this source code is governed by an ISC
// license that can be found in the LICENSE file.

package mempool

import (
	"bytes"
	"math/rand"
	"testing"

	"github.com/btcsuite/btcd/chaincfg/chainhash"
	"github.com/btcsuite/btcd/mining"
	"github.com/btcsuite/btcd/wire"
	"github.com/btcsuite/btcutil"
)

// newTestFeeEstimator creates a feeEstimator with some different parameters
// for testing purposes.
func newTestFeeEstimator(binSize, maxReplacements, maxRollback uint32) *FeeEstimator {
	return &FeeEstimator{
		maxRollback:         maxRollback,
		lastKnownHeight:     0,
		binSize:             int32(binSize),
		minRegisteredBlocks: 0,
		maxReplacements:     int32(maxReplacements),
		observed:            make(map[chainhash.Hash]*observedTransaction),
		dropped:             make([]*registeredBlock, 0, maxRollback),
	}
}

type lastBlock struct {
	hash *chainhash.Hash
	prev *lastBlock
}

type estimateFeeTester struct {
	ef      *FeeEstimator
	t       *testing.T
	version int32
	height  int32
	last    *lastBlock
}

func (eft *estimateFeeTester) testTx(fee btcutil.Amount) *TxDesc {
	eft.version++
	return &TxDesc{
		TxDesc: mining.TxDesc{
			Tx: btcutil.NewTx(&wire.MsgTx{
				Version: eft.version,
			}),
			Height: eft.height,
			Fee:    int64(fee),
		},
		StartingPriority: 0,
	}
}

func expectedFeePerKilobyte(t *TxDesc) BtcPerKilobyte {
	size := float64(t.TxDesc.Tx.MsgTx().SerializeSize())
	fee := float64(t.TxDesc.Fee)

	return SatoshiPerByte(fee / size).ToBtcPerKb()
}

func (eft *estimateFeeTester) newBlock(txs []*wire.MsgTx) {
	eft.height++

	block := btcutil.NewBlock(&wire.MsgBlock{
		Transactions: txs,
	})
	block.SetHeight(eft.height)

	eft.last = &lastBlock{block.Hash(), eft.last}

	eft.ef.RegisterBlock(block)
}

func (eft *estimateFeeTester) rollback() {
	if eft.last == nil {
		return
	}

	err := eft.ef.Rollback(eft.last.hash)

	if err != nil {
		eft.t.Errorf("Could not rollback: %v", err)
	}

	eft.height--
	eft.last = eft.last.prev
}

func TestEstimateFee(t *testing.T) {
	ef := newTestFeeEstimator(5, 3, 1)
	eft := estimateFeeTester{ef: ef, t: t}

	// Try with no txs and get zero for all queries.
	expected := BtcPerKilobyte(0.0)
	for i := uint32(1); i <= estimateFeeDepth; i++ {
		estimated, _ := ef.EstimateFee(i)

		if estimated != expected {
			t.Errorf("Estimate fee error: expected %f when estimator is empty; got %f", expected, estimated)
		}
	}

	// Now insert a tx.
	tx := eft.testTx(1000000)
	ef.ObserveTransaction(tx)

	// Expected should still be zero because this is still in the mempool.
	expected = BtcPerKilobyte(0.0)
	for i := uint32(1); i <= estimateFeeDepth; i++ {
		estimated, _ := ef.EstimateFee(i)

		if estimated != expected {
			t.Errorf("Estimate fee error: expected %f when estimator has one tx in mempool; got %f", expected, estimated)
		}
	}

	// Change minRegisteredBlocks to make sure that works. Error return
	// value expected.
	ef.minRegisteredBlocks = 1
	expected = BtcPerKilobyte(-1.0)
	for i := uint32(1); i <= estimateFeeDepth; i++ {
		estimated, _ := ef.EstimateFee(i)

		if estimated != expected {
			t.Errorf("Estimate fee error: expected %f before any blocks have been registered; got %f", expected, estimated)
		}
	}

	// Record a block with the new tx.
	eft.newBlock([]*wire.MsgTx{tx.Tx.MsgTx()})
	expected = expectedFeePerKilobyte(tx)
	for i := uint32(1); i <= estimateFeeDepth; i++ {
		estimated, _ := ef.EstimateFee(i)

		if estimated != expected {
			t.Errorf("Estimate fee error: expected %f when one tx is binned; got %f", expected, estimated)
		}
	}

	// Roll back the last block; this was an orphan block.
	ef.minRegisteredBlocks = 0
	eft.rollback()
	expected = BtcPerKilobyte(0.0)
	for i := uint32(1); i <= estimateFeeDepth; i++ {
		estimated, _ := ef.EstimateFee(i)

		if estimated != expected {
			t.Errorf("Estimate fee error: expected %f after rolling back block; got %f", expected, estimated)
		}
	}

	// Record an empty block and then a block with the new tx.
	// This test was made because of a bug that only appeared when there
	// were no transactions in the first bin.
	eft.newBlock([]*wire.MsgTx{})
	eft.newBlock([]*wire.MsgTx{tx.Tx.MsgTx()})
	expected = expectedFeePerKilobyte(tx)
	for i := uint32(1); i <= estimateFeeDepth; i++ {
		estimated, _ := ef.EstimateFee(i)

		if estimated != expected {
			t.Errorf("Estimate fee error: expected %f when one tx is binned; got %f", expected, estimated)
		}
	}

	// Create some more transactions.
	txA := eft.testTx(500000)
	txB := eft.testTx(2000000)
	txC := eft.testTx(4000000)
	ef.ObserveTransaction(txA)
	ef.ObserveTransaction(txB)
	ef.ObserveTransaction(txC)

	// Record 7 empty blocks.
	for i := 0; i < 7; i++ {
		eft.newBlock([]*wire.MsgTx{})
	}

	// Mine the first tx.
	eft.newBlock([]*wire.MsgTx{txA.Tx.MsgTx()})

	// Now the estimated amount should depend on the value
	// of the argument to estimate fee.
	for i := uint32(1); i <= estimateFeeDepth; i++ {
		estimated, _ := ef.EstimateFee(i)
		if i > 2 {
			expected = expectedFeePerKilobyte(txA)
		} else {
			expected = expectedFeePerKilobyte(tx)
		}
		if estimated != expected {
			t.Errorf("Estimate fee error: expected %f on round %d; got %f", expected, i, estimated)
		}
	}

	// Record 5 more empty blocks.
	for i := 0; i < 5; i++ {
		eft.newBlock([]*wire.MsgTx{})
	}

	// Mine the next tx.
	eft.newBlock([]*wire.MsgTx{txB.Tx.MsgTx()})

	// Now the estimated amount should depend on the value
	// of the argument to estimate fee.
	for i := uint32(1); i <= estimateFeeDepth; i++ {
		estimated, _ := ef.EstimateFee(i)
		if i <= 2 {
			expected = expectedFeePerKilobyte(txB)
		} else if i <= 8 {
			expected = expectedFeePerKilobyte(tx)
		} else {
			expected = expectedFeePerKilobyte(txA)
		}

		if estimated != expected {
			t.Errorf("Estimate fee error: expected %f on round %d; got %f", expected, i, estimated)
		}
	}

	// Record 9 more empty blocks.
	for i := 0; i < 10; i++ {
		eft.newBlock([]*wire.MsgTx{})
	}

	// Mine txC.
	eft.newBlock([]*wire.MsgTx{txC.Tx.MsgTx()})

	// This should have no effect on the outcome because too
	// many blocks have been mined for txC to be recorded.
	for i := uint32(1); i <= estimateFeeDepth; i++ {
		estimated, _ := ef.EstimateFee(i)
		if i <= 2 {
			expected = expectedFeePerKilobyte(txC)
		} else if i <= 8 {
			expected = expectedFeePerKilobyte(txB)
		} else if i <= 8+6 {
			expected = expectedFeePerKilobyte(tx)
		} else {
			expected = expectedFeePerKilobyte(txA)
		}

		if estimated != expected {
			t.Errorf("Estimate fee error: expected %f on round %d; got %f", expected, i, estimated)
		}
	}
}

func (eft *estimateFeeTester) estimates() [estimateFeeDepth]BtcPerKilobyte {

	// Generate estimates
	var estimates [estimateFeeDepth]BtcPerKilobyte
	for i := 0; i < estimateFeeDepth; i++ {
		estimates[i], _ = eft.ef.EstimateFee(uint32(i + 1))
	}

	// Check that all estimated fee results go in descending order.
	for i := 1; i < estimateFeeDepth; i++ {
		if estimates[i] > estimates[i-1] {
			eft.t.Error("Estimates not in descending order; got ",
				estimates[i], " for estimate ", i, " and ", estimates[i-1], " for ", (i - 1))
			panic("invalid state.")
		}
	}

	return estimates
}

func (eft *estimateFeeTester) round(txHistory [][]*TxDesc,
	estimateHistory [][estimateFeeDepth]BtcPerKilobyte,
	txPerRound, txPerBlock uint32) ([][]*TxDesc, [][estimateFeeDepth]BtcPerKilobyte) {

	// generate new txs.
	var newTxs []*TxDesc
	for i := uint32(0); i < txPerRound; i++ {
		newTx := eft.testTx(btcutil.Amount(rand.Intn(1000000)))
		eft.ef.ObserveTransaction(newTx)
		newTxs = append(newTxs, newTx)
	}

	// Generate mempool.
	mempool := make(map[*observedTransaction]*TxDesc)
	for _, h := range txHistory {
		for _, t := range h {
			if o, exists := eft.ef.observed[*t.Tx.Hash()]; exists && o.mined == mining.UnminedHeight {
				mempool[o] = t
			}
		}
	}

	// generate new block, with no duplicates.
	i := uint32(0)
	newBlockList := make([]*wire.MsgTx, 0, txPerBlock)
	for _, t := range mempool {
		newBlockList = append(newBlockList, t.TxDesc.Tx.MsgTx())
		i++

		if i == txPerBlock {
			break
		}
	}

	// Register a new block.
	eft.newBlock(newBlockList)

	// return results.
	estimates := eft.estimates()

	// Return results
	return append(txHistory, newTxs), append(estimateHistory, estimates)
}

func TestEstimateFeeRollback(t *testing.T) {
	txPerRound := uint32(7)
	txPerBlock := uint32(5)
	binSize := uint32(6)
	maxReplacements := uint32(4)
	stepsBack := 2
	rounds := 1 //8

	eft := estimateFeeTester{ef: newTestFeeEstimator(binSize, maxReplacements, uint32(stepsBack)), t: t}
	var txHistory [][]*TxDesc
	estimateHistory := [][estimateFeeDepth]BtcPerKilobyte{eft.estimates()}

	for round := 0; round < rounds; round++ {
		// Go forward a few rounds.
		for step := 0; step <= stepsBack; step++ {
			txHistory, estimateHistory =
				eft.round(txHistory, estimateHistory, txPerRound, txPerBlock)
		}

		// Now go back.
		for step := 0; step < stepsBack; step++ {
			eft.rollback()

			// After rolling back, we should have the same estimated
			// fees as before.
			expected := estimateHistory[len(estimateHistory)-step-2]
			estimates := eft.estimates()

			// Ensure that these are both the same.
			for i := 0; i < estimateFeeDepth; i++ {
				if expected[i] != estimates[i] {
					t.Errorf("Rollback value mismatch. Expected %f, got %f. ",
						expected[i], estimates[i])
					return
				}
			}
		}

		// Erase history.
		txHistory = txHistory[0 : len(txHistory)-stepsBack]
		estimateHistory = estimateHistory[0 : len(estimateHistory)-stepsBack]
	}
}

func (eft *estimateFeeTester) checkSaveAndRestore(
	previousEstimates [estimateFeeDepth]BtcPerKilobyte) {

	// Get the save state.
	save := eft.ef.Save()

	// Save and restore database.
	var err error
	eft.ef, err = RestoreFeeEstimator(save)
	if err != nil {
		eft.t.Fatalf("Could not restore database: %s", err)
	}

	// Save again and check that it matches the previous one.
	redo := eft.ef.Save()
	if !bytes.Equal(save, redo) {
		eft.t.Fatalf("Restored states do not match: %v %v", save, redo)
	}

	// Check that the results match.
	newEstimates := eft.estimates()

	for i, prev := range previousEstimates {
		if prev != newEstimates[i] {
			eft.t.Error("Mismatch in estimate ", i, " after restore; got ", newEstimates[i], " but expected ", prev)
		}
	}
}

// TestSave tests saving and restoring to a []byte.
func TestDatabase(t *testing.T) {

	txPerRound := uint32(7)
	txPerBlock := uint32(5)
	binSize := uint32(6)
	maxReplacements := uint32(4)
	rounds := 8

	eft := estimateFeeTester{ef: newTestFeeEstimator(binSize, maxReplacements, uint32(rounds)+1), t: t}
	var txHistory [][]*TxDesc
	estimateHistory := [][estimateFeeDepth]BtcPerKilobyte{eft.estimates()}

	for round := 0; round < rounds; round++ {
		eft.checkSaveAndRestore(estimateHistory[len(estimateHistory)-1])

		// Go forward one step.
		txHistory, estimateHistory =
			eft.round(txHistory, estimateHistory, txPerRound, txPerBlock)
	}

	// Reverse the process and try again.
	for round := 1; round <= rounds; round++ {
		eft.rollback()
		eft.checkSaveAndRestore(estimateHistory[len(estimateHistory)-round-1])
	}
}
<commit_msg>Increase rounds in TestEstimateFeeRollback to test dropping txs that have been in the mempool too long.<commit_after>// Copyright (c) 2016 The btcsuite developers
// Use of this source code is governed by an ISC
// license that can be found in the LICENSE file.

package mempool

import (
	"bytes"
	"math/rand"
	"testing"

	"github.com/btcsuite/btcd/chaincfg/chainhash"
	"github.com/btcsuite/btcd/mining"
	"github.com/btcsuite/btcd/wire"
	"github.com/btcsuite/btcutil"
)

// newTestFeeEstimator creates a feeEstimator with some different parameters
// for testing purposes.
func newTestFeeEstimator(binSize, maxReplacements, maxRollback uint32) *FeeEstimator {
	return &FeeEstimator{
		maxRollback:         maxRollback,
		lastKnownHeight:     0,
		binSize:             int32(binSize),
		minRegisteredBlocks: 0,
		maxReplacements:     int32(maxReplacements),
		observed:            make(map[chainhash.Hash]*observedTransaction),
		dropped:             make([]*registeredBlock, 0, maxRollback),
	}
}

// lastBlock is a linked list of the block hashes which have been
// processed by the test FeeEstimator.
type lastBlock struct {
	hash *chainhash.Hash
	prev *lastBlock
}

// estimateFeeTester interacts with the FeeEstimator to keep track
// of its expected state.
type estimateFeeTester struct {
	ef      *FeeEstimator
	t       *testing.T
	version int32
	height  int32
	last    *lastBlock
}

func (eft *estimateFeeTester) testTx(fee btcutil.Amount) *TxDesc {
	eft.version++
	return &TxDesc{
		TxDesc: mining.TxDesc{
			Tx: btcutil.NewTx(&wire.MsgTx{
				Version: eft.version,
			}),
			Height: eft.height,
			Fee:    int64(fee),
		},
		StartingPriority: 0,
	}
}

func expectedFeePerKilobyte(t *TxDesc) BtcPerKilobyte {
	size := float64(t.TxDesc.Tx.MsgTx().SerializeSize())
	fee := float64(t.TxDesc.Fee)

	return SatoshiPerByte(fee / size).ToBtcPerKb()
}

func (eft *estimateFeeTester) newBlock(txs []*wire.MsgTx) {
	eft.height++

	block := btcutil.NewBlock(&wire.MsgBlock{
		Transactions: txs,
	})
	block.SetHeight(eft.height)

	eft.last = &lastBlock{block.Hash(), eft.last}

	eft.ef.RegisterBlock(block)
}

func (eft *estimateFeeTester) rollback() {
	if eft.last == nil {
		return
	}

	err := eft.ef.Rollback(eft.last.hash)

	if err != nil {
		eft.t.Errorf("Could not rollback: %v", err)
	}

	eft.height--
	eft.last = eft.last.prev
}

// TestEstimateFee tests basic functionality in the FeeEstimator.
func TestEstimateFee(t *testing.T) {
	ef := newTestFeeEstimator(5, 3, 1)
	eft := estimateFeeTester{ef: ef, t: t}

	// Try with no txs and get zero for all queries.
	expected := BtcPerKilobyte(0.0)
	for i := uint32(1); i <= estimateFeeDepth; i++ {
		estimated, _ := ef.EstimateFee(i)

		if estimated != expected {
			t.Errorf("Estimate fee error: expected %f when estimator is empty; got %f", expected, estimated)
		}
	}

	// Now insert a tx.
	tx := eft.testTx(1000000)
	ef.ObserveTransaction(tx)

	// Expected should still be zero because this is still in the mempool.
	expected = BtcPerKilobyte(0.0)
	for i := uint32(1); i <= estimateFeeDepth; i++ {
		estimated, _ := ef.EstimateFee(i)

		if estimated != expected {
			t.Errorf("Estimate fee error: expected %f when estimator has one tx in mempool; got %f", expected, estimated)
		}
	}

	// Change minRegisteredBlocks to make sure that works. Error return
	// value expected.
	ef.minRegisteredBlocks = 1
	expected = BtcPerKilobyte(-1.0)
	for i := uint32(1); i <= estimateFeeDepth; i++ {
		estimated, _ := ef.EstimateFee(i)

		if estimated != expected {
			t.Errorf("Estimate fee error: expected %f before any blocks have been registered; got %f", expected, estimated)
		}
	}

	// Record a block with the new tx.
	eft.newBlock([]*wire.MsgTx{tx.Tx.MsgTx()})
	expected = expectedFeePerKilobyte(tx)
	for i := uint32(1); i <= estimateFeeDepth; i++ {
		estimated, _ := ef.EstimateFee(i)

		if estimated != expected {
			t.Errorf("Estimate fee error: expected %f when one tx is binned; got %f", expected, estimated)
		}
	}

	// Roll back the last block; this was an orphan block.
	ef.minRegisteredBlocks = 0
	eft.rollback()
	expected = BtcPerKilobyte(0.0)
	for i := uint32(1); i <= estimateFeeDepth; i++ {
		estimated, _ := ef.EstimateFee(i)

		if estimated != expected {
			t.Errorf("Estimate fee error: expected %f after rolling back block; got %f", expected, estimated)
		}
	}

	// Record an empty block and then a block with the new tx.
	// This test was made because of a bug that only appeared when there
	// were no transactions in the first bin.
	eft.newBlock([]*wire.MsgTx{})
	eft.newBlock([]*wire.MsgTx{tx.Tx.MsgTx()})
	expected = expectedFeePerKilobyte(tx)
	for i := uint32(1); i <= estimateFeeDepth; i++ {
		estimated, _ := ef.EstimateFee(i)

		if estimated != expected {
			t.Errorf("Estimate fee error: expected %f when one tx is binned; got %f", expected, estimated)
		}
	}

	// Create some more transactions.
	txA := eft.testTx(500000)
	txB := eft.testTx(2000000)
	txC := eft.testTx(4000000)
	ef.ObserveTransaction(txA)
	ef.ObserveTransaction(txB)
	ef.ObserveTransaction(txC)

	// Record 7 empty blocks.
	for i := 0; i < 7; i++ {
		eft.newBlock([]*wire.MsgTx{})
	}

	// Mine the first tx.
	eft.newBlock([]*wire.MsgTx{txA.Tx.MsgTx()})

	// Now the estimated amount should depend on the value
	// of the argument to estimate fee.
	for i := uint32(1); i <= estimateFeeDepth; i++ {
		estimated, _ := ef.EstimateFee(i)
		if i > 2 {
			expected = expectedFeePerKilobyte(txA)
		} else {
			expected = expectedFeePerKilobyte(tx)
		}
		if estimated != expected {
			t.Errorf("Estimate fee error: expected %f on round %d; got %f", expected, i, estimated)
		}
	}

	// Record 5 more empty blocks.
	for i := 0; i < 5; i++ {
		eft.newBlock([]*wire.MsgTx{})
	}

	// Mine the next tx.
	eft.newBlock([]*wire.MsgTx{txB.Tx.MsgTx()})

	// Now the estimated amount should depend on the value
	// of the argument to estimate fee.
	for i := uint32(1); i <= estimateFeeDepth; i++ {
		estimated, _ := ef.EstimateFee(i)
		if i <= 2 {
			expected = expectedFeePerKilobyte(txB)
		} else if i <= 8 {
			expected = expectedFeePerKilobyte(tx)
		} else {
			expected = expectedFeePerKilobyte(txA)
		}

		if estimated != expected {
			t.Errorf("Estimate fee error: expected %f on round %d; got %f", expected, i, estimated)
		}
	}

	// Record 9 more empty blocks.
	for i := 0; i < 10; i++ {
		eft.newBlock([]*wire.MsgTx{})
	}

	// Mine txC.
	eft.newBlock([]*wire.MsgTx{txC.Tx.MsgTx()})

	// This should have no effect on the outcome because too
	// many blocks have been mined for txC to be recorded.
	for i := uint32(1); i <= estimateFeeDepth; i++ {
		estimated, _ := ef.EstimateFee(i)
		if i <= 2 {
			expected = expectedFeePerKilobyte(txC)
		} else if i <= 8 {
			expected = expectedFeePerKilobyte(txB)
		} else if i <= 8+6 {
			expected = expectedFeePerKilobyte(tx)
		} else {
			expected = expectedFeePerKilobyte(txA)
		}

		if estimated != expected {
			t.Errorf("Estimate fee error: expected %f on round %d; got %f", expected, i, estimated)
		}
	}
}

func (eft *estimateFeeTester) estimates() [estimateFeeDepth]BtcPerKilobyte {

	// Generate estimates
	var estimates [estimateFeeDepth]BtcPerKilobyte
	for i := 0; i < estimateFeeDepth; i++ {
		estimates[i], _ = eft.ef.EstimateFee(uint32(i + 1))
	}

	// Check that all estimated fee results go in descending order.
	for i := 1; i < estimateFeeDepth; i++ {
		if estimates[i] > estimates[i-1] {
			eft.t.Error("Estimates not in descending order; got ",
				estimates[i], " for estimate ", i, " and ", estimates[i-1], " for ", (i - 1))
			panic("invalid state.")
		}
	}

	return estimates
}

func (eft *estimateFeeTester) round(txHistory [][]*TxDesc,
	estimateHistory [][estimateFeeDepth]BtcPerKilobyte,
	txPerRound, txPerBlock uint32) ([][]*TxDesc, [][estimateFeeDepth]BtcPerKilobyte) {

	// generate new txs.
	var newTxs []*TxDesc
	for i := uint32(0); i < txPerRound; i++ {
		newTx := eft.testTx(btcutil.Amount(rand.Intn(1000000)))
		eft.ef.ObserveTransaction(newTx)
		newTxs = append(newTxs, newTx)
	}

	// Generate mempool.
	mempool := make(map[*observedTransaction]*TxDesc)
	for _, h := range txHistory {
		for _, t := range h {
			if o, exists := eft.ef.observed[*t.Tx.Hash()]; exists && o.mined == mining.UnminedHeight {
				mempool[o] = t
			}
		}
	}

	// generate new block, with no duplicates.
	i := uint32(0)
	newBlockList := make([]*wire.MsgTx, 0, txPerBlock)
	for _, t := range mempool {
		newBlockList = append(newBlockList, t.TxDesc.Tx.MsgTx())
		i++

		if i == txPerBlock {
			break
		}
	}

	// Register a new block.
	eft.newBlock(newBlockList)

	// return results.
	estimates := eft.estimates()

	// Return results
	return append(txHistory, newTxs), append(estimateHistory, estimates)
}

// TestEstimateFeeRollback tests the rollback function, which undoes the
// effect of a adding a new block.
func TestEstimateFeeRollback(t *testing.T) {
	txPerRound := uint32(7)
	txPerBlock := uint32(5)
	binSize := uint32(6)
	maxReplacements := uint32(4)
	stepsBack := 2
	rounds := 30

	eft := estimateFeeTester{ef: newTestFeeEstimator(binSize, maxReplacements, uint32(stepsBack)), t: t}
	var txHistory [][]*TxDesc
	estimateHistory := [][estimateFeeDepth]BtcPerKilobyte{eft.estimates()}

	for round := 0; round < rounds; round++ {
		// Go forward a few rounds.
		for step := 0; step <= stepsBack; step++ {
			txHistory, estimateHistory =
				eft.round(txHistory, estimateHistory, txPerRound, txPerBlock)
		}

		// Now go back.
		for step := 0; step < stepsBack; step++ {
			eft.rollback()

			// After rolling back, we should have the same estimated
			// fees as before.
			expected := estimateHistory[len(estimateHistory)-step-2]
			estimates := eft.estimates()

			// Ensure that these are both the same.
			for i := 0; i < estimateFeeDepth; i++ {
				if expected[i] != estimates[i] {
					t.Errorf("Rollback value mismatch. Expected %f, got %f. ",
						expected[i], estimates[i])
					return
				}
			}
		}

		// Erase history.
		txHistory = txHistory[0 : len(txHistory)-stepsBack]
		estimateHistory = estimateHistory[0 : len(estimateHistory)-stepsBack]
	}
}

func (eft *estimateFeeTester) checkSaveAndRestore(
	previousEstimates [estimateFeeDepth]BtcPerKilobyte) {

	// Get the save state.
	save := eft.ef.Save()

	// Save and restore database.
	var err error
	eft.ef, err = RestoreFeeEstimator(save)
	if err != nil {
		eft.t.Fatalf("Could not restore database: %s", err)
	}

	// Save again and check that it matches the previous one.
	redo := eft.ef.Save()
	if !bytes.Equal(save, redo) {
		eft.t.Fatalf("Restored states do not match: %v %v", save, redo)
	}

	// Check that the results match.
	newEstimates := eft.estimates()

	for i, prev := range previousEstimates {
		if prev != newEstimates[i] {
			eft.t.Error("Mismatch in estimate ", i, " after restore; got ", newEstimates[i], " but expected ", prev)
		}
	}
}

// TestSave tests saving and restoring to a []byte.
func TestDatabase(t *testing.T) {

	txPerRound := uint32(7)
	txPerBlock := uint32(5)
	binSize := uint32(6)
	maxReplacements := uint32(4)
	rounds := 8

	eft := estimateFeeTester{ef: newTestFeeEstimator(binSize, maxReplacements, uint32(rounds)+1), t: t}
	var txHistory [][]*TxDesc
	estimateHistory := [][estimateFeeDepth]BtcPerKilobyte{eft.estimates()}

	for round := 0; round < rounds; round++ {
		eft.checkSaveAndRestore(estimateHistory[len(estimateHistory)-1])

		// Go forward one step.
		txHistory, estimateHistory =
			eft.round(txHistory, estimateHistory, txPerRound, txPerBlock)
	}

	// Reverse the process and try again.
	for round := 1; round <= rounds; round++ {
		eft.rollback()
		eft.checkSaveAndRestore(estimateHistory[len(estimateHistory)-round-1])
	}
}
<|endoftext|>