bitcoin_node.go

package bitcoin_reader

import (
	"context"
	"io"
	"net"
	"sync"
	"sync/atomic"
	"time"

	"github.com/tokenized/logger"
	"github.com/tokenized/pkg/bitcoin"
	"github.com/tokenized/pkg/wire"
	"github.com/tokenized/threads"

	"github.com/google/uuid"
	"github.com/pkg/errors"
)

const (
	ServiceFull = 0x01
)

var (
	ErrTimeout = errors.New("Timeout")
	ErrBusy    = errors.New("Busy")

	// ErrNotFullService node is not a full service node.
	ErrNotFullService = errors.New("Not Full Service")
)

// BitcoinNode is a connection to a Bitcoin node in the peer to peer network that can be used to
// send requests.
type BitcoinNode struct {
	id        uuid.UUID
	address   string
	userAgent string
	config    *Config

	headers HeaderRepository
	peers   PeerRepository

	connection              net.Conn // Connection to trusted full node
	connectionClosedLocally bool
	connectionLock          sync.Mutex
	// IP             net.IP
	// Port           uint16

	pingNonce uint64
	pingSent  time.Time

	handlers          MessageHandlers
	headerHandler     MessageHandlerFunction
	lastHeaderHash    *bitcoin.Hash32 // last header received from the node
	lastHeaderRequest []bitcoin.Hash32

	requestTime        *time.Time
	blockRequest       *bitcoin.Hash32
	blockHandler       HandleBlock
	blockReader        io.ReadCloser
	blockOnStop        OnStop
	lastRequestedBlock *bitcoin.Hash32

	txManager       *TxManager
	txReceivedCount uint64
	txReceivedSize  uint64

	outgoingMsgChannel MessageChannel
	handshakeChannel   chan wire.Message

	handshakeIsComplete atomic.Value
	isReady             atomic.Value
	isStopped           atomic.Value
	verified            atomic.Value
	protoconfCount      int

	isVerifyOnly bool // disconnect after chain verification

	interrupt <-chan interface{}

	sync.Mutex
}

type HeaderRepository interface {
	GetNewHeadersAvailableChannel() <-chan *wire.BlockHeader
	Height() int
	Hash(ctx context.Context, height int) (*bitcoin.Hash32, error)
	HashHeight(hash bitcoin.Hash32) int
	LastHash() bitcoin.Hash32
	LastTime() uint32
	PreviousHash(bitcoin.Hash32) (*bitcoin.Hash32, int)
	GetLocatorHashes(ctx context.Context, max int) ([]bitcoin.Hash32, error)
	GetVerifyOnlyLocatorHashes(ctx context.Context) ([]bitcoin.Hash32, error)
	VerifyHeader(ctx context.Context, header *wire.BlockHeader) error
	ProcessHeader(ctx context.Context, header *wire.BlockHeader) error
	Stop(ctx context.Context)
}

type PeerRepository interface {
	Add(ctx context.Context, address string) (bool, error)
	Get(ctx context.Context, minScore, maxScore int32) (PeerList, error)

	UpdateTime(ctx context.Context, address string) bool
	UpdateScore(ctx context.Context, address string, delta int32) bool
}

func NewBitcoinNode(address, userAgent string, config *Config, headers HeaderRepository,
	peers PeerRepository) *BitcoinNode {

	result := &BitcoinNode{
		id:               uuid.New(),
		address:          address,
		userAgent:        userAgent,
		config:           config,
		headers:          headers,
		peers:            peers,
		handlers:         make(MessageHandlers),
		handshakeChannel: make(chan wire.Message, 10),
	}

	result.handshakeIsComplete.Store(false)
	result.isReady.Store(false)
	result.isStopped.Store(false)
	result.verified.Store(false)

	// Only enable messages that are required for handshake and verification.
	result.handlers[wire.CmdVersion] = result.handleVersion
	result.handlers[wire.CmdVerAck] = result.handleVerack
	result.handlers[wire.CmdHeaders] = result.handleHeadersVerify
	result.handlers[wire.CmdProtoconf] = result.handleProtoconf
	result.handlers[wire.CmdPing] = result.handlePing
	result.handlers[wire.CmdReject] = result.handleReject

	// Extended messages must be handled to properly get the size of the message. The payload
	// message will still be ignored if the tx and block handlers aren't enabled.
	result.handlers[wire.CmdExtended] = result.handleExtended

	return result
}

