<commit_before>// Copyright (c) 2013-2015 Conformal Systems LLC.
// Use of this source code is governed by an ISC
// license that can be found in the LICENSE file.

package wire

import (
	"bytes"
	"fmt"
	"io"
	"unicode/utf8"
)

// MessageHeaderSize is the number of bytes in a bitcoin message header.
// Bitcoin network (magic) 4 bytes + command 12 bytes + payload length 4 bytes +
// checksum 4 bytes.
const MessageHeaderSize = 24

// CommandSize is the fixed size of all commands in the common bitcoin message
// header.  Shorter commands must be zero padded.
const CommandSize = 15

// MaxMessagePayload is the maximum bytes a message can be regardless of other
// individual limits imposed by messages themselves.
const MaxMessagePayload = (1024 * 1024 * 32) // 32MB

// Commands used in bitcoin message headers which describe the type of message.
const (
	CmdVersion     = "version"
	CmdVerAck      = "verack"
	CmdGetAddr     = "getaddr"
	CmdAddr        = "addr"
	CmdGetBlocks   = "getblocks"
	CmdInv         = "inv"
	CmdGetData     = "getdata"
	CmdNotFound    = "notfound"
	CmdBlock       = "block" // Factoid Block
	CmdTx          = "tx"    // Factoid Tx
	CmdGetHeaders  = "getheaders"
	CmdHeaders     = "headers"
	CmdPing        = "ping"
	CmdPong        = "pong"
	CmdAlert       = "alert"
	CmdMemPool     = "mempool"
	CmdFilterAdd   = "filteradd"
	CmdFilterClear = "filterclear"
	CmdFilterLoad  = "filterload"
	CmdMerkleBlock = "merkleblock"
	CmdReject      = "reject"

	// Factom additions:

	// Factom downloads: block & entry
	CmdGetDirBlocks  = "getdirblocks"
	CmdDirInv        = "dirInv"
	CmdGetDirData    = "getdirdata"
	CmdDirBlock      = "dirblock"
	CmdGetNonDirData = "getnondirdata"
	CmdCBlock        = "ecblock"    // Entry Credit Block
	CmdEBlock        = "entryblock" // Entry Block
	CmdGetEntryData  = "getentrydata"
	CmdEntry         = "entry"

	CmdCommitChain = "commitchain"
	CmdRevealChain = "revealchain"
	CmdCommitEntry = "commitentry"
	CmdRevealEntry = "revealentry"

	CmdAcknowledgement = "confirmation"
	CmdMHashReveal     = "mhashreveal"
)

// MaxAppMsgPayload is the maximum bytes a factom app message can be in bytes.
const MaxAppMsgPayload = (10 * 1024) // 10Kib

// Message is an interface that describes a bitcoin message.  A type that
// implements Message has complete control over the representation of its data
// and may therefore contain additional or fewer fields than those which
// are used directly in the protocol encoded message.
type Message interface {
	BtcDecode(io.Reader, uint32) error
	BtcEncode(io.Writer, uint32) error
	Command() string
	MaxPayloadLength(uint32) uint32
}

// makeEmptyMessage creates a message of the appropriate concrete type based
// on the command.
func makeEmptyMessage(command string) (Message, error) {
	var msg Message
	switch command {
	case CmdVersion:
		msg = &MsgVersion{}

	case CmdVerAck:
		msg = &MsgVerAck{}

	case CmdGetAddr:
		msg = &MsgGetAddr{}

	case CmdAddr:
		msg = &MsgAddr{}

	case CmdGetBlocks:
		msg = &MsgGetBlocks{}

	case CmdBlock:
		msg = &MsgBlock{}

	case CmdInv:
		msg = &MsgInv{}

	case CmdGetData:
		msg = &MsgGetData{}

	case CmdNotFound:
		msg = &MsgNotFound{}

	case CmdTx:
		msg = &MsgTx{}

	case CmdPing:
		msg = &MsgPing{}

	case CmdPong:
		msg = &MsgPong{}

	case CmdGetHeaders:
		msg = &MsgGetHeaders{}

	case CmdHeaders:
		msg = &MsgHeaders{}

	case CmdAlert:
		msg = &MsgAlert{}

	case CmdMemPool:
		msg = &MsgMemPool{}

	case CmdFilterAdd:
		msg = &MsgFilterAdd{}

	case CmdFilterClear:
		msg = &MsgFilterClear{}

	case CmdFilterLoad:
		msg = &MsgFilterLoad{}

	case CmdMerkleBlock:
		msg = &MsgMerkleBlock{}

	case CmdReject:
		msg = &MsgReject{}

	case CmdRevealEntry:
		msg = &MsgRevealEntry{}

	case CmdCommitEntry:
		msg = &MsgCommitEntry{}

	case CmdRevealChain:
		msg = &MsgRevealChain{}

	case CmdCommitChain:
		msg = &MsgCommitChain{}

	case CmdAcknowledgement:
		msg = &MsgAcknowledgement{}

	case CmdGetDirBlocks:
		msg = &MsgGetDirBlocks{}

	case CmdDirInv:
		msg = &MsgDirInv{}

	case CmdGetDirData:
		msg = &MsgGetDirData{}

	case CmdDirBlock:
		msg = &MsgDirBlock{}

	case CmdCBlock:
		msg = &MsgCBlock{}

	case CmdEBlock:
		msg = &MsgEBlock{}

	case CmdEntry:
		msg = &MsgEntry{}

	default:
		return nil, fmt.Errorf("unhandled command [%s]", command)
	}
	return msg, nil
}

// messageHeader defines the header structure for all bitcoin protocol messages.
type messageHeader struct {
	magic    BitcoinNet // 4 bytes
	command  string     // 12 bytes
	length   uint32     // 4 bytes
	checksum [4]byte    // 4 bytes
}

// readMessageHeader reads a bitcoin message header from r.
func readMessageHeader(r io.Reader) (int, *messageHeader, error) {
	// Since readElements doesn't return the amount of bytes read, attempt
	// to read the entire header into a buffer first in case there is a
	// short read so the proper amount of read bytes are known.  This works
	// since the header is a fixed size.
	var headerBytes [MessageHeaderSize]byte
	n, err := io.ReadFull(r, headerBytes[:])
	if err != nil {
		return n, nil, err
	}
	hr := bytes.NewReader(headerBytes[:])

	// Create and populate a messageHeader struct from the raw header bytes.
	hdr := messageHeader{}
	var command [CommandSize]byte
	readElements(hr, &hdr.magic, &command, &hdr.length, &hdr.checksum)

	// Strip trailing zeros from command string.
	hdr.command = string(bytes.TrimRight(command[:], string(0)))

	return n, &hdr, nil
}

// discardInput reads n bytes from reader r in chunks and discards the read
// bytes.  This is used to skip payloads when various errors occur and helps
// prevent rogue nodes from causing massive memory allocation through forging
// header length.
func discardInput(r io.Reader, n uint32) {
	maxSize := uint32(10 * 1024) // 10k at a time
	numReads := n / maxSize
	bytesRemaining := n % maxSize
	if n > 0 {
		buf := make([]byte, maxSize)
		for i := uint32(0); i < numReads; i++ {
			io.ReadFull(r, buf)
		}
	}
	if bytesRemaining > 0 {
		buf := make([]byte, bytesRemaining)
		io.ReadFull(r, buf)
	}
}

// WriteMessageN writes a bitcoin Message to w including the necessary header
// information and returns the number of bytes written.    This function is the
// same as WriteMessage except it also returns the number of bytes written.
func WriteMessageN(w io.Writer, msg Message, pver uint32, btcnet BitcoinNet) (int, error) {
	totalBytes := 0

	// Enforce max command size.
	var command [CommandSize]byte
	cmd := msg.Command()
	if len(cmd) > CommandSize {
		str := fmt.Sprintf("command [%s] is too long [max %v]",
			cmd, CommandSize)
		return totalBytes, messageError("WriteMessage", str)
	}
	copy(command[:], []byte(cmd))

	// Encode the message payload.
	var bw bytes.Buffer
	err := msg.BtcEncode(&bw, pver)
	if err != nil {
		return totalBytes, err
	}
	payload := bw.Bytes()
	lenp := len(payload)

	// Enforce maximum overall message payload.
	if lenp > MaxMessagePayload {
		str := fmt.Sprintf("message payload is too large - encoded "+
			"%d bytes, but maximum message payload is %d bytes",
			lenp, MaxMessagePayload)
		return totalBytes, messageError("WriteMessage", str)
	}

	// Enforce maximum message payload based on the message type.
	mpl := msg.MaxPayloadLength(pver)
	if uint32(lenp) > mpl {
		str := fmt.Sprintf("message payload is too large - encoded "+
			"%d bytes, but maximum message payload size for "+
			"messages of type [%s] is %d.", lenp, cmd, mpl)
		return totalBytes, messageError("WriteMessage", str)
	}

	// Create header for the message.
	hdr := messageHeader{}
	hdr.magic = btcnet
	hdr.command = cmd
	hdr.length = uint32(lenp)
	copy(hdr.checksum[:], DoubleSha256(payload)[0:4])

	// Encode the header for the message.  This is done to a buffer
	// rather than directly to the writer since writeElements doesn't
	// return the number of bytes written.
	hw := bytes.NewBuffer(make([]byte, 0, MessageHeaderSize))
	writeElements(hw, hdr.magic, command, hdr.length, hdr.checksum)

	// Write header.
	n, err := w.Write(hw.Bytes())
	if err != nil {
		totalBytes += n
		return totalBytes, err
	}
	totalBytes += n

	// Write payload.
	n, err = w.Write(payload)
	if err != nil {
		totalBytes += n
		return totalBytes, err
	}
	totalBytes += n

	return totalBytes, nil
}

// WriteMessage writes a bitcoin Message to w including the necessary header
// information.  This function is the same as WriteMessageN except it doesn't
// doesn't return the number of bytes written.  This function is mainly provided
// for backwards compatibility with the original API, but it's also useful for
// callers that don't care about byte counts.
func WriteMessage(w io.Writer, msg Message, pver uint32, btcnet BitcoinNet) error {
	_, err := WriteMessageN(w, msg, pver, btcnet)
	return err
}

// ReadMessageN reads, validates, and parses the next bitcoin Message from r for
// the provided protocol version and bitcoin network.  It returns the number of
// bytes read in addition to the parsed Message and raw bytes which comprise the
// message.  This function is the same as ReadMessage except it also returns the
// number of bytes read.
func ReadMessageN(r io.Reader, pver uint32, btcnet BitcoinNet) (int, Message, []byte, error) {
	totalBytes := 0
	n, hdr, err := readMessageHeader(r)
	if err != nil {
		totalBytes += n
		return totalBytes, nil, nil, err
	}
	totalBytes += n

	// Enforce maximum message payload.
	if hdr.length > MaxMessagePayload {
		str := fmt.Sprintf("message payload is too large - header "+
			"indicates %d bytes, but max message payload is %d "+
			"bytes.", hdr.length, MaxMessagePayload)
		return totalBytes, nil, nil, messageError("ReadMessage", str)

	}

	// Check for messages from the wrong bitcoin network.
	if hdr.magic != btcnet {
		discardInput(r, hdr.length)
		str := fmt.Sprintf("message from other network [%v]", hdr.magic)
		return totalBytes, nil, nil, messageError("ReadMessage", str)
	}

	// Check for malformed commands.
	command := hdr.command
	if !utf8.ValidString(command) {
		discardInput(r, hdr.length)
		str := fmt.Sprintf("invalid command %v", []byte(command))
		return totalBytes, nil, nil, messageError("ReadMessage", str)
	}

	// Create struct of appropriate message type based on the command.
	msg, err := makeEmptyMessage(command)
	if err != nil {
		discardInput(r, hdr.length)
		return totalBytes, nil, nil, messageError("ReadMessage",
			err.Error())
	}

	// Check for maximum length based on the message type as a malicious client
	// could otherwise create a well-formed header and set the length to max
	// numbers in order to exhaust the machine's memory.
	mpl := msg.MaxPayloadLength(pver)
	if hdr.length > mpl {
		discardInput(r, hdr.length)
		str := fmt.Sprintf("payload exceeds max length - header "+
			"indicates %v bytes, but max payload size for "+
			"messages of type [%v] is %v.", hdr.length, command, mpl)
		return totalBytes, nil, nil, messageError("ReadMessage", str)
	}

	// Read payload.
	payload := make([]byte, hdr.length)
	n, err = io.ReadFull(r, payload)
	if err != nil {
		totalBytes += n
		return totalBytes, nil, nil, err
	}
	totalBytes += n

	// Test checksum.
	checksum := DoubleSha256(payload)[0:4]
	if !bytes.Equal(checksum[:], hdr.checksum[:]) {
		str := fmt.Sprintf("payload checksum failed - header "+
			"indicates %v, but actual checksum is %v.",
			hdr.checksum, checksum)
		return totalBytes, nil, nil, messageError("ReadMessage", str)
	}

	// Unmarshal message.  NOTE: This must be a *bytes.Buffer since the
	// MsgVersion BtcDecode function requires it.
	pr := bytes.NewBuffer(payload)
	err = msg.BtcDecode(pr, pver)
	if err != nil {
		return totalBytes, nil, nil, err
	}

	return totalBytes, msg, payload, nil
}

// ReadMessage reads, validates, and parses the next bitcoin Message from r for
// the provided protocol version and bitcoin network.  It returns the parsed
// Message and raw bytes which comprise the message.  This function only differs
// from ReadMessageN in that it doesn't return the number of bytes read.  This
// function is mainly provided for backwards compatibility with the original
// API, but it's also useful for callers that don't care about byte counts.
func ReadMessage(r io.Reader, pver uint32, btcnet BitcoinNet) (Message, []byte, error) {
	_, msg, buf, err := ReadMessageN(r, pver, btcnet)
	return msg, buf, err
}
<commit_msg>add CmdGetNonDirData & CmdGetEntryData<commit_after>// Copyright (c) 2013-2015 Conformal Systems LLC.
// Use of this source code is governed by an ISC
// license that can be found in the LICENSE file.

package wire

import (
	"bytes"
	"fmt"
	"io"
	"unicode/utf8"
)

// MessageHeaderSize is the number of bytes in a bitcoin message header.
// Bitcoin network (magic) 4 bytes + command 12 bytes + payload length 4 bytes +
// checksum 4 bytes.
const MessageHeaderSize = 24

// CommandSize is the fixed size of all commands in the common bitcoin message
// header.  Shorter commands must be zero padded.
const CommandSize = 12

// MaxMessagePayload is the maximum bytes a message can be regardless of other
// individual limits imposed by messages themselves.
const MaxMessagePayload = (1024 * 1024 * 32) // 32MB

// Commands used in bitcoin message headers which describe the type of message.
const (
	CmdVersion     = "version"
	CmdVerAck      = "verack"
	CmdGetAddr     = "getaddr"
	CmdAddr        = "addr"
	CmdGetBlocks   = "getblocks"
	CmdInv         = "inv"
	CmdGetData     = "getdata"
	CmdNotFound    = "notfound"
	CmdBlock       = "block" // Factoid Block
	CmdTx          = "tx"    // Factoid Tx
	CmdGetHeaders  = "getheaders"
	CmdHeaders     = "headers"
	CmdPing        = "ping"
	CmdPong        = "pong"
	CmdAlert       = "alert"
	CmdMemPool     = "mempool"
	CmdFilterAdd   = "filteradd"
	CmdFilterClear = "filterclear"
	CmdFilterLoad  = "filterload"
	CmdMerkleBlock = "merkleblock"
	CmdReject      = "reject"

	// Factom additions:

	// Factom downloads: block & entry
	CmdGetDirBlocks  = "getdirblocks"
	CmdDirInv        = "dirInv"
	CmdGetDirData    = "getdirdata"
	CmdDirBlock      = "dirblock"
	CmdGetNonDirData = "getnodirdata"
	CmdCBlock        = "ecblock"    // Entry Credit Block
	CmdEBlock        = "entryblock" // Entry Block
	CmdGetEntryData  = "getentrydata"
	CmdEntry         = "entry"

	CmdCommitChain = "commitchain"
	CmdRevealChain = "revealchain"
	CmdCommitEntry = "commitentry"
	CmdRevealEntry = "revealentry"

	CmdAcknowledgement = "confirmation"
	CmdMHashReveal     = "mhashreveal"
)

// MaxAppMsgPayload is the maximum bytes a factom app message can be in bytes.
const MaxAppMsgPayload = (10 * 1024) // 10Kib

// Message is an interface that describes a bitcoin message.  A type that
// implements Message has complete control over the representation of its data
// and may therefore contain additional or fewer fields than those which
// are used directly in the protocol encoded message.
type Message interface {
	BtcDecode(io.Reader, uint32) error
	BtcEncode(io.Writer, uint32) error
	Command() string
	MaxPayloadLength(uint32) uint32
}

// makeEmptyMessage creates a message of the appropriate concrete type based
// on the command.
func makeEmptyMessage(command string) (Message, error) {
	var msg Message
	switch command {
	case CmdVersion:
		msg = &MsgVersion{}

	case CmdVerAck:
		msg = &MsgVerAck{}

	case CmdGetAddr:
		msg = &MsgGetAddr{}

	case CmdAddr:
		msg = &MsgAddr{}

	case CmdGetBlocks:
		msg = &MsgGetBlocks{}

	case CmdBlock:
		msg = &MsgBlock{}

	case CmdInv:
		msg = &MsgInv{}

	case CmdGetData:
		msg = &MsgGetData{}

	case CmdNotFound:
		msg = &MsgNotFound{}

	case CmdTx:
		msg = &MsgTx{}

	case CmdPing:
		msg = &MsgPing{}

	case CmdPong:
		msg = &MsgPong{}

	case CmdGetHeaders:
		msg = &MsgGetHeaders{}

	case CmdHeaders:
		msg = &MsgHeaders{}

	case CmdAlert:
		msg = &MsgAlert{}

	case CmdMemPool:
		msg = &MsgMemPool{}

	case CmdFilterAdd:
		msg = &MsgFilterAdd{}

	case CmdFilterClear:
		msg = &MsgFilterClear{}

	case CmdFilterLoad:
		msg = &MsgFilterLoad{}

	case CmdMerkleBlock:
		msg = &MsgMerkleBlock{}

	case CmdReject:
		msg = &MsgReject{}

	case CmdRevealEntry:
		msg = &MsgRevealEntry{}

	case CmdCommitEntry:
		msg = &MsgCommitEntry{}

	case CmdRevealChain:
		msg = &MsgRevealChain{}

	case CmdCommitChain:
		msg = &MsgCommitChain{}

	case CmdAcknowledgement:
		msg = &MsgAcknowledgement{}

	case CmdGetDirBlocks:
		msg = &MsgGetDirBlocks{}

	case CmdDirInv:
		msg = &MsgDirInv{}

	case CmdGetDirData:
		msg = &MsgGetDirData{}

	case CmdDirBlock:
		msg = &MsgDirBlock{}

	case CmdGetNonDirData:
		msg = &MsgGetNonDirData{}

	case CmdCBlock:
		msg = &MsgCBlock{}

	case CmdEBlock:
		msg = &MsgEBlock{}

	case CmdGetEntryData:
		msg = &MsgGetEntryData{}

	case CmdEntry:
		msg = &MsgEntry{}

	default:
		return nil, fmt.Errorf("unhandled command [%s]", command)
	}
	return msg, nil
}

// messageHeader defines the header structure for all bitcoin protocol messages.
type messageHeader struct {
	magic    BitcoinNet // 4 bytes
	command  string     // 12 bytes
	length   uint32     // 4 bytes
	checksum [4]byte    // 4 bytes
}

// readMessageHeader reads a bitcoin message header from r.
func readMessageHeader(r io.Reader) (int, *messageHeader, error) {
	// Since readElements doesn't return the amount of bytes read, attempt
	// to read the entire header into a buffer first in case there is a
	// short read so the proper amount of read bytes are known.  This works
	// since the header is a fixed size.
	var headerBytes [MessageHeaderSize]byte
	n, err := io.ReadFull(r, headerBytes[:])
	if err != nil {
		return n, nil, err
	}
	hr := bytes.NewReader(headerBytes[:])

	// Create and populate a messageHeader struct from the raw header bytes.
	hdr := messageHeader{}
	var command [CommandSize]byte
	readElements(hr, &hdr.magic, &command, &hdr.length, &hdr.checksum)

	// Strip trailing zeros from command string.
	hdr.command = string(bytes.TrimRight(command[:], string(0)))

	return n, &hdr, nil
}

// discardInput reads n bytes from reader r in chunks and discards the read
// bytes.  This is used to skip payloads when various errors occur and helps
// prevent rogue nodes from causing massive memory allocation through forging
// header length.
func discardInput(r io.Reader, n uint32) {
	maxSize := uint32(10 * 1024) // 10k at a time
	numReads := n / maxSize
	bytesRemaining := n % maxSize
	if n > 0 {
		buf := make([]byte, maxSize)
		for i := uint32(0); i < numReads; i++ {
			io.ReadFull(r, buf)
		}
	}
	if bytesRemaining > 0 {
		buf := make([]byte, bytesRemaining)
		io.ReadFull(r, buf)
	}
}

// WriteMessageN writes a bitcoin Message to w including the necessary header
// information and returns the number of bytes written.    This function is the
// same as WriteMessage except it also returns the number of bytes written.
func WriteMessageN(w io.Writer, msg Message, pver uint32, btcnet BitcoinNet) (int, error) {
	totalBytes := 0

	// Enforce max command size.
	var command [CommandSize]byte
	cmd := msg.Command()
	if len(cmd) > CommandSize {
		str := fmt.Sprintf("command [%s] is too long [max %v]",
			cmd, CommandSize)
		return totalBytes, messageError("WriteMessage", str)
	}
	copy(command[:], []byte(cmd))

	// Encode the message payload.
	var bw bytes.Buffer
	err := msg.BtcEncode(&bw, pver)
	if err != nil {
		return totalBytes, err
	}
	payload := bw.Bytes()
	lenp := len(payload)

	// Enforce maximum overall message payload.
	if lenp > MaxMessagePayload {
		str := fmt.Sprintf("message payload is too large - encoded "+
			"%d bytes, but maximum message payload is %d bytes",
			lenp, MaxMessagePayload)
		return totalBytes, messageError("WriteMessage", str)
	}

	// Enforce maximum message payload based on the message type.
	mpl := msg.MaxPayloadLength(pver)
	if uint32(lenp) > mpl {
		str := fmt.Sprintf("message payload is too large - encoded "+
			"%d bytes, but maximum message payload size for "+
			"messages of type [%s] is %d.", lenp, cmd, mpl)
		return totalBytes, messageError("WriteMessage", str)
	}

	// Create header for the message.
	hdr := messageHeader{}
	hdr.magic = btcnet
	hdr.command = cmd
	hdr.length = uint32(lenp)
	copy(hdr.checksum[:], DoubleSha256(payload)[0:4])

	// Encode the header for the message.  This is done to a buffer
	// rather than directly to the writer since writeElements doesn't
	// return the number of bytes written.
	hw := bytes.NewBuffer(make([]byte, 0, MessageHeaderSize))
	writeElements(hw, hdr.magic, command, hdr.length, hdr.checksum)

	// Write header.
	n, err := w.Write(hw.Bytes())
	if err != nil {
		totalBytes += n
		return totalBytes, err
	}
	totalBytes += n

	// Write payload.
	n, err = w.Write(payload)
	if err != nil {
		totalBytes += n
		return totalBytes, err
	}
	totalBytes += n

	return totalBytes, nil
}

// WriteMessage writes a bitcoin Message to w including the necessary header
// information.  This function is the same as WriteMessageN except it doesn't
// doesn't return the number of bytes written.  This function is mainly provided
// for backwards compatibility with the original API, but it's also useful for
// callers that don't care about byte counts.
func WriteMessage(w io.Writer, msg Message, pver uint32, btcnet BitcoinNet) error {
	_, err := WriteMessageN(w, msg, pver, btcnet)
	return err
}

// ReadMessageN reads, validates, and parses the next bitcoin Message from r for
// the provided protocol version and bitcoin network.  It returns the number of
// bytes read in addition to the parsed Message and raw bytes which comprise the
// message.  This function is the same as ReadMessage except it also returns the
// number of bytes read.
func ReadMessageN(r io.Reader, pver uint32, btcnet BitcoinNet) (int, Message, []byte, error) {
	totalBytes := 0
	n, hdr, err := readMessageHeader(r)
	if err != nil {
		totalBytes += n
		return totalBytes, nil, nil, err
	}
	totalBytes += n

	// Enforce maximum message payload.
	if hdr.length > MaxMessagePayload {
		str := fmt.Sprintf("message payload is too large - header "+
			"indicates %d bytes, but max message payload is %d "+
			"bytes.", hdr.length, MaxMessagePayload)
		return totalBytes, nil, nil, messageError("ReadMessage", str)

	}

	// Check for messages from the wrong bitcoin network.
	if hdr.magic != btcnet {
		discardInput(r, hdr.length)
		str := fmt.Sprintf("message from other network [%v]", hdr.magic)
		return totalBytes, nil, nil, messageError("ReadMessage", str)
	}

	// Check for malformed commands.
	command := hdr.command
	if !utf8.ValidString(command) {
		discardInput(r, hdr.length)
		str := fmt.Sprintf("invalid command %v", []byte(command))
		return totalBytes, nil, nil, messageError("ReadMessage", str)
	}

	// Create struct of appropriate message type based on the command.
	msg, err := makeEmptyMessage(command)
	if err != nil {
		discardInput(r, hdr.length)
		return totalBytes, nil, nil, messageError("ReadMessage",
			err.Error())
	}

	// Check for maximum length based on the message type as a malicious client
	// could otherwise create a well-formed header and set the length to max
	// numbers in order to exhaust the machine's memory.
	mpl := msg.MaxPayloadLength(pver)
	if hdr.length > mpl {
		discardInput(r, hdr.length)
		str := fmt.Sprintf("payload exceeds max length - header "+
			"indicates %v bytes, but max payload size for "+
			"messages of type [%v] is %v.", hdr.length, command, mpl)
		return totalBytes, nil, nil, messageError("ReadMessage", str)
	}

	// Read payload.
	payload := make([]byte, hdr.length)
	n, err = io.ReadFull(r, payload)
	if err != nil {
		totalBytes += n
		return totalBytes, nil, nil, err
	}
	totalBytes += n

	// Test checksum.
	checksum := DoubleSha256(payload)[0:4]
	if !bytes.Equal(checksum[:], hdr.checksum[:]) {
		str := fmt.Sprintf("payload checksum failed - header "+
			"indicates %v, but actual checksum is %v.",
			hdr.checksum, checksum)
		return totalBytes, nil, nil, messageError("ReadMessage", str)
	}

	// Unmarshal message.  NOTE: This must be a *bytes.Buffer since the
	// MsgVersion BtcDecode function requires it.
	pr := bytes.NewBuffer(payload)
	err = msg.BtcDecode(pr, pver)
	if err != nil {
		return totalBytes, nil, nil, err
	}

	return totalBytes, msg, payload, nil
}