func (n *BitcoinNode) ID() uuid.UUID {
	n.Lock()
	defer n.Unlock()

	return n.id
}

// SetVerifyOnly sets the node to only verify the correct chain and then disconnect.
func (n *BitcoinNode) SetVerifyOnly() {
	n.Lock()
	defer n.Unlock()

	n.isVerifyOnly = true
}

func (n *BitcoinNode) SetTxManager(txManager *TxManager) {
	n.Lock()
	defer n.Unlock()

	n.txManager = txManager
}

func (n *BitcoinNode) GetAndResetTxReceivedCount() (uint64, uint64) {
	n.Lock()
	defer n.Unlock()

	resultCount := n.txReceivedCount
	resultSize := n.txReceivedSize
	n.txReceivedCount = 0
	n.txReceivedSize = 0
	return resultCount, resultSize
}

func (n *BitcoinNode) IsBusy() bool {
	n.Lock()
	defer n.Unlock()

	return n.requestTime != nil
}

func (n *BitcoinNode) HasBlock(ctx context.Context, hash bitcoin.Hash32, height int) bool {
	n.Lock()
	id := n.id
	lastRequestedBlock := n.lastRequestedBlock
	lastHeaderHash := n.lastHeaderHash
	n.Unlock()

	ctx = logger.ContextWithLogFields(ctx, logger.Stringer("connection", id))

	if lastRequestedBlock != nil && lastRequestedBlock.Equal(&hash) {
		return false // already requested this block and failed
	}

	if lastHeaderHash == nil {
		return false
	}

	if lastHeaderHash.Equal(&hash) {
		return true
	}

	lastHeight := n.headers.HashHeight(*lastHeaderHash)
	if lastHeight == -1 {
		logger.WarnWithFields(ctx, []logger.Field{
			logger.Stringer("last_hash", lastHeaderHash),
		}, "Last header height not found")
		return false // node's last header isn't in our chain
	}

	return lastHeight >= height
}

func (n *BitcoinNode) RequestBlock(ctx context.Context, hash bitcoin.Hash32, handler HandleBlock,
	onStop OnStop) error {
	n.Lock()
	ctx = logger.ContextWithLogFields(ctx, logger.Stringer("connection", n.id))

	if n.requestTime != nil {
		n.Unlock()
		return ErrBusy
	}

	now := time.Now()
	n.requestTime = &now
	n.blockRequest = &hash
	n.handlers[wire.CmdBlock] = n.handleBlock
	n.blockHandler = handler
	n.blockReader = nil
	n.lastRequestedBlock = &hash
	n.Unlock()

	logger.InfoWithFields(ctx, []logger.Field{
		logger.Stringer("block_hash", hash),
		logger.Int("block_height", n.headers.HashHeight(hash)),
	}, "Requesting block")
	getBlocks := wire.NewMsgGetData() // Block request message
	getBlocks.AddInvVect(wire.NewInvVect(wire.InvTypeBlock, &hash))
	if err := n.sendMessage(ctx, getBlocks); err != nil {
		return errors.Wrap(err, "send block request")
	}

	n.Lock()
	n.blockOnStop = onStop
	n.Unlock()

	return nil
}