// ReadMessage reads, validates, and parses the next bitcoin Message from r for
// the provided protocol version and bitcoin network.  It returns the parsed
// Message and raw bytes which comprise the message.  This function only differs
// from ReadMessageN in that it doesn't return the number of bytes read.  This
// function is mainly provided for backwards compatibility with the original
// API, but it's also useful for callers that don't care about byte counts.
func ReadMessage(r io.Reader, pver uint32, btcnet BitcoinNet) (Message, []byte, error) {
	_, msg, buf, err := ReadMessageN(r, pver, btcnet)
	return msg, buf, err
}
<|endoftext|>


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

package wire

import (
	"bytes"
	"fmt"
	"io"
	"unicode/utf8"

	"github.com/btcsuite/btcd/chaincfg/chainhash"
)

// MessageHeaderSize is the number of bytes in a bitcoin message header.
// Bitcoin network (magic) 4 bytes + command 12 bytes + payload length 4 bytes +
// checksum 4 bytes.
const MessageHeaderSize = 24

// CommandSize is the fixed size of all commands in the common bitcoin message
// header.  Shorter commands must be zero padded.
const CommandSize = 12

// MaxMessagePayload is the maximum bytes a message can be regardless of other
// individual limits imposed by messages themselves.
const MaxMessagePayload = (1024 * 1024 * 32) // 32MB

// Commands used in bitcoin message headers which describe the type of message.
const (
	CmdVersion      = "version"
	CmdVerAck       = "verack"
	CmdGetAddr      = "getaddr"
	CmdAddr         = "addr"
	CmdGetBlocks    = "getblocks"
	CmdInv          = "inv"
	CmdGetData      = "getdata"
	CmdNotFound     = "notfound"
	CmdBlock        = "block"
	CmdTx           = "tx"
	CmdGetHeaders   = "getheaders"
	CmdHeaders      = "headers"
	CmdPing         = "ping"
	CmdPong         = "pong"
	CmdAlert        = "alert"
	CmdMemPool      = "mempool"
	CmdFilterAdd    = "filteradd"
	CmdFilterClear  = "filterclear"
	CmdFilterLoad   = "filterload"
	CmdMerkleBlock  = "merkleblock"
	CmdReject       = "reject"
	CmdSendHeaders  = "sendheaders"
	CmdFeeFilter    = "feefilter"
	CmdGetCFilters  = "getcfilters"
	CmdGetCFHeaders = "getcfheaders"
	CmdGetCFCheckpt = "getcfcheckpt"
	CmdCFilter      = "cfilter"
	CmdCFHeaders    = "cfheaders"
	CmdCFCheckpt    = "cfcheckpt"
)

// MessageEncoding represents the wire message encoding format to be used.
type MessageEncoding uint32

const (
	// BaseEncoding encodes all messages in the default format specified
	// for the Bitcoin wire protocol.
	BaseEncoding MessageEncoding = 1 << iota

	// WitnessEncoding encodes all messages other than transaction messages
	// using the default Bitcoin wire protocol specification. For transaction
	// messages, the new encoding format detailed in BIP0144 will be used.
	WitnessEncoding
)

// LatestEncoding is the most recently specified encoding for the Bitcoin wire
// protocol.
var LatestEncoding = WitnessEncoding

// Message is an interface that describes a bitcoin message.  A type that
// implements Message has complete control over the representation of its data
// and may therefore contain additional or fewer fields than those which
// are used directly in the protocol encoded message.
type Message interface {
	BtcDecode(io.Reader, uint32, MessageEncoding) error
	BtcEncode(io.Writer, uint32, MessageEncoding) error
	Command() string
	MaxPayloadLength(uint32) uint32
}

// makeEmptyMessage creates a message of the appropriate concrete type based
// on the command.
func makeEmptyMessage(command string) (Message, error) {
	var msg Message
	switch command {
	case CmdVersion:
		msg = &MsgVersion{}

	case CmdVerAck:
		msg = &MsgVerAck{}

	case CmdGetAddr:
		msg = &MsgGetAddr{}

	case CmdAddr:
		msg = &MsgAddr{}

	case CmdGetBlocks:
		msg = &MsgGetBlocks{}

	case CmdBlock:
		msg = &MsgBlock{}

	case CmdInv:
		msg = &MsgInv{}

	case CmdGetData:
		msg = &MsgGetData{}

	case CmdNotFound:
		msg = &MsgNotFound{}

	case CmdTx:
		msg = &MsgTx{}

	case CmdPing:
		msg = &MsgPing{}

	case CmdPong:
		msg = &MsgPong{}

	case CmdGetHeaders:
		msg = &MsgGetHeaders{}

	case CmdHeaders:
		msg = &MsgHeaders{}

	case CmdAlert:
		msg = &MsgAlert{}

	case CmdMemPool:
		msg = &MsgMemPool{}

	case CmdFilterAdd:
		msg = &MsgFilterAdd{}

	case CmdFilterClear:
		msg = &MsgFilterClear{}

	case CmdFilterLoad:
		msg = &MsgFilterLoad{}

	case CmdMerkleBlock:
		msg = &MsgMerkleBlock{}

	case CmdReject:
		msg = &MsgReject{}

	case CmdSendHeaders:
		msg = &MsgSendHeaders{}

	case CmdFeeFilter:
		msg = &MsgFeeFilter{}

	case CmdGetCFilters:
		msg = &MsgGetCFilters{}

	case CmdGetCFHeaders:
		msg = &MsgGetCFHeaders{}

	case CmdGetCFCheckpt:
		msg = &MsgGetCFCheckpt{}

	case CmdCFilter:
		msg = &MsgCFilter{}

	case CmdCFHeaders:
		msg = &MsgCFHeaders{}

	case CmdCFCheckpt:
		msg = &MsgCFCheckpt{}

	default:
		return nil, fmt.Errorf("unhandled command [%s]", command)
	}
	return msg, nil
}

// messageHeader defines the header structure for all bitcoin protocol messages.
type messageHeader struct {
	magic    BitcoinNet // 4 bytes
	command  string     // 12 bytes
	length   uint32     // 4 bytes
	checksum [4]byte    // 4 bytes
}

// readMessageHeader reads a bitcoin message header from r.
func readMessageHeader(r io.Reader) (int, *messageHeader, error) {
	// Since readElements doesn't return the amount of bytes read, attempt
	// to read the entire header into a buffer first in case there is a
	// short read so the proper amount of read bytes are known.  This works
	// since the header is a fixed size.
	var headerBytes [MessageHeaderSize]byte
	n, err := io.ReadFull(r, headerBytes[:])
	if err != nil {
		return n, nil, err
	}
	hr := bytes.NewReader(headerBytes[:])

	// Create and populate a messageHeader struct from the raw header bytes.
	hdr := messageHeader{}
	var command [CommandSize]byte
	readElements(hr, &hdr.magic, &command, &hdr.length, &hdr.checksum)

	// Strip trailing zeros from command string.
	hdr.command = string(bytes.TrimRight(command[:], string(0)))

	return n, &hdr, nil
}

// discardInput reads n bytes from reader r in chunks and discards the read
// bytes.  This is used to skip payloads when various errors occur and helps
// prevent rogue nodes from causing massive memory allocation through forging
// header length.
func discardInput(r io.Reader, n uint32) {
	maxSize := uint32(10 * 1024) // 10k at a time
	numReads := n / maxSize
	bytesRemaining := n % maxSize
	if n > 0 {
		buf := make([]byte, maxSize)
		for i := uint32(0); i < numReads; i++ {
			io.ReadFull(r, buf)
		}
	}
	if bytesRemaining > 0 {
		buf := make([]byte, bytesRemaining)
		io.ReadFull(r, buf)
	}
}

// WriteMessageN writes a bitcoin Message to w including the necessary header
// information and returns the number of bytes written.    This function is the
// same as WriteMessage except it also returns the number of bytes written.
func WriteMessageN(w io.Writer, msg Message, pver uint32, btcnet BitcoinNet) (int, error) {
	return WriteMessageWithEncodingN(w, msg, pver, btcnet, BaseEncoding)
}

// WriteMessage writes a bitcoin Message to w including the necessary header
// information.  This function is the same as WriteMessageN except it doesn't
// doesn't return the number of bytes written.  This function is mainly provided
// for backwards compatibility with the original API, but it's also useful for
// callers that don't care about byte counts.
func WriteMessage(w io.Writer, msg Message, pver uint32, btcnet BitcoinNet) error {
	_, err := WriteMessageN(w, msg, pver, btcnet)
	return err
}

// WriteMessageWithEncodingN writes a bitcoin Message to w including the
// necessary header information and returns the number of bytes written.
// This function is the same as WriteMessageN except it also allows the caller
// to specify the message encoding format to be used when serializing wire
// messages.
func WriteMessageWithEncodingN(w io.Writer, msg Message, pver uint32,
	btcnet BitcoinNet, encoding MessageEncoding) (int, error) {

	totalBytes := 0

	// Enforce max command size.
	var command [CommandSize]byte
	cmd := msg.Command()
	if len(cmd) > CommandSize {
		str := fmt.Sprintf("command [%s] is too long [max %v]",
			cmd, CommandSize)
		return totalBytes, messageError("WriteMessage", str)
	}
	copy(command[:], []byte(cmd))

	// Encode the message payload.
	var bw bytes.Buffer
	err := msg.BtcEncode(&bw, pver, encoding)
	if err != nil {
		return totalBytes, err
	}
	payload := bw.Bytes()
	lenp := len(payload)

	// Enforce maximum overall message payload.
	if lenp > MaxMessagePayload {
		str := fmt.Sprintf("message payload is too large - encoded "+
			"%d bytes, but maximum message payload is %d bytes",
			lenp, MaxMessagePayload)
		return totalBytes, messageError("WriteMessage", str)
	}

	// Enforce maximum message payload based on the message type.
	mpl := msg.MaxPayloadLength(pver)
	if uint32(lenp) > mpl {
		str := fmt.Sprintf("message payload is too large - encoded "+
			"%d bytes, but maximum message payload size for "+
			"messages of type [%s] is %d.", lenp, cmd, mpl)
		return totalBytes, messageError("WriteMessage", str)
	}

	// Create header for the message.
	hdr := messageHeader{}
	hdr.magic = btcnet
	hdr.command = cmd
	hdr.length = uint32(lenp)
	copy(hdr.checksum[:], chainhash.DoubleHashB(payload)[0:4])

	// Encode the header for the message.  This is done to a buffer
	// rather than directly to the writer since writeElements doesn't
	// return the number of bytes written.
	hw := bytes.NewBuffer(make([]byte, 0, MessageHeaderSize))
	writeElements(hw, hdr.magic, command, hdr.length, hdr.checksum)

	// Write header.
	n, err := w.Write(hw.Bytes())
	totalBytes += n
	if err != nil {
		return totalBytes, err
	}

	// Write payload.
	n, err = w.Write(payload)
	totalBytes += n
	return totalBytes, err
}

// ReadMessageWithEncodingN reads, validates, and parses the next bitcoin Message
// from r for the provided protocol version and bitcoin network.  It returns the
// number of bytes read in addition to the parsed Message and raw bytes which
// comprise the message.  This function is the same as ReadMessageN except it
// allows the caller to specify which message encoding is to to consult when
// decoding wire messages.
func ReadMessageWithEncodingN(r io.Reader, pver uint32, btcnet BitcoinNet,
	enc MessageEncoding) (int, Message, []byte, error) {

	totalBytes := 0
	n, hdr, err := readMessageHeader(r)
	totalBytes += n
	if err != nil {
		return totalBytes, nil, nil, err
	}

	// Enforce maximum message payload.
	if hdr.length > MaxMessagePayload {
		str := fmt.Sprintf("message payload is too large - header "+
			"indicates %d bytes, but max message payload is %d "+
			"bytes.", hdr.length, MaxMessagePayload)
		return totalBytes, nil, nil, messageError("ReadMessage", str)

	}

	// Check for messages from the wrong bitcoin network.
	if hdr.magic != btcnet {
		discardInput(r, hdr.length)
		str := fmt.Sprintf("message from other network [%v]", hdr.magic)
		return totalBytes, nil, nil, messageError("ReadMessage", str)
	}

	// Check for malformed commands.
	command := hdr.command
	if !utf8.ValidString(command) {
		discardInput(r, hdr.length)
		str := fmt.Sprintf("invalid command %v", []byte(command))
		return totalBytes, nil, nil, messageError("ReadMessage", str)
	}

	// Create struct of appropriate message type based on the command.
	msg, err := makeEmptyMessage(command)
	if err != nil {
		discardInput(r, hdr.length)
		return totalBytes, nil, nil, messageError("ReadMessage",
			err.Error())
	}

	// Check for maximum length based on the message type as a malicious client
	// could otherwise create a well-formed header and set the length to max
	// numbers in order to exhaust the machine's memory.
	mpl := msg.MaxPayloadLength(pver)
	if hdr.length > mpl {
		discardInput(r, hdr.length)
		str := fmt.Sprintf("payload exceeds max length - header "+
			"indicates %v bytes, but max payload size for "+
			"messages of type [%v] is %v.", hdr.length, command, mpl)
		return totalBytes, nil, nil, messageError("ReadMessage", str)
	}

	// Read payload.
	payload := make([]byte, hdr.length)
	n, err = io.ReadFull(r, payload)
	totalBytes += n
	if err != nil {
		return totalBytes, nil, nil, err
	}

	// Test checksum.
	checksum := chainhash.DoubleHashB(payload)[0:4]
	if !bytes.Equal(checksum[:], hdr.checksum[:]) {
		str := fmt.Sprintf("payload checksum failed - header "+
			"indicates %v, but actual checksum is %v.",
			hdr.checksum, checksum)
		return totalBytes, nil, nil, messageError("ReadMessage", str)
	}

	// Unmarshal message.  NOTE: This must be a *bytes.Buffer since the
	// MsgVersion BtcDecode function requires it.
	pr := bytes.NewBuffer(payload)
	err = msg.BtcDecode(pr, pver, enc)
	if err != nil {
		return totalBytes, nil, nil, err
	}

	return totalBytes, msg, payload, nil
}

// ReadMessageN reads, validates, and parses the next bitcoin Message from r for
// the provided protocol version and bitcoin network.  It returns the number of
// bytes read in addition to the parsed Message and raw bytes which comprise the
// message.  This function is the same as ReadMessage except it also returns the
// number of bytes read.
func ReadMessageN(r io.Reader, pver uint32, btcnet BitcoinNet) (int, Message, []byte, error) {
	return ReadMessageWithEncodingN(r, pver, btcnet, BaseEncoding)
}

// ReadMessage reads, validates, and parses the next bitcoin Message from r for
// the provided protocol version and bitcoin network.  It returns the parsed
// Message and raw bytes which comprise the message.  This function only differs
// from ReadMessageN in that it doesn't return the number of bytes read.  This
// function is mainly provided for backwards compatibility with the original
// API, but it's also useful for callers that don't care about byte counts.
func ReadMessage(r io.Reader, pver uint32, btcnet BitcoinNet) (Message, []byte, error) {
	_, msg, buf, err := ReadMessageN(r, pver, btcnet)
	return msg, buf, err
}
<commit_msg>wire: only write message payload if there actually is one<commit_after>// Copyright (c) 2013-2016 The btcsuite developers
// Use of this source code is governed by an ISC
// license that can be found in the LICENSE file.

package wire

import (
	"bytes"
	"fmt"
	"io"
	"unicode/utf8"

	"github.com/btcsuite/btcd/chaincfg/chainhash"
)

// MessageHeaderSize is the number of bytes in a bitcoin message header.
// Bitcoin network (magic) 4 bytes + command 12 bytes + payload length 4 bytes +
// checksum 4 bytes.
const MessageHeaderSize = 24

// CommandSize is the fixed size of all commands in the common bitcoin message
// header.  Shorter commands must be zero padded.
const CommandSize = 12

// MaxMessagePayload is the maximum bytes a message can be regardless of other
// individual limits imposed by messages themselves.
const MaxMessagePayload = (1024 * 1024 * 32) // 32MB

// Commands used in bitcoin message headers which describe the type of message.
const (
	CmdVersion      = "version"
	CmdVerAck       = "verack"
	CmdGetAddr      = "getaddr"
	CmdAddr         = "addr"
	CmdGetBlocks    = "getblocks"
	CmdInv          = "inv"
	CmdGetData      = "getdata"
	CmdNotFound     = "notfound"
	CmdBlock        = "block"
	CmdTx           = "tx"
	CmdGetHeaders   = "getheaders"
	CmdHeaders      = "headers"
	CmdPing         = "ping"
	CmdPong         = "pong"
	CmdAlert        = "alert"
	CmdMemPool      = "mempool"
	CmdFilterAdd    = "filteradd"
	CmdFilterClear  = "filterclear"
	CmdFilterLoad   = "filterload"
	CmdMerkleBlock  = "merkleblock"
	CmdReject       = "reject"
	CmdSendHeaders  = "sendheaders"
	CmdFeeFilter    = "feefilter"
	CmdGetCFilters  = "getcfilters"
	CmdGetCFHeaders = "getcfheaders"
	CmdGetCFCheckpt = "getcfcheckpt"
	CmdCFilter      = "cfilter"
	CmdCFHeaders    = "cfheaders"
	CmdCFCheckpt    = "cfcheckpt"
)

// MessageEncoding represents the wire message encoding format to be used.
type MessageEncoding uint32

const (
	// BaseEncoding encodes all messages in the default format specified
	// for the Bitcoin wire protocol.
	BaseEncoding MessageEncoding = 1 << iota

	// WitnessEncoding encodes all messages other than transaction messages
	// using the default Bitcoin wire protocol specification. For transaction
	// messages, the new encoding format detailed in BIP0144 will be used.
	WitnessEncoding
)

// LatestEncoding is the most recently specified encoding for the Bitcoin wire
// protocol.
var LatestEncoding = WitnessEncoding

// Message is an interface that describes a bitcoin message.  A type that
// implements Message has complete control over the representation of its data
// and may therefore contain additional or fewer fields than those which
// are used directly in the protocol encoded message.
type Message interface {
	BtcDecode(io.Reader, uint32, MessageEncoding) error
	BtcEncode(io.Writer, uint32, MessageEncoding) error
	Command() string
	MaxPayloadLength(uint32) uint32
}

// makeEmptyMessage creates a message of the appropriate concrete type based
// on the command.
func makeEmptyMessage(command string) (Message, error) {
	var msg Message
	switch command {
	case CmdVersion:
		msg = &MsgVersion{}

	case CmdVerAck:
		msg = &MsgVerAck{}

	case CmdGetAddr:
		msg = &MsgGetAddr{}

	case CmdAddr:
		msg = &MsgAddr{}

	case CmdGetBlocks:
		msg = &MsgGetBlocks{}

	case CmdBlock:
		msg = &MsgBlock{}

	case CmdInv:
		msg = &MsgInv{}

	case CmdGetData:
		msg = &MsgGetData{}

	case CmdNotFound:
		msg = &MsgNotFound{}

	case CmdTx:
		msg = &MsgTx{}

	case CmdPing:
		msg = &MsgPing{}

	case CmdPong:
		msg = &MsgPong{}

	case CmdGetHeaders:
		msg = &MsgGetHeaders{}

	case CmdHeaders:
		msg = &MsgHeaders{}

	case CmdAlert:
		msg = &MsgAlert{}

	case CmdMemPool:
		msg = &MsgMemPool{}

	case CmdFilterAdd:
		msg = &MsgFilterAdd{}

	case CmdFilterClear:
		msg = &MsgFilterClear{}

	case CmdFilterLoad:
		msg = &MsgFilterLoad{}

	case CmdMerkleBlock:
		msg = &MsgMerkleBlock{}

	case CmdReject:
		msg = &MsgReject{}

	case CmdSendHeaders:
		msg = &MsgSendHeaders{}

	case CmdFeeFilter:
		msg = &MsgFeeFilter{}

	case CmdGetCFilters:
		msg = &MsgGetCFilters{}

	case CmdGetCFHeaders:
		msg = &MsgGetCFHeaders{}

	case CmdGetCFCheckpt:
		msg = &MsgGetCFCheckpt{}

	case CmdCFilter:
		msg = &MsgCFilter{}

	case CmdCFHeaders:
		msg = &MsgCFHeaders{}

	case CmdCFCheckpt:
		msg = &MsgCFCheckpt{}

	default:
		return nil, fmt.Errorf("unhandled command [%s]", command)
	}
	return msg, nil
}

// messageHeader defines the header structure for all bitcoin protocol messages.
type messageHeader struct {
	magic    BitcoinNet // 4 bytes
	command  string     // 12 bytes
	length   uint32     // 4 bytes
	checksum [4]byte    // 4 bytes
}

// readMessageHeader reads a bitcoin message header from r.
func readMessageHeader(r io.Reader) (int, *messageHeader, error) {
	// Since readElements doesn't return the amount of bytes read, attempt
	// to read the entire header into a buffer first in case there is a
	// short read so the proper amount of read bytes are known.  This works
	// since the header is a fixed size.
	var headerBytes [MessageHeaderSize]byte
	n, err := io.ReadFull(r, headerBytes[:])
	if err != nil {
		return n, nil, err
	}
	hr := bytes.NewReader(headerBytes[:])

	// Create and populate a messageHeader struct from the raw header bytes.
	hdr := messageHeader{}
	var command [CommandSize]byte
	readElements(hr, &hdr.magic, &command, &hdr.length, &hdr.checksum)

	// Strip trailing zeros from command string.
	hdr.command = string(bytes.TrimRight(command[:], string(0)))

	return n, &hdr, nil
}

// discardInput reads n bytes from reader r in chunks and discards the read
// bytes.  This is used to skip payloads when various errors occur and helps
// prevent rogue nodes from causing massive memory allocation through forging
// header length.
func discardInput(r io.Reader, n uint32) {
	maxSize := uint32(10 * 1024) // 10k at a time
	numReads := n / maxSize
	bytesRemaining := n % maxSize
	if n > 0 {
		buf := make([]byte, maxSize)
		for i := uint32(0); i < numReads; i++ {
			io.ReadFull(r, buf)
		}
	}
	if bytesRemaining > 0 {
		buf := make([]byte, bytesRemaining)
		io.ReadFull(r, buf)
	}
}

// WriteMessageN writes a bitcoin Message to w including the necessary header
// information and returns the number of bytes written.    This function is the
// same as WriteMessage except it also returns the number of bytes written.
func WriteMessageN(w io.Writer, msg Message, pver uint32, btcnet BitcoinNet) (int, error) {
	return WriteMessageWithEncodingN(w, msg, pver, btcnet, BaseEncoding)
}

// WriteMessage writes a bitcoin Message to w including the necessary header
// information.  This function is the same as WriteMessageN except it doesn't
// doesn't return the number of bytes written.  This function is mainly provided
// for backwards compatibility with the original API, but it's also useful for
// callers that don't care about byte counts.
func WriteMessage(w io.Writer, msg Message, pver uint32, btcnet BitcoinNet) error {
	_, err := WriteMessageN(w, msg, pver, btcnet)
	return err
}

// WriteMessageWithEncodingN writes a bitcoin Message to w including the
// necessary header information and returns the number of bytes written.
// This function is the same as WriteMessageN except it also allows the caller
// to specify the message encoding format to be used when serializing wire
// messages.
func WriteMessageWithEncodingN(w io.Writer, msg Message, pver uint32,
	btcnet BitcoinNet, encoding MessageEncoding) (int, error) {

	totalBytes := 0

	// Enforce max command size.
	var command [CommandSize]byte
	cmd := msg.Command()
	if len(cmd) > CommandSize {
		str := fmt.Sprintf("command [%s] is too long [max %v]",
			cmd, CommandSize)
		return totalBytes, messageError("WriteMessage", str)
	}
	copy(command[:], []byte(cmd))

	// Encode the message payload.
	var bw bytes.Buffer
	err := msg.BtcEncode(&bw, pver, encoding)
	if err != nil {
		return totalBytes, err
	}
	payload := bw.Bytes()
	lenp := len(payload)

	// Enforce maximum overall message payload.
	if lenp > MaxMessagePayload {
		str := fmt.Sprintf("message payload is too large - encoded "+
			"%d bytes, but maximum message payload is %d bytes",
			lenp, MaxMessagePayload)
		return totalBytes, messageError("WriteMessage", str)
	}

	// Enforce maximum message payload based on the message type.
	mpl := msg.MaxPayloadLength(pver)
	if uint32(lenp) > mpl {
		str := fmt.Sprintf("message payload is too large - encoded "+
			"%d bytes, but maximum message payload size for "+
			"messages of type [%s] is %d.", lenp, cmd, mpl)
		return totalBytes, messageError("WriteMessage", str)
	}

	// Create header for the message.
	hdr := messageHeader{}
	hdr.magic = btcnet
	hdr.command = cmd
	hdr.length = uint32(lenp)
	copy(hdr.checksum[:], chainhash.DoubleHashB(payload)[0:4])

	// Encode the header for the message.  This is done to a buffer
	// rather than directly to the writer since writeElements doesn't
	// return the number of bytes written.
	hw := bytes.NewBuffer(make([]byte, 0, MessageHeaderSize))
	writeElements(hw, hdr.magic, command, hdr.length, hdr.checksum)

	// Write header.
	n, err := w.Write(hw.Bytes())
	totalBytes += n
	if err != nil {
		return totalBytes, err
	}

	// Only write the payload if there is one, e.g., verack messages don't
	// have one.
	if len(payload) > 0 {
		n, err = w.Write(payload)
		totalBytes += n
	}

	return totalBytes, err
}