// CancelBlockRequest cancels a request for a block. It returns true if the block handler has
// already been called and started handling the block.
func (n *BitcoinNode) CancelBlockRequest(ctx context.Context, hash bitcoin.Hash32) bool {
	n.Lock()
	defer n.Unlock()

	if n.blockRequest == nil {
		logger.Warn(ctx, "Block request not found to cancel")
		return false
	}

	if !n.blockRequest.Equal(&hash) {
		logger.WarnWithFields(ctx, []logger.Field{
			logger.Stringer("current_block_hash", n.blockRequest),
		}, "Wrong block request found to cancel")
		return false
	}

	if n.blockReader != nil {
		// Stop in progress handling of block
		n.blockReader.Close()
		n.blockReader = nil
		n.blockOnStop = nil
		n.blockHandler = nil
		logger.Info(ctx, "Cancelled in progress block")
		return true
	}

	// Stop handling a block before it happens
	n.blockOnStop = nil
	n.blockHandler = nil
	logger.Info(ctx, "Cancelled block request before download started")
	return false
}

func (n *BitcoinNode) RequestHeaders(ctx context.Context) error {
	n.Lock()
	ctx = logger.ContextWithLogFields(ctx, logger.Stringer("connection", n.id))

	if n.requestTime != nil {
		n.Unlock()
		return ErrBusy
	}
	n.Unlock()

	logger.Verbose(ctx, "Requesting headers")
	if err := n.sendHeaderRequest(ctx); err != nil {
		return errors.Wrap(err, "send header request")
	}

	return nil
}

func (n *BitcoinNode) RequestTxs(ctx context.Context, txids []bitcoin.Hash32) error {
	ctx = logger.ContextWithLogFields(ctx, logger.Stringer("connection", n.ID()))
	logger.Info(ctx, "Requesting %d previous txs", len(txids))

	invRequest := wire.NewMsgGetData()
	for _, txid := range txids {
		hash := txid
		item := wire.NewInvVect(wire.InvTypeTx, &hash)

		if err := invRequest.AddInvVect(item); err != nil {
			// Too many requests for one message, send it and start a new message.
			if err := n.sendMessage(ctx, invRequest); err != nil {
				return errors.Wrap(err, "send tx request")
			}

			invRequest = wire.NewMsgGetData()
			if err := invRequest.AddInvVect(item); err != nil {
				return errors.Wrap(err, "add tx to request")
			}
		}
	}

	if len(invRequest.InvList) > 0 {
		if err := n.sendMessage(ctx, invRequest); err != nil {
			return errors.Wrap(err, "send tx request")
		}
	}

	return nil
}

func (n *BitcoinNode) SetBlockHandler(handler MessageHandlerFunction) {
	n.Lock()
	defer n.Unlock()

	if handler == nil {
		delete(n.handlers, wire.CmdBlock)
	} else {
		n.handlers[wire.CmdBlock] = handler
	}
}

func (n *BitcoinNode) SetHeaderHandler(handler MessageHandlerFunction) {
	n.Lock()
	defer n.Unlock()

	n.headerHandler = handler
}

func (n *BitcoinNode) SetTxHandler(handler MessageHandlerFunction) {
	n.Lock()
	defer n.Unlock()

	if handler == nil {
		delete(n.handlers, wire.CmdTx)
	} else {
		n.handlers[wire.CmdTx] = handler
	}
}

func (n *BitcoinNode) Run(ctx context.Context, interrupt <-chan interface{}) error {
	logger.VerboseWithFields(ctx, []logger.Field{
		logger.String("address", n.address),
	}, "Connecting to node")

	n.interrupt = interrupt

	if err := n.connect(ctx); err != nil {
		n.isReady.Store(false)
		n.isStopped.Store(true)
		logger.VerboseWithFields(ctx, []logger.Field{
			logger.String("address", n.address),
		}, "Failed to connect to node : %s", err)
		return nil
	}

	return n.run(ctx, interrupt)
}