// ReadMessageWithEncodingN reads, validates, and parses the next bitcoin Message
// from r for the provided protocol version and bitcoin network.  It returns the
// number of bytes read in addition to the parsed Message and raw bytes which
// comprise the message.  This function is the same as ReadMessageN except it
// allows the caller to specify which message encoding is to to consult when
// decoding wire messages.
func ReadMessageWithEncodingN(r io.Reader, pver uint32, btcnet BitcoinNet,
	enc MessageEncoding) (int, Message, []byte, error) {

	totalBytes := 0
	n, hdr, err := readMessageHeader(r)
	totalBytes += n
	if err != nil {
		return totalBytes, nil, nil, err
	}

	// Enforce maximum message payload.
	if hdr.length > MaxMessagePayload {
		str := fmt.Sprintf("message payload is too large - header "+
			"indicates %d bytes, but max message payload is %d "+
			"bytes.", hdr.length, MaxMessagePayload)
		return totalBytes, nil, nil, messageError("ReadMessage", str)

	}

	// Check for messages from the wrong bitcoin network.
	if hdr.magic != btcnet {
		discardInput(r, hdr.length)
		str := fmt.Sprintf("message from other network [%v]", hdr.magic)
		return totalBytes, nil, nil, messageError("ReadMessage", str)
	}

	// Check for malformed commands.
	command := hdr.command
	if !utf8.ValidString(command) {
		discardInput(r, hdr.length)
		str := fmt.Sprintf("invalid command %v", []byte(command))
		return totalBytes, nil, nil, messageError("ReadMessage", str)
	}

	// Create struct of appropriate message type based on the command.
	msg, err := makeEmptyMessage(command)
	if err != nil {
		discardInput(r, hdr.length)
		return totalBytes, nil, nil, messageError("ReadMessage",
			err.Error())
	}

	// Check for maximum length based on the message type as a malicious client
	// could otherwise create a well-formed header and set the length to max
	// numbers in order to exhaust the machine's memory.
	mpl := msg.MaxPayloadLength(pver)
	if hdr.length > mpl {
		discardInput(r, hdr.length)
		str := fmt.Sprintf("payload exceeds max length - header "+
			"indicates %v bytes, but max payload size for "+
			"messages of type [%v] is %v.", hdr.length, command, mpl)
		return totalBytes, nil, nil, messageError("ReadMessage", str)
	}

	// Read payload.
	payload := make([]byte, hdr.length)
	n, err = io.ReadFull(r, payload)
	totalBytes += n
	if err != nil {
		return totalBytes, nil, nil, err
	}

	// Test checksum.
	checksum := chainhash.DoubleHashB(payload)[0:4]
	if !bytes.Equal(checksum[:], hdr.checksum[:]) {
		str := fmt.Sprintf("payload checksum failed - header "+
			"indicates %v, but actual checksum is %v.",
			hdr.checksum, checksum)
		return totalBytes, nil, nil, messageError("ReadMessage", str)
	}

	// Unmarshal message.  NOTE: This must be a *bytes.Buffer since the
	// MsgVersion BtcDecode function requires it.
	pr := bytes.NewBuffer(payload)
	err = msg.BtcDecode(pr, pver, enc)
	if err != nil {
		return totalBytes, nil, nil, err
	}

	return totalBytes, msg, payload, nil
}

// ReadMessageN reads, validates, and parses the next bitcoin Message from r for
// the provided protocol version and bitcoin network.  It returns the number of
// bytes read in addition to the parsed Message and raw bytes which comprise the
// message.  This function is the same as ReadMessage except it also returns the
// number of bytes read.
func ReadMessageN(r io.Reader, pver uint32, btcnet BitcoinNet) (int, Message, []byte, error) {
	return ReadMessageWithEncodingN(r, pver, btcnet, BaseEncoding)
}

// ReadMessage reads, validates, and parses the next bitcoin Message from r for
// the provided protocol version and bitcoin network.  It returns the parsed
// Message and raw bytes which comprise the message.  This function only differs
// from ReadMessageN in that it doesn't return the number of bytes read.  This
// function is mainly provided for backwards compatibility with the original
// API, but it's also useful for callers that don't care about byte counts.
func ReadMessage(r io.Reader, pver uint32, btcnet BitcoinNet) (Message, []byte, error) {
	_, msg, buf, err := ReadMessageN(r, pver, btcnet)
	return msg, buf, err
}
<|endoftext|>


<commit_before>package util

import (
	"fmt"
	"github.com/fatih/color"
	"github.com/sirupsen/logrus"
	"os"
	"sync"
	"time"
)

type Log interface {
	Debug(format string, a ...interface{})
	Output(format string, a ...interface{})
	Message(format string, a ...interface{})
	Info(format string, a ...interface{})
	Warning(format string, a ...interface{})
	Error(format string, a ...interface{})
}

const (
	timeFormat = "2006/01/02 15:04:05"
)

type ColoredLog struct {
	l     sync.Mutex
	debug bool
}

func (c *ColoredLog) Debug(format string, a ...interface{}) {
	if c.debug {
		now := time.Now()
		c.l.Lock()
		defer c.l.Unlock()
		color.White(c.appendTime(now, format), a...)
	}
}

func (c *ColoredLog) Output(format string, a ...interface{}) {
	c.l.Lock()
	defer c.l.Unlock()
	color.White(format, a...)
}

func (ColoredLog) appendTime(stamp time.Time, str string) string {
	return fmt.Sprintf("%s [rivi] %s", stamp.Format(timeFormat), str)
}

func (c *ColoredLog) Message(format string, a ...interface{}) {
	now := time.Now()
	c.l.Lock()
	defer c.l.Unlock()
	color.White(c.appendTime(now, format), a...)
}

func (c *ColoredLog) Warning(format string, a ...interface{}) {
	now := time.Now()
	c.l.Lock()
	defer c.l.Unlock()
	color.Yellow(c.appendTime(now, format), a...)
}

func (c *ColoredLog) Error(format string, a ...interface{}) {
	now := time.Now()
	c.l.Lock()
	defer c.l.Unlock()
	color.Red(c.appendTime(now, format), a...)
}

func (c *ColoredLog) Info(format string, a ...interface{}) {
	now := time.Now()
	c.l.Lock()
	defer c.l.Unlock()
	color.Cyan(c.appendTime(now, format), a...)
}

type LogrusLog struct {
	log *logrus.Logger
}

func (l *LogrusLog) Debug(format string, a ...interface{}) {
	l.log.Debugf(format, a...)
}

func (l *LogrusLog) Output(format string, a ...interface{}) {
	l.log.Printf(format, a...)
}

func (l *LogrusLog) Message(format string, a ...interface{}) {
	l.log.Printf(format, a...)
}

func (l *LogrusLog) Info(format string, a ...interface{}) {
	l.log.Infof(format, a...)
}

func (l *LogrusLog) Warning(format string, a ...interface{}) {
	l.log.Warningf(format, a...)
}

func (l *LogrusLog) Error(format string, a ...interface{}) {
	l.log.Errorf(format, a...)
}

var Logger Log

func init() {
	logger := "logrus"
	debug := false
	if os.Getenv("BOT_LOGGER") != "" {
		logger = os.Getenv("BOT_LOGGER")
	}
	if len(os.Getenv("BOT_DEBUG")) > 0 {
		debug = true
	}
	switch logger {
	case "logrus":
		l := &LogrusLog{log: logrus.New()}
		Logger = l
		l.log.Out = os.Stdout

		if debug {
			l.log.Level = logrus.DebugLevel
		} else {
			l.log.Level = logrus.InfoLevel
		}
	case "color":
		c := &ColoredLog{}
		if debug {
			c.debug = true
		}
		Logger = c
	}
}
<commit_msg>go vet<commit_after>package util

import (
	"fmt"
	"github.com/fatih/color"
	"github.com/sirupsen/logrus"
	"os"
	"sync"
	"time"
)

type Log interface {
	Debug(format string, a ...interface{})
	Output(format string, a ...interface{})
	Message(format string, a ...interface{})
	Info(format string, a ...interface{})
	Warning(format string, a ...interface{})
	Error(format string, a ...interface{})
}

const (
	timeFormat = time.RFC3339
)

type ColoredLog struct {
	l     sync.Mutex
	debug bool
}

func (c *ColoredLog) Debug(format string, a ...interface{}) {
	if c.debug {
		now := time.Now()
		c.l.Lock()
		defer c.l.Unlock()
		color.White(c.appendTime(now, format), a...)
	}
}

func (c *ColoredLog) Output(format string, a ...interface{}) {
	c.l.Lock()
	defer c.l.Unlock()
	color.White(format, a...)
}

func (*ColoredLog) appendTime(stamp time.Time, str string) string {
	return fmt.Sprintf("%s [rivi] %s", stamp.Format(timeFormat), str)
}

func (c *ColoredLog) Message(format string, a ...interface{}) {
	now := time.Now()
	c.l.Lock()
	defer c.l.Unlock()
	color.White(c.appendTime(now, format), a...)
}

func (c *ColoredLog) Warning(format string, a ...interface{}) {
	now := time.Now()
	c.l.Lock()
	defer c.l.Unlock()
	color.Yellow(c.appendTime(now, format), a...)
}

func (c *ColoredLog) Error(format string, a ...interface{}) {
	now := time.Now()
	c.l.Lock()
	defer c.l.Unlock()
	color.Red(c.appendTime(now, format), a...)
}

func (c *ColoredLog) Info(format string, a ...interface{}) {
	now := time.Now()
	c.l.Lock()
	defer c.l.Unlock()
	color.Cyan(c.appendTime(now, format), a...)
}

type LogrusLog struct {
	log *logrus.Logger
}

func (l *LogrusLog) Debug(format string, a ...interface{}) {
	l.log.Debugf(format, a...)
}

func (l *LogrusLog) Output(format string, a ...interface{}) {
	l.log.Printf(format, a...)
}

func (l *LogrusLog) Message(format string, a ...interface{}) {
	l.log.Printf(format, a...)
}

func (l *LogrusLog) Info(format string, a ...interface{}) {
	l.log.Infof(format, a...)
}

func (l *LogrusLog) Warning(format string, a ...interface{}) {
	l.log.Warningf(format, a...)
}

func (l *LogrusLog) Error(format string, a ...interface{}) {
	l.log.Errorf(format, a...)
}

var Logger Log

func init() {
	logger := "logrus"
	debug := false
	if os.Getenv("BOT_LOGGER") != "" {
		logger = os.Getenv("BOT_LOGGER")
	}
	if len(os.Getenv("BOT_DEBUG")) > 0 {
		debug = true
	}
	switch logger {
	case "logrus":
		l := &LogrusLog{log: logrus.New()}
		Logger = l
		l.log.Out = os.Stdout

		if debug {
			l.log.Level = logrus.DebugLevel
		} else {
			l.log.Level = logrus.InfoLevel
		}
	case "color":
		c := &ColoredLog{}
		if debug {
			c.debug = true
		}
		Logger = c
	}
}
<|endoftext|>


<commit_before>package upsert

import (
	"bytes"
	"errors"
	"fmt"
	"log"
	"reflect"
	"strings"
	"time"

	"github.com/jmoiron/sqlx"
)

const (
	Inserted = iota
	Updated
	NoChange
)

var (
	ErrNoIDReturned = errors.New("no id returned")

	// LongQuery will log long queries if set to a non-zero time
	LongQuery time.Duration
)

// Upserter is an interface specific to sqlx and PostgreSQL that can save a
// single row of data via Upsert(), Update() or Insert(). It doesn't try to
// know anything about relationships between tables. The behavior of Upserter
// depends on three struct tags.
//
//  * db: As with sqlx, this tag is the database column name for the field.
//     If db is not defined, the default is the lowercase value of the field
//     name.
//
//  * upsert: This may either be "key" or "omit". If it's "key", the
//     field/column is part of the where clause when attempting to update
//     an existing column. If it's "omit", the field is ignored completely.
//     By default, the field is considered a non-key value that should be
//     updated/set in the db.
//
//  * upsert_value: This is the placeholder for the value of the field for
//     use by sqlx.NamedExec(). By default, this is :column_name and typically
//     doesn't need to be changed. However, if the field needs to be
//     transformed by an SQL function before storing in the database,
//     this tag can be set. For example, if you had "lat" and "lon" columns,
//     you wouldn't want to store them in the db. Instead you'd want a
//     "location" column tagged with `upsert_value:"ll_to_earth(:lat, :lon)`
//
type Upserter interface {
	// Table returns table name we should save to
	Table() string
}

type columnSpec struct {
	name  string
	value string
}

func newColumnSpec(fieldName string, tag reflect.StructTag) columnSpec {
	cs := columnSpec{}

	// The name of the column is either the value of the "db" struct tag
	// or a lowercase version of the field name.
	dbTag := tag.Get("db")
	if len(dbTag) > 0 {
		cs.name = dbTag
	} else {
		cs.name = strings.ToLower(fieldName)
	}

	// The value placeholder of the column is typically just ":column_name"
	// but can be overriden with upsert_value.
	val := tag.Get("upsert_value")
	if len(val) > 0 {
		cs.value = val
	} else {
		cs.value = ":" + cs.name
	}

	return cs
}

// updateColumns returns the fields that are read from the struct and set
// on upserting in the db. Typically this should include everything except the
// key fields and any composite (array, nested struct) types or any
// field that doesn't map directly into a db column. Tag a field with
// `upsert:"omit"` to explicitly exclude from this list.
func updateColumns(u Upserter) (columns []columnSpec) {
	ut := reflect.TypeOf(u)

	if ut.Kind() == reflect.Ptr {
		ut = ut.Elem()
	}

	if ut.Kind() != reflect.Struct {
		return
	}

	for i := 0; i < ut.NumField(); i++ {
		field := ut.Field(i)
		tag := field.Tag

		// Include any column that isn't tagged with upsert:omit
		if !strings.Contains(tag.Get("upsert"), "omit") {
			columns = append(columns, newColumnSpec(field.Name, tag))
		}
	}

	return
}

// uniqueKeyColumns returns the fields of the struct that together are
// naturally unique. For example, an md5 hash of the content. Or a
// foreign key plus an internal value. This is used in where clause
// when trying to find existing rows. Tag a field with `"upsert:"key"`
// to include in the unique key.
func uniqueKeyColumns(u Upserter) (columns []columnSpec) {
	ut := reflect.TypeOf(u)

	if ut.Kind() == reflect.Ptr {
		ut = ut.Elem()
	}

	if ut.Kind() != reflect.Struct {
		return
	}

	for i := 0; i < ut.NumField(); i++ {
		field := ut.Field(i)
		tag := field.Tag
		// Check if upsert tag contains "key". This wouldn't work
		// if possible options were substrings of one another. For a
		// better implementation, look at src/encoding/json/tags.go
		if strings.Contains(tag.Get("upsert"), "key") {
			columns = append(columns, newColumnSpec(field.Name, tag))
		}
	}

	return
}

// set returns a string like "SET "col1" = :col1, "col2" = :col2" for
// use with sqlx.NamedExec() and friends.
func set(u Upserter) string {
	cols := updateColumns(u)
	n := len(cols)

	b := bytes.Buffer{}

	b.WriteString("SET ")
	for i := 0; i < n; i++ {
		fmt.Fprintf(&b, `"%s" = %s`, cols[i].name, cols[i].value)

		// If we are not at the last value, add a comma
		if i < n-1 {
			b.WriteRune(',')
		}
	}

	return b.String()
}

// values returns a string like `("col1", "col2") VALUES(:col1, :col2)`
// for use with sqlx.NamedExec() etc.
func values(u Upserter) string {
	cols := updateColumns(u)
	n := len(cols)

	b := bytes.Buffer{}

	b.WriteRune('(')
	for i := 0; i < n; i++ {
		fmt.Fprintf(&b, `"%s"`, cols[i].name)

		// If we are not at the last value, add a comma
		if i < n-1 {
			b.WriteRune(',')
		}
	}
	b.WriteRune(')')

	b.WriteString("VALUES (")
	for i := 0; i < n; i++ {
		fmt.Fprintf(&b, `%s`, cols[i].value)

		// If we are not at the last value, add a comma
		if i < n-1 {
			b.WriteRune(',')
		}
	}
	b.WriteRune(')')

	return b.String()
}

// where returns an SQL where clause with all the key columns of
// this Upserter
func where(u Upserter) string {
	keycols := uniqueKeyColumns(u)
	b := bytes.Buffer{}
	n := len(keycols)

	fmt.Fprintf(&b, "WHERE ")

	for i := 0; i < n; i++ {
		// If we need to support NULLs here, the best option may be
		// something like (x = y OR (x is null and y is null))
		// rather than "IS NOT DISTINCT FROM" which doesn't use indexes
		// it seems
		fmt.Fprintf(&b, `%s = %s`, keycols[i].name, keycols[i].value)

		if i < n-1 {
			fmt.Fprint(&b, " AND ")
		}
	}

	return b.String()
}

// updateSQL returns a full SQL command to update this Upserter u
func updateSQL(u Upserter) string {
	q := fmt.Sprintf(`UPDATE "%s" %s %s RETURNING *`,
		u.Table(), set(u), where(u))

	return q
}

// insertSQL returns a full SQL command to insert this Upserter u
func insertSQL(u Upserter) string {
	q := fmt.Sprintf(`INSERT INTO "%s" %s RETURNING *`, u.Table(), values(u))

	return q
}

// getSQL returns a full SQL command to retrieve this Upserter u
func getSQL(u Upserter) string {
	q := fmt.Sprintf(`SELECT * FROM %s %s`, u.Table(), where(u))

	return q
}

func Update(ext sqlx.Ext, u Upserter) (status int, err error) {
	q := updateSQL(u)

	if LongQuery > time.Duration(0) {
		t1 := time.Now()
		defer func() {
			t2 := time.Now()
			if t2.Sub(t1) > LongQuery {
				log.Println(t2.Sub(t1), q, u)
			}
		}()
	}

	otherPtr := reflect.New(reflect.TypeOf(u).Elem())
	other := reflect.Indirect(otherPtr)

	rows, err := sqlx.NamedQuery(ext, getSQL(u), u)
	if err != nil {
		log.Println("error getting", err)
		return
	}

	if rows.Next() {
		err = rows.StructScan(other.Addr().Interface())
		if err != nil {
			log.Println("error scanning", err)
			return
		}

		if reflect.DeepEqual(other.Addr().Interface(), u) {
			log.Println("no change, ignoring update")
			status = NoChange
			return
		} else {
			log.Println("yes there is a change")
		}
	}

	status = Updated

	// Try to update an existing row
	rows, err = sqlx.NamedQuery(ext, q, u)
	if err != nil {
		log.Println(updateSQL(u), err)
		return
	}
	defer rows.Close()

	if rows.Next() {
		err = rows.StructScan(u)
		if err != nil {
			log.Println(err)
		}
	} else {
		// We could not find anything to update.
		err = ErrNoIDReturned
		return
	}

	return
}

// Insert takes either an sqlx.DB or sqlx.Tx as ext, along with a value
// that implements the Upserter() interface. We attempt to insert it
// and set its primary key id value.
func Insert(ext sqlx.Ext, u Upserter) (err error) {
	q := insertSQL(u)

	if LongQuery > time.Duration(0) {
		t1 := time.Now()
		defer func() {
			t2 := time.Now()
			if t2.Sub(t1) > LongQuery {
				log.Println(t2.Sub(t1), q, u)
			}
		}()
	}

	// Try to insert a row
	rows, err := sqlx.NamedQuery(ext, q, u)
	if err != nil {
		log.Println(err)
		return
	}
	defer rows.Close()

	if rows.Next() {
		err = rows.StructScan(u)
		if err != nil {
			log.Println(err)
		}
	} else {
		// No rows were returned but no SQL error. Weird, return generic
		// error.
		err = ErrNoIDReturned
		return
	}

	return
}

// Upsert takes either an sqlx.DB or sqlx.Tx as ext, along with a value
// that implements the Upserter() interface. We attempt to insert/update it
// and set the new primary key id if that succeeds. inserted returns true
// if a new row was inserted. The client is responsible for wrapping
// in a transaction when needed. This can be used when running a transaction
// at a higher level (upserting multiple items).
func Upsert(ext sqlx.Ext, u Upserter) (status int, err error) {
	// Try to update, return immediately if succcesful
	status, err = Update(ext, u)
	if err == nil {
		return
	}

	// Can't update? Try insert
	err = Insert(ext, u)
	if err != nil {
		log.Println(err)
		return
	}

	status = Inserted

	return
}

// UpsertTx takes only an sqlx.DB and wraps the upsert attempt into a
// a transaction.
func UpsertTx(db *sqlx.DB, u Upserter) (status int, err error) {
	tx, err := db.Beginx()
	if err != nil {
		log.Println("can't start transaction", err)
		return
	}
	defer func() {
		if err != nil {
			tx.Rollback()
		} else {
			tx.Commit()
		}
	}()

	// Try to update
	status, err = Update(tx, u)

	// If we have a nil error, we successfully updated. If we have
	// an err other than ErrNoIDReturned, we couldn't update for an
	// unexpected reason. In either case return.
	if err != ErrNoIDReturned {
		return
	}

	// No ID returned in the update? Try insert
	err = Insert(tx, u)
	if err != nil {
		log.Println(err)
		return
	}

	status = Inserted

	return
}

func Delete(ext sqlx.Ext, u Upserter) (err error) {
	q := fmt.Sprintf(`DELETE FROM "%s" %s`,
		u.Table(), where(u))
	_, err = sqlx.NamedExec(ext, q, u)

	if err != nil {
		log.Println("can't delete", err)
		return
	}

	return
}
<commit_msg>fix lack of rows.close<commit_after>package upsert

import (
	"bytes"
	"errors"
	"fmt"
	"log"
	"reflect"
	"strings"
	"time"

	"github.com/jmoiron/sqlx"
)

const (
	Inserted = iota
	Updated
	NoChange
)

var (
	ErrNoIDReturned = errors.New("no id returned")

	// LongQuery will log long queries if set to a non-zero time
	LongQuery time.Duration
)

// Upserter is an interface specific to sqlx and PostgreSQL that can save a
// single row of data via Upsert(), Update() or Insert(). It doesn't try to
// know anything about relationships between tables. The behavior of Upserter
// depends on three struct tags.
//
//  * db: As with sqlx, this tag is the database column name for the field.
//     If db is not defined, the default is the lowercase value of the field
//     name.
//
//  * upsert: This may either be "key" or "omit". If it's "key", the
//     field/column is part of the where clause when attempting to update
//     an existing column. If it's "omit", the field is ignored completely.
//     By default, the field is considered a non-key value that should be
//     updated/set in the db.
//
//  * upsert_value: This is the placeholder for the value of the field for
//     use by sqlx.NamedExec(). By default, this is :column_name and typically
//     doesn't need to be changed. However, if the field needs to be
//     transformed by an SQL function before storing in the database,
//     this tag can be set. For example, if you had "lat" and "lon" columns,
//     you wouldn't want to store them in the db. Instead you'd want a
//     "location" column tagged with `upsert_value:"ll_to_earth(:lat, :lon)`
//
type Upserter interface {
	// Table returns table name we should save to
	Table() string
}

type columnSpec struct {
	name  string
	value string
}

func newColumnSpec(fieldName string, tag reflect.StructTag) columnSpec {
	cs := columnSpec{}

	// The name of the column is either the value of the "db" struct tag
	// or a lowercase version of the field name.
	dbTag := tag.Get("db")
	if len(dbTag) > 0 {
		cs.name = dbTag
	} else {
		cs.name = strings.ToLower(fieldName)
	}

	// The value placeholder of the column is typically just ":column_name"
	// but can be overriden with upsert_value.
	val := tag.Get("upsert_value")
	if len(val) > 0 {
		cs.value = val
	} else {
		cs.value = ":" + cs.name
	}

	return cs
}

// updateColumns returns the fields that are read from the struct and set
// on upserting in the db. Typically this should include everything except the
// key fields and any composite (array, nested struct) types or any
// field that doesn't map directly into a db column. Tag a field with
// `upsert:"omit"` to explicitly exclude from this list.
func updateColumns(u Upserter) (columns []columnSpec) {
	ut := reflect.TypeOf(u)

	if ut.Kind() == reflect.Ptr {
		ut = ut.Elem()
	}

	if ut.Kind() != reflect.Struct {
		return
	}

	for i := 0; i < ut.NumField(); i++ {
		field := ut.Field(i)
		tag := field.Tag

		// Include any column that isn't tagged with upsert:omit
		if !strings.Contains(tag.Get("upsert"), "omit") {
			columns = append(columns, newColumnSpec(field.Name, tag))
		}
	}

	return
}

// uniqueKeyColumns returns the fields of the struct that together are
// naturally unique. For example, an md5 hash of the content. Or a
// foreign key plus an internal value. This is used in where clause
// when trying to find existing rows. Tag a field with `"upsert:"key"`
// to include in the unique key.
func uniqueKeyColumns(u Upserter) (columns []columnSpec) {
	ut := reflect.TypeOf(u)

	if ut.Kind() == reflect.Ptr {
		ut = ut.Elem()
	}

	if ut.Kind() != reflect.Struct {
		return
	}

	for i := 0; i < ut.NumField(); i++ {
		field := ut.Field(i)
		tag := field.Tag
		// Check if upsert tag contains "key". This wouldn't work
		// if possible options were substrings of one another. For a
		// better implementation, look at src/encoding/json/tags.go
		if strings.Contains(tag.Get("upsert"), "key") {
			columns = append(columns, newColumnSpec(field.Name, tag))
		}
	}

	return
}

// set returns a string like "SET "col1" = :col1, "col2" = :col2" for
// use with sqlx.NamedExec() and friends.
func set(u Upserter) string {
	cols := updateColumns(u)
	n := len(cols)

	b := bytes.Buffer{}

	b.WriteString("SET ")
	for i := 0; i < n; i++ {
		fmt.Fprintf(&b, `"%s" = %s`, cols[i].name, cols[i].value)

		// If we are not at the last value, add a comma
		if i < n-1 {
			b.WriteRune(',')
		}
	}

	return b.String()
}

// values returns a string like `("col1", "col2") VALUES(:col1, :col2)`
// for use with sqlx.NamedExec() etc.
func values(u Upserter) string {
	cols := updateColumns(u)
	n := len(cols)

	b := bytes.Buffer{}

	b.WriteRune('(')
	for i := 0; i < n; i++ {
		fmt.Fprintf(&b, `"%s"`, cols[i].name)

		// If we are not at the last value, add a comma
		if i < n-1 {
			b.WriteRune(',')
		}
	}
	b.WriteRune(')')

	b.WriteString("VALUES (")
	for i := 0; i < n; i++ {
		fmt.Fprintf(&b, `%s`, cols[i].value)

		// If we are not at the last value, add a comma
		if i < n-1 {
			b.WriteRune(',')
		}
	}
	b.WriteRune(')')

	return b.String()
}