func (n *BitcoinNode) run(ctx context.Context, interrupt <-chan interface{}) error {
	n.interrupt = interrupt
	n.outgoingMsgChannel.Open(1000)

	var stopper threads.StopCombiner
	var wait sync.WaitGroup

	stopper.Add(n) // close connection and outgoing channel to stop incoming and outgoing threads

	readIncomingThread, readIncomingComplete := threads.NewUninterruptableThreadComplete("Read Incoming",
		n.readIncoming, &wait)

	sendOutgoingThread, sendOutgoingComplete := threads.NewUninterruptableThreadComplete("Send Outgoing",
		n.sendOutgoing, &wait)

	pingThread, pingComplete := threads.NewPeriodicThreadComplete("Ping", n.sendPing,
		10*time.Minute, &wait)
	stopper.Add(pingThread)

	handshakeThread := threads.NewInterruptableThread("Handshake", n.handshake)
	handshakeThread.SetWait(&wait)
	stopper.Add(handshakeThread)

	// Start threads
	readIncomingThread.Start(ctx)
	sendOutgoingThread.Start(ctx)
	pingThread.Start(ctx)
	handshakeThread.Start(ctx)

	// Wait for a thread to complete
	select {
	case <-interrupt:
	case <-readIncomingComplete:
	case <-sendOutgoingComplete:
	case <-pingComplete:
	case <-time.After(n.config.Timeout.Duration):
		logger.Verbose(ctx, "Node reached timeout")
	}

	stopper.Stop(ctx)

	n.Lock()
	n.isReady.Store(false)
	blockOnStop := n.blockOnStop
	n.Unlock()

	if blockOnStop != nil {
		logger.Info(ctx, "Calling block request \"on stop\" function")
		waitWarning := logger.NewWaitingWarning(ctx, time.Second, "Call block \"on stop\"")
		blockOnStop(ctx)
		waitWarning.Cancel()
	}

	waitWarning := logger.NewWaitingWarning(ctx, 3*time.Second, "Node Shutdown")
	wait.Wait()
	waitWarning.Cancel()

	n.isStopped.Store(true)

	return threads.CombineErrors(
		handshakeThread.Error(),
		readIncomingThread.Error(),
		sendOutgoingThread.Error(),
	)
}

func (n *BitcoinNode) Stop(ctx context.Context) {
	logger.Info(ctx, "Stopping: %s", n.Address())
	n.connectionLock.Lock()
	if n.connection != nil {
		n.connection.Close()
		n.connection = nil
		n.connectionClosedLocally = true
	}
	n.connectionLock.Unlock()

	n.outgoingMsgChannel.Close()
}

func (n *BitcoinNode) HandshakeIsComplete() bool {
	return n.handshakeIsComplete.Load().(bool)
}

func (n *BitcoinNode) IsReady() bool {
	return n.isReady.Load().(bool)
}

func (n *BitcoinNode) IsStopped() bool {
	return n.isStopped.Load().(bool)
}

func (n *BitcoinNode) Verified() bool {
	return n.verified.Load().(bool)
}

func (n *BitcoinNode) Address() string {
	n.Lock()
	defer n.Unlock()

	return n.address
}

// handshake performs the initial handshake with the node.
func (n *BitcoinNode) handshake(ctx context.Context, interrupt <-chan interface{}) error {
	versionReceived := false
	verAckSent := false
	verAckReceived := false

	n.Lock()
	address := n.address
	userAgent := n.userAgent
	receiveTxs := n.txManager != nil
	n.Unlock()

	if err := n.sendMessage(ctx, buildVersionMsg(address, userAgent, n.headers.Height(),
		receiveTxs)); err != nil {
		return errors.Wrap(err, "send version")
	}

	for {
		select {
		case msg, ok := <-n.handshakeChannel:
			if !ok {
				return nil
			}

			switch message := msg.(type) {
			case *wire.MsgVersion:
				logger.VerboseWithFields(ctx, []logger.Field{
					logger.String("address", address),
					logger.String("user_agent", message.UserAgent),
					logger.Int32("protocol", message.ProtocolVersion),
					logger.Formatter("services", "%016x", message.Services),
					logger.Int32("block_height", message.LastBlock),
				}, "Version")
				versionReceived = true

				if !verAckSent {
					if err := n.sendMessage(ctx, &wire.MsgVerAck{}); err != nil {
						return errors.Wrap(err, "send ver ack")
					}
					verAckSent = true
				}

				if verAckReceived {
					return n.sendVerifyInitiation(ctx)
				}

			case *wire.MsgVerAck:
				verAckReceived = true
				if versionReceived {
					return n.sendVerifyInitiation(ctx)
				}
			}

		case <-time.After(3 * time.Second):
			logger.Verbose(ctx, "Handshake timed out")
			n.Stop(ctx)
			return nil

		case <-interrupt:
			return nil
		}
	}
}

func (n *BitcoinNode) sendVerifyInitiation(ctx context.Context) error {
	n.handshakeIsComplete.Store(true)

	if err := n.sendMessage(ctx, wire.NewMsgProtoconf()); err != nil {
		return errors.Wrap(err, "send protoconf")
	}

	// Send header request to check the node is on the same chain
	if err := n.sendVerifyHeaderRequest(ctx); err != nil {
		return errors.Wrap(err, "send verify header request")
	}

	return nil
}

func (n *BitcoinNode) accept(ctx context.Context) error {
	n.Lock()

	// Switch headers handler to tracking mode.
	n.handlers[wire.CmdHeaders] = n.handleHeadersTrack

	// Enable more commands. These messages are ignored before this point.
	n.handlers[wire.CmdAddr] = n.handleAddress
	n.handlers[wire.CmdPong] = n.handlePong
	n.handlers[wire.CmdGetAddr] = n.handleGetAddresses

	// Enable tx handling
	if n.txManager != nil {
		n.handlers[wire.CmdInv] = n.handleInventory
		n.handlers[wire.CmdTx] = n.handleTx
	}

	isVerifyOnly := n.isVerifyOnly

	n.isReady.Store(true)
	n.verified.Store(true)
	n.Unlock()

	if isVerifyOnly {
		logger.Verbose(ctx, "Disconnecting after chain verification")
		n.Stop(ctx)
		return nil
	}

	if err := n.sendMessage(ctx, wire.NewMsgSendHeaders()); err != nil {
		return errors.Wrap(err, "send \"sendheaders\" request")
	}

	if err := n.sendMessage(ctx, wire.NewMsgGetAddr()); err != nil {
		return errors.Wrap(err, "send peer request")
	}

	// Send initial header request to get any new headers the node might have.
	if err := n.sendInitialHeaderRequest(ctx); err != nil {
		return errors.Wrap(err, "send initial header request")
	}

	addresses, err := buildAddressesMessage(ctx, n.peers)
	if err != nil {
		return errors.Wrap(err, "build addresses")
	}

	logger.Verbose(ctx, "Sending %d addresses", len(addresses.AddrList))
	if err := n.sendMessage(ctx, addresses); err != nil {
		return errors.Wrap(err, "send addresses")
	}

	return nil
}

func (n *BitcoinNode) sendPing(ctx context.Context) error {
	n.Lock()
	defer n.Unlock()

	n.pingNonce = nonce()
	n.pingSent = time.Now()

	logger.Debug(ctx, "Sending ping 0x%16x", n.pingNonce)
	return n.sendMessage(ctx, wire.NewMsgPing(n.pingNonce))
}

func (n *BitcoinNode) connect(ctx context.Context) error {
	connection, err := net.DialTimeout("tcp", n.address, 5*time.Second)
	if err != nil {
		return err
	}

	addr := connection.RemoteAddr()
	ip, port := parseAddress(addr.String())
	if ip == nil {
		logger.Info(ctx, "Connected to unknown IP")
	} else {
		logger.Verbose(ctx, "Connected to %s:%d", ip.String(), port)
	}

	n.connectionLock.Lock()
	n.connection = connection
	n.connectionLock.Unlock()
	n.peers.UpdateTime(ctx, n.address)
	return nil
}

func (n *BitcoinNode) mockConnect(ctx context.Context, connection net.Conn) error {
	logger.Info(ctx, "Connected to mock connection")

	n.connectionLock.Lock()
	n.connection = connection
	n.connectionLock.Unlock()
	return nil
}