// where returns an SQL where clause with all the key columns of
// this Upserter
func where(u Upserter) string {
	keycols := uniqueKeyColumns(u)
	b := bytes.Buffer{}
	n := len(keycols)

	fmt.Fprintf(&b, "WHERE ")

	for i := 0; i < n; i++ {
		// If we need to support NULLs here, the best option may be
		// something like (x = y OR (x is null and y is null))
		// rather than "IS NOT DISTINCT FROM" which doesn't use indexes
		// it seems
		fmt.Fprintf(&b, `%s = %s`, keycols[i].name, keycols[i].value)

		if i < n-1 {
			fmt.Fprint(&b, " AND ")
		}
	}

	return b.String()
}

// updateSQL returns a full SQL command to update this Upserter u
func updateSQL(u Upserter) string {
	q := fmt.Sprintf(`UPDATE "%s" %s %s RETURNING *`,
		u.Table(), set(u), where(u))

	return q
}

// insertSQL returns a full SQL command to insert this Upserter u
func insertSQL(u Upserter) string {
	q := fmt.Sprintf(`INSERT INTO "%s" %s RETURNING *`, u.Table(), values(u))

	return q
}

// getSQL returns a full SQL command to retrieve this Upserter u
func getSQL(u Upserter) string {
	q := fmt.Sprintf(`SELECT * FROM %s %s`, u.Table(), where(u))

	return q
}

func Update(ext sqlx.Ext, u Upserter) (status int, err error) {
	q := updateSQL(u)

	if LongQuery > time.Duration(0) {
		t1 := time.Now()
		defer func() {
			t2 := time.Now()
			if t2.Sub(t1) > LongQuery {
				log.Println(t2.Sub(t1), q, u)
			}
		}()
	}

	otherPtr := reflect.New(reflect.TypeOf(u).Elem())
	other := reflect.Indirect(otherPtr)
	otherInterface := other.Addr().Interface()

	rows, err := sqlx.NamedQuery(ext, getSQL(u), u)
	if err != nil {
		log.Println("error getting", err)
		return
	}

	if rows.Next() {
		err = rows.StructScan(otherInterface)
		if err != nil {
			log.Println("error scanning", err)
			return
		}

		if reflect.DeepEqual(otherInterface, u) {
			log.Println("no change, ignoring update")
			status = NoChange
			return
		} else {
			log.Println("yes there is a change", otherInterface, u)
		}
	}
	rows.Close()

	status = Updated

	// Try to update an existing row
	rows, err = sqlx.NamedQuery(ext, q, u)
	if err != nil {
		log.Println(updateSQL(u), err)
		return
	}
	defer rows.Close()

	if rows.Next() {
		err = rows.StructScan(u)
		if err != nil {
			log.Println(err)
		}
	} else {
		// We could not find anything to update.
		err = ErrNoIDReturned
		return
	}

	return
}

// Insert takes either an sqlx.DB or sqlx.Tx as ext, along with a value
// that implements the Upserter() interface. We attempt to insert it
// and set its primary key id value.
func Insert(ext sqlx.Ext, u Upserter) (err error) {
	q := insertSQL(u)

	if LongQuery > time.Duration(0) {
		t1 := time.Now()
		defer func() {
			t2 := time.Now()
			if t2.Sub(t1) > LongQuery {
				log.Println(t2.Sub(t1), q, u)
			}
		}()
	}

	// Try to insert a row
	rows, err := sqlx.NamedQuery(ext, q, u)
	if err != nil {
		log.Println(err)
		return
	}
	defer rows.Close()

	if rows.Next() {
		err = rows.StructScan(u)
		if err != nil {
			log.Println(err)
		}
	} else {
		// No rows were returned but no SQL error. Weird, return generic
		// error.
		err = ErrNoIDReturned
		return
	}

	return
}

// Upsert takes either an sqlx.DB or sqlx.Tx as ext, along with a value
// that implements the Upserter() interface. We attempt to insert/update it
// and set the new primary key id if that succeeds. inserted returns true
// if a new row was inserted. The client is responsible for wrapping
// in a transaction when needed. This can be used when running a transaction
// at a higher level (upserting multiple items).
func Upsert(ext sqlx.Ext, u Upserter) (status int, err error) {
	// Try to update, return immediately if succcesful
	status, err = Update(ext, u)
	if err == nil {
		return
	}

	// Can't update? Try insert
	err = Insert(ext, u)
	if err != nil {
		log.Println(err)
		return
	}

	status = Inserted

	return
}

// UpsertTx takes only an sqlx.DB and wraps the upsert attempt into a
// a transaction.
func UpsertTx(db *sqlx.DB, u Upserter) (status int, err error) {
	tx, err := db.Beginx()
	if err != nil {
		log.Println("can't start transaction", err)
		return
	}
	defer func() {
		if err != nil {
			tx.Rollback()
		} else {
			tx.Commit()
		}
	}()

	// Try to update
	status, err = Update(tx, u)

	// If we have a nil error, we successfully updated. If we have
	// an err other than ErrNoIDReturned, we couldn't update for an
	// unexpected reason. In either case return.
	if err != ErrNoIDReturned {
		return
	}

	// No ID returned in the update? Try insert
	err = Insert(tx, u)
	if err != nil {
		log.Println(err)
		return
	}

	status = Inserted

	return
}

func Delete(ext sqlx.Ext, u Upserter) (err error) {
	q := fmt.Sprintf(`DELETE FROM "%s" %s`,
		u.Table(), where(u))
	_, err = sqlx.NamedExec(ext, q, u)

	if err != nil {
		log.Println("can't delete", err)
		return
	}

	return
}
<|endoftext|>


<commit_before>/*
Copyright IBM Corp. 2018 All Rights Reserved.

SPDX-License-Identifier: Apache-2.0
*/

package endorser

import (
	"fmt"
	"sync"

	"github.com/hyperledger/fabric/core/chaincode/shim"
	"github.com/hyperledger/fabric/core/handlers/endorsement/api"
	endorsement3 "github.com/hyperledger/fabric/core/handlers/endorsement/api/identities"
	"github.com/hyperledger/fabric/core/transientstore"
	pb "github.com/hyperledger/fabric/protos/peer"
	putils "github.com/hyperledger/fabric/protos/utils"
	"github.com/pkg/errors"
)

//go:generate mockery -dir . -name TransientStoreRetriever -case underscore -output mocks/
//go:generate mockery -dir ../transientstore/ -name Store -case underscore -output mocks/

// TransientStoreRetriever retrieves transient stores
type TransientStoreRetriever interface {
	// StoreForChannel returns the transient store for the given channel
	StoreForChannel(channel string) transientstore.Store
}

//go:generate mockery -dir . -name ChannelStateRetriever -case underscore -output mocks/

// ChannelStateRetriever retrieves Channel state
type ChannelStateRetriever interface {
	// ChannelState returns a QueryCreator for the given Channel
	NewQueryCreator(channel string) (QueryCreator, error)
}

//go:generate mockery -dir . -name PluginMapper -case underscore -output mocks/

// PluginMapper maps plugin names to their corresponding factories
type PluginMapper interface {
	PluginFactoryByName(name PluginName) endorsement.PluginFactory
}

// MapBasedPluginMapper maps plugin names to their corresponding factories
type MapBasedPluginMapper map[string]endorsement.PluginFactory

// PluginFactoryByName returns a plugin factory for the given plugin name, or nil if not found
func (m MapBasedPluginMapper) PluginFactoryByName(name PluginName) endorsement.PluginFactory {
	return m[string(name)]
}

// Context defines the data that is related to an in-flight endorsement
type Context struct {
	PluginName     string
	Channel        string
	TxID           string
	Proposal       *pb.Proposal
	SignedProposal *pb.SignedProposal
	Visibility     []byte
	Response       *pb.Response
	Event          []byte
	ChaincodeID    *pb.ChaincodeID
	SimRes         []byte
}

// String returns a text representation of this context
func (c Context) String() string {
	return fmt.Sprintf("{plugin: %s, channel: %s, tx: %s, chaincode: %s}", c.PluginName, c.Channel, c.TxID, c.ChaincodeID.Name)
}

// PluginSupport aggregates the support interfaces
// needed for the operation of the plugin endorser
type PluginSupport struct {
	ChannelStateRetriever
	endorsement3.SigningIdentityFetcher
	PluginMapper
	TransientStoreRetriever
}

// NewPluginEndorser endorses with using a plugin
func NewPluginEndorser(ps *PluginSupport) *PluginEndorser {
	return &PluginEndorser{
		SigningIdentityFetcher:  ps.SigningIdentityFetcher,
		PluginMapper:            ps.PluginMapper,
		pluginChannelMapping:    make(map[PluginName]*pluginsByChannel),
		ChannelStateRetriever:   ps.ChannelStateRetriever,
		TransientStoreRetriever: ps.TransientStoreRetriever,
	}
}

// PluginName defines the name of the plugin as it appears in the configuration
type PluginName string

type pluginsByChannel struct {
	sync.RWMutex
	pluginFactory    endorsement.PluginFactory
	channels2Plugins map[string]endorsement.Plugin
	pe               *PluginEndorser
}

func (pbc *pluginsByChannel) createPluginIfAbsent(channel string) (endorsement.Plugin, error) {
	pbc.RLock()
	plugin, exists := pbc.channels2Plugins[channel]
	pbc.RUnlock()
	if exists {
		return plugin, nil
	}

	pbc.Lock()
	defer pbc.Unlock()
	plugin, exists = pbc.channels2Plugins[channel]
	if exists {
		return plugin, nil
	}

	pluginInstance := pbc.pluginFactory.New()
	plugin, err := pbc.initPlugin(pluginInstance, channel)
	if err != nil {
		return nil, err
	}
	pbc.channels2Plugins[channel] = plugin
	return plugin, nil
}

func (pbc *pluginsByChannel) initPlugin(plugin endorsement.Plugin, channel string) (endorsement.Plugin, error) {
	var dependencies []endorsement.Dependency
	var err error
	// If this is a channel endorsement, add the channel state as a dependency
	if channel != "" {
		query, err := pbc.pe.NewQueryCreator(channel)
		if err != nil {
			return nil, errors.Wrap(err, "failed obtaining channel state")
		}
		store := pbc.pe.TransientStoreRetriever.StoreForChannel(channel)
		if store == nil {
			return nil, errors.Errorf("transient store for channel %s was not initialized", channel)
		}
		dependencies = append(dependencies, &ChannelState{QueryCreator: query, Store: store})
	}
	// Add the SigningIdentityFetcher as a dependency
	dependencies = append(dependencies, pbc.pe.SigningIdentityFetcher)
	err = plugin.Init(dependencies...)
	if err != nil {
		return nil, err
	}
	return plugin, nil
}

// PluginEndorser endorsers proposal responses using plugins
type PluginEndorser struct {
	sync.Mutex
	PluginMapper
	pluginChannelMapping map[PluginName]*pluginsByChannel
	ChannelStateRetriever
	endorsement3.SigningIdentityFetcher
	TransientStoreRetriever
}

// EndorseWithPlugin endorses the response with a plugin
func (pe *PluginEndorser) EndorseWithPlugin(ctx Context) (*pb.ProposalResponse, error) {
	endorserLogger.Debug("Entering endorsement for", ctx)

	if ctx.Response == nil {
		return nil, errors.New("response is nil")
	}

	if ctx.Response.Status >= shim.ERRORTHRESHOLD {
		return &pb.ProposalResponse{Response: ctx.Response}, nil
	}

	plugin, err := pe.getOrCreatePlugin(PluginName(ctx.PluginName), ctx.Channel)
	if err != nil {
		endorserLogger.Warning("Endorsement with plugin for", ctx, " failed:", err)
		return nil, errors.Errorf("plugin with name %s could not be used: %v", ctx.PluginName, err)
	}

	prpBytes, err := proposalResponsePayloadFromContext(ctx)
	if err != nil {
		endorserLogger.Warning("Endorsement with plugin for", ctx, " failed:", err)
		return nil, errors.Wrap(err, "failed assembling proposal response payload")
	}

	endorsement, prpBytes, err := plugin.Endorse(prpBytes, ctx.SignedProposal)
	if err != nil {
		endorserLogger.Warning("Endorsement with plugin for", ctx, " failed:", err)
		return nil, errors.WithStack(err)
	}

	resp := &pb.ProposalResponse{
		Version:     1,
		Endorsement: endorsement,
		Payload:     prpBytes,
		Response:    ctx.Response,
	}
	endorserLogger.Debug("Exiting", ctx)
	return resp, nil
}

// getAndStorePlugin returns a plugin instance for the given plugin name and channel
func (pe *PluginEndorser) getOrCreatePlugin(plugin PluginName, channel string) (endorsement.Plugin, error) {
	pluginFactory := pe.PluginFactoryByName(plugin)
	if pluginFactory == nil {
		return nil, errors.Errorf("plugin with name %s wasn't found", plugin)
	}

	pluginsByChannel := pe.getOrCreatePluginChannelMapping(PluginName(plugin), pluginFactory)
	return pluginsByChannel.createPluginIfAbsent(channel)
}

func (pe *PluginEndorser) getOrCreatePluginChannelMapping(plugin PluginName, pf endorsement.PluginFactory) *pluginsByChannel {
	pe.Lock()
	defer pe.Unlock()
	endorserChannelMapping, exists := pe.pluginChannelMapping[PluginName(plugin)]
	if !exists {
		endorserChannelMapping = &pluginsByChannel{
			pluginFactory:    pf,
			channels2Plugins: make(map[string]endorsement.Plugin),
			pe:               pe,
		}
		pe.pluginChannelMapping[PluginName(plugin)] = endorserChannelMapping
	}
	return endorserChannelMapping
}

func proposalResponsePayloadFromContext(ctx Context) ([]byte, error) {
	hdr, err := putils.GetHeader(ctx.Proposal.Header)
	if err != nil {
		endorserLogger.Warning("Failed parsing header", err)
		return nil, errors.Wrap(err, "failed parsing header")
	}

	pHashBytes, err := putils.GetProposalHash1(hdr, ctx.Proposal.Payload, ctx.Visibility)
	if err != nil {
		endorserLogger.Warning("Failed computing proposal hash", err)
		return nil, errors.Wrap(err, "could not compute proposal hash")
	}

	prpBytes, err := putils.GetBytesProposalResponsePayload(pHashBytes, ctx.Response, ctx.SimRes, ctx.Event, ctx.ChaincodeID)
	if err != nil {
		endorserLogger.Warning("Failed marshaling the proposal response payload to bytes", err)
		return nil, errors.New("failure while marshaling the ProposalResponsePayload")
	}
	return prpBytes, nil
}
<commit_msg>[FAB-14165] typo<commit_after>/*
Copyright IBM Corp. 2018 All Rights Reserved.

SPDX-License-Identifier: Apache-2.0
*/

package endorser

import (
	"fmt"
	"sync"

	"github.com/hyperledger/fabric/core/chaincode/shim"
	"github.com/hyperledger/fabric/core/handlers/endorsement/api"
	endorsement3 "github.com/hyperledger/fabric/core/handlers/endorsement/api/identities"
	"github.com/hyperledger/fabric/core/transientstore"
	pb "github.com/hyperledger/fabric/protos/peer"
	putils "github.com/hyperledger/fabric/protos/utils"
	"github.com/pkg/errors"
)

//go:generate mockery -dir . -name TransientStoreRetriever -case underscore -output mocks/
//go:generate mockery -dir ../transientstore/ -name Store -case underscore -output mocks/

// TransientStoreRetriever retrieves transient stores
type TransientStoreRetriever interface {
	// StoreForChannel returns the transient store for the given channel
	StoreForChannel(channel string) transientstore.Store
}

//go:generate mockery -dir . -name ChannelStateRetriever -case underscore -output mocks/

// ChannelStateRetriever retrieves Channel state
type ChannelStateRetriever interface {
	// ChannelState returns a QueryCreator for the given Channel
	NewQueryCreator(channel string) (QueryCreator, error)
}

//go:generate mockery -dir . -name PluginMapper -case underscore -output mocks/

// PluginMapper maps plugin names to their corresponding factories
type PluginMapper interface {
	PluginFactoryByName(name PluginName) endorsement.PluginFactory
}

// MapBasedPluginMapper maps plugin names to their corresponding factories
type MapBasedPluginMapper map[string]endorsement.PluginFactory

// PluginFactoryByName returns a plugin factory for the given plugin name, or nil if not found
func (m MapBasedPluginMapper) PluginFactoryByName(name PluginName) endorsement.PluginFactory {
	return m[string(name)]
}

// Context defines the data that is related to an in-flight endorsement
type Context struct {
	PluginName     string
	Channel        string
	TxID           string
	Proposal       *pb.Proposal
	SignedProposal *pb.SignedProposal
	Visibility     []byte
	Response       *pb.Response
	Event          []byte
	ChaincodeID    *pb.ChaincodeID
	SimRes         []byte
}

// String returns a text representation of this context
func (c Context) String() string {
	return fmt.Sprintf("{plugin: %s, channel: %s, tx: %s, chaincode: %s}", c.PluginName, c.Channel, c.TxID, c.ChaincodeID.Name)
}

// PluginSupport aggregates the support interfaces
// needed for the operation of the plugin endorser
type PluginSupport struct {
	ChannelStateRetriever
	endorsement3.SigningIdentityFetcher
	PluginMapper
	TransientStoreRetriever
}

// NewPluginEndorser endorses with using a plugin
func NewPluginEndorser(ps *PluginSupport) *PluginEndorser {
	return &PluginEndorser{
		SigningIdentityFetcher:  ps.SigningIdentityFetcher,
		PluginMapper:            ps.PluginMapper,
		pluginChannelMapping:    make(map[PluginName]*pluginsByChannel),
		ChannelStateRetriever:   ps.ChannelStateRetriever,
		TransientStoreRetriever: ps.TransientStoreRetriever,
	}
}

// PluginName defines the name of the plugin as it appears in the configuration
type PluginName string

type pluginsByChannel struct {
	sync.RWMutex
	pluginFactory    endorsement.PluginFactory
	channels2Plugins map[string]endorsement.Plugin
	pe               *PluginEndorser
}

func (pbc *pluginsByChannel) createPluginIfAbsent(channel string) (endorsement.Plugin, error) {
	pbc.RLock()
	plugin, exists := pbc.channels2Plugins[channel]
	pbc.RUnlock()
	if exists {
		return plugin, nil
	}

	pbc.Lock()
	defer pbc.Unlock()
	plugin, exists = pbc.channels2Plugins[channel]
	if exists {
		return plugin, nil
	}

	pluginInstance := pbc.pluginFactory.New()
	plugin, err := pbc.initPlugin(pluginInstance, channel)
	if err != nil {
		return nil, err
	}
	pbc.channels2Plugins[channel] = plugin
	return plugin, nil
}

func (pbc *pluginsByChannel) initPlugin(plugin endorsement.Plugin, channel string) (endorsement.Plugin, error) {
	var dependencies []endorsement.Dependency
	var err error
	// If this is a channel endorsement, add the channel state as a dependency
	if channel != "" {
		query, err := pbc.pe.NewQueryCreator(channel)
		if err != nil {
			return nil, errors.Wrap(err, "failed obtaining channel state")
		}
		store := pbc.pe.TransientStoreRetriever.StoreForChannel(channel)
		if store == nil {
			return nil, errors.Errorf("transient store for channel %s was not initialized", channel)
		}
		dependencies = append(dependencies, &ChannelState{QueryCreator: query, Store: store})
	}
	// Add the SigningIdentityFetcher as a dependency
	dependencies = append(dependencies, pbc.pe.SigningIdentityFetcher)
	err = plugin.Init(dependencies...)
	if err != nil {
		return nil, err
	}
	return plugin, nil
}

// PluginEndorser endorsers proposal responses using plugins
type PluginEndorser struct {
	sync.Mutex
	PluginMapper
	pluginChannelMapping map[PluginName]*pluginsByChannel
	ChannelStateRetriever
	endorsement3.SigningIdentityFetcher
	TransientStoreRetriever
}

// EndorseWithPlugin endorses the response with a plugin
func (pe *PluginEndorser) EndorseWithPlugin(ctx Context) (*pb.ProposalResponse, error) {
	endorserLogger.Debug("Entering endorsement for", ctx)

	if ctx.Response == nil {
		return nil, errors.New("response is nil")
	}

	if ctx.Response.Status >= shim.ERRORTHRESHOLD {
		return &pb.ProposalResponse{Response: ctx.Response}, nil
	}

	plugin, err := pe.getOrCreatePlugin(PluginName(ctx.PluginName), ctx.Channel)
	if err != nil {
		endorserLogger.Warning("Endorsement with plugin for", ctx, " failed:", err)
		return nil, errors.Errorf("plugin with name %s could not be used: %v", ctx.PluginName, err)
	}

	prpBytes, err := proposalResponsePayloadFromContext(ctx)
	if err != nil {
		endorserLogger.Warning("Endorsement with plugin for", ctx, " failed:", err)
		return nil, errors.Wrap(err, "failed assembling proposal response payload")
	}

	endorsement, prpBytes, err := plugin.Endorse(prpBytes, ctx.SignedProposal)
	if err != nil {
		endorserLogger.Warning("Endorsement with plugin for", ctx, " failed:", err)
		return nil, errors.WithStack(err)
	}

	resp := &pb.ProposalResponse{
		Version:     1,
		Endorsement: endorsement,
		Payload:     prpBytes,
		Response:    ctx.Response,
	}
	endorserLogger.Debug("Exiting", ctx)
	return resp, nil
}

// getOrCreatePlugin returns a plugin instance for the given plugin name and channel
func (pe *PluginEndorser) getOrCreatePlugin(plugin PluginName, channel string) (endorsement.Plugin, error) {
	pluginFactory := pe.PluginFactoryByName(plugin)
	if pluginFactory == nil {
		return nil, errors.Errorf("plugin with name %s wasn't found", plugin)
	}

	pluginsByChannel := pe.getOrCreatePluginChannelMapping(PluginName(plugin), pluginFactory)
	return pluginsByChannel.createPluginIfAbsent(channel)
}

func (pe *PluginEndorser) getOrCreatePluginChannelMapping(plugin PluginName, pf endorsement.PluginFactory) *pluginsByChannel {
	pe.Lock()
	defer pe.Unlock()
	endorserChannelMapping, exists := pe.pluginChannelMapping[PluginName(plugin)]
	if !exists {
		endorserChannelMapping = &pluginsByChannel{
			pluginFactory:    pf,
			channels2Plugins: make(map[string]endorsement.Plugin),
			pe:               pe,
		}
		pe.pluginChannelMapping[PluginName(plugin)] = endorserChannelMapping
	}
	return endorserChannelMapping
}

func proposalResponsePayloadFromContext(ctx Context) ([]byte, error) {
	hdr, err := putils.GetHeader(ctx.Proposal.Header)
	if err != nil {
		endorserLogger.Warning("Failed parsing header", err)
		return nil, errors.Wrap(err, "failed parsing header")
	}

	pHashBytes, err := putils.GetProposalHash1(hdr, ctx.Proposal.Payload, ctx.Visibility)
	if err != nil {
		endorserLogger.Warning("Failed computing proposal hash", err)
		return nil, errors.Wrap(err, "could not compute proposal hash")
	}

	prpBytes, err := putils.GetBytesProposalResponsePayload(pHashBytes, ctx.Response, ctx.SimRes, ctx.Event, ctx.ChaincodeID)
	if err != nil {
		endorserLogger.Warning("Failed marshaling the proposal response payload to bytes", err)
		return nil, errors.New("failure while marshaling the ProposalResponsePayload")
	}
	return prpBytes, nil
}
<|endoftext|>


<commit_before>package utils

import (
	"errors"
	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/credentials"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/s3"
	"github.com/aws/aws-sdk-go/service/s3/s3manager"
	"net/http"
	"os"

	"github.com/techjanitor/pram-post/config"
)

// Authenticate to AWS and return handler
func getS3() (svc *session.Session, err error) {

	// new credentials from settings
	creds := credentials.NewStaticCredentials(config.Settings.Amazon.Id, config.Settings.Amazon.Key, "")

	// create our session
	svc = session.New(&aws.Config{
		Region:      aws.String(config.Settings.Amazon.Region),
		Credentials: creds,
	})

	return

}

// Upload a file to S3
func UploadS3(filepath, filename string) (err error) {

	session, err := getS3()
	if err != nil {
		return
	}

	file, err := os.Open(filepath)
	if err != nil {
		return errors.New("problem opening file for s3")
	}
	defer file.Close()

	uploader := s3manager.NewUploader(session)

	params := &s3manager.UploadInput{
		Bucket:               aws.String(config.Settings.Amazon.Bucket),
		Key:                  aws.String(filename),
		Body:                 file,
		ServerSideEncryption: aws.String(s3.ServerSideEncryptionAes256),
	}

	_, err = uploader.Upload(params)

	return

}

// Delete a file from S3
func DeleteS3(object string) (err error) {

	session, err := getS3()
	if err != nil {
		return
	}

	svc := s3.New(session)

	params := &s3.DeleteObjectInput{
		Bucket: aws.String(config.Settings.Amazon.Bucket),
		Key:    aws.String(object),
	}

	_, err = svc.DeleteObject(params)

	return

}
<commit_msg>use http mime sniffer<commit_after>package utils

import (
	"errors"
	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/credentials"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/s3"
	"github.com/aws/aws-sdk-go/service/s3/s3manager"
	"os"

	"github.com/techjanitor/pram-post/config"
)

// Authenticate to AWS and return handler
func getS3() (svc *session.Session, err error) {

	// new credentials from settings
	creds := credentials.NewStaticCredentials(config.Settings.Amazon.Id, config.Settings.Amazon.Key, "")

	// create our session
	svc = session.New(&aws.Config{
		Region:      aws.String(config.Settings.Amazon.Region),
		Credentials: creds,
	})

	return

}

// Upload a file to S3
func UploadS3(filepath, filename string) (err error) {

	session, err := getS3()
	if err != nil {
		return
	}

	file, err := os.Open(filepath)
	if err != nil {
		return errors.New("problem opening file for s3")
	}
	defer file.Close()

	uploader := s3manager.NewUploader(session)

	params := &s3manager.UploadInput{
		Bucket:               aws.String(config.Settings.Amazon.Bucket),
		Key:                  aws.String(filename),
		Body:                 file,
		ServerSideEncryption: aws.String(s3.ServerSideEncryptionAes256),
	}

	_, err = uploader.Upload(params)

	return

}

// Delete a file from S3
func DeleteS3(object string) (err error) {

	session, err := getS3()
	if err != nil {
		return
	}

	svc := s3.New(session)

	params := &s3.DeleteObjectInput{
		Bucket: aws.String(config.Settings.Amazon.Bucket),
		Key:    aws.String(object),
	}

	_, err = svc.DeleteObject(params)

	return

}
<|endoftext|>