17991712134

Committed 24 Sep 2025 10:49PM UTC coverage: 57.143% (-9.5%) from 66.653%

Build # 17991712134

Build Type

push

github

Committed by

web-flow

Commit Message

Merge pull request #10183 from Roasbeef/brontide-static-alloac

brontide: eliminate all allocations from WriteMessage+Flush

Run Details

59 of 67 new or added lines in 3 files covered. (88.06%)

28705 existing lines in 459 files now uncovered.

99466 of 174065 relevant lines covered (57.14%)

1.77 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

87.07

/brontide/noise.go

package brontide

import (
        "crypto/cipher"
        "crypto/sha256"
        "encoding/binary"
        "errors"
        "fmt"
        "io"
        "math"
        "sync"
        "time"

        "github.com/btcsuite/btcd/btcec/v2"
        "github.com/lightningnetwork/lnd/keychain"
        "golang.org/x/crypto/chacha20poly1305"
        "golang.org/x/crypto/hkdf"
)

const (
        // protocolName is the precise instantiation of the Noise protocol
        // handshake at the center of Brontide. This value will be used as part
        // of the prologue. If the initiator and responder aren't using the
        // exact same string for this value, along with prologue of the Bitcoin
        // network, then the initial handshake will fail.
        protocolName = "Noise_XK_secp256k1_ChaChaPoly_SHA256"

        // macSize is the length in bytes of the tags generated by poly1305.
        macSize = 16

        // lengthHeaderSize is the number of bytes used to prefix encode the
        // length of a message payload.
        lengthHeaderSize = 2

        // encHeaderSize is the number of bytes required to hold an encrypted
        // header and it's MAC.
        encHeaderSize = lengthHeaderSize + macSize

        // maxMessageSize is the maximum size of an encrypted message including
        // the MAC. This is the max payload (65535) plus the MAC size (16).
        maxMessageSize = math.MaxUint16 + macSize

        // keyRotationInterval is the number of messages sent on a single
        // cipher stream before the keys are rotated forwards.
        keyRotationInterval = 1000

        // handshakeReadTimeout is a read timeout that will be enforced when
        // waiting for data payloads during the various acts of Brontide. If
        // the remote party fails to deliver the proper payload within this
        // time frame, then we'll fail the connection.
        handshakeReadTimeout = time.Second * 5
)

var (
        // ErrMaxMessageLengthExceeded is returned when a message to be written to
        // the cipher session exceeds the maximum allowed message payload.
        ErrMaxMessageLengthExceeded = errors.New("the generated payload exceeds " +
                "the max allowed message length of (2^16)-1")

        // ErrMessageNotFlushed signals that the connection cannot accept a new
        // message because the prior message has not been fully flushed.
        ErrMessageNotFlushed = errors.New("prior message not flushed")

        // lightningPrologue is the noise prologue that is used to initialize
        // the brontide noise handshake.
        lightningPrologue = []byte("lightning")

        // ephemeralGen is the default ephemeral key generator, used to derive a
        // unique ephemeral key for each brontide handshake.
        ephemeralGen = func() (*btcec.PrivateKey, error) {
                return btcec.NewPrivateKey()
        }

        // headerBufferPool is a pool for encrypted header buffers.
        headerBufferPool = &sync.Pool{
                New: func() interface{} {
                        b := make([]byte, 0, encHeaderSize)
                        return &b
                },
        }

        // bodyBufferPool is a pool for encrypted message body buffers.
        bodyBufferPool = &sync.Pool{
                New: func() interface{} {
                        // Allocate max size to avoid reallocation.
                        // maxMessageSize already includes the MAC.
                        b := make([]byte, 0, maxMessageSize)
                        return &b
                },
        }
)

// ecdh performs an ECDH operation between pub and priv. The returned value is
// the sha256 of the compressed shared point.
func ecdh(pub *btcec.PublicKey, priv keychain.SingleKeyECDH) ([]byte, error) {
        hash, err := priv.ECDH(pub)
        return hash[:], err
}

// cipherState encapsulates the state for the AEAD which will be used to
// encrypt+authenticate any payloads sent during the handshake, and messages
// sent once the handshake has completed.
type cipherState struct {
        // nonce is the nonce passed into the chacha20-poly1305 instance for
        // encryption+decryption. The nonce is incremented after each successful
        // encryption/decryption.
        //
        // TODO(roasbeef): this should actually be 96 bit
        nonce uint64

        // nonceBuffer is a reusable buffer for the nonce to avoid allocations.
        nonceBuffer [12]byte

        // secretKey is the shared symmetric key which will be used to
        // instantiate the cipher.
        //
        // TODO(roasbeef): m-lock??
        secretKey [32]byte

        // salt is an additional secret which is used during key rotation to
        // generate new keys.
        salt [32]byte

        // cipher is an instance of the ChaCha20-Poly1305 AEAD construction
        // created using the secretKey above.
        cipher cipher.AEAD
}

// Encrypt returns a ciphertext which is the encryption of the plainText
// observing the passed associatedData within the AEAD construction.
func (c *cipherState) Encrypt(associatedData, cipherText, plainText []byte) []byte {
        defer func() {
                c.nonce++

                if c.nonce == keyRotationInterval {
                        c.rotateKey()
                }
        }()

        // Write the nonce counter to the buffer (bytes 4-11).
        binary.LittleEndian.PutUint64(c.nonceBuffer[4:], c.nonce)

        return c.cipher.Seal(
                cipherText, c.nonceBuffer[:], plainText, associatedData,
        )
}

// Decrypt attempts to decrypt the passed ciphertext observing the specified
// associatedData within the AEAD construction. In the case that the final MAC
// check fails, then a non-nil error will be returned.
func (c *cipherState) Decrypt(associatedData, plainText, cipherText []byte) ([]byte, error) {
        defer func() {
                c.nonce++

                if c.nonce == keyRotationInterval {
                        c.rotateKey()
                }
        }()

        // Write the nonce counter to the buffer (bytes 4-11).
        binary.LittleEndian.PutUint64(c.nonceBuffer[4:], c.nonce)

        return c.cipher.Open(
                plainText, c.nonceBuffer[:], cipherText, associatedData,
        )
}

// InitializeKey initializes the secret key and AEAD cipher scheme based off of
// the passed key.
func (c *cipherState) InitializeKey(key [32]byte) {
        c.secretKey = key
        c.nonce = 0

        // Safe to ignore the error here as our key is properly sized
        // (32-bytes).
        c.cipher, _ = chacha20poly1305.New(c.secretKey[:])
}

// InitializeKeyWithSalt is identical to InitializeKey however it also sets the
// cipherState's salt field which is used for key rotation.
func (c *cipherState) InitializeKeyWithSalt(salt, key [32]byte) {
        c.salt = salt
        c.InitializeKey(key)
}

// rotateKey rotates the current encryption/decryption key for this cipherState
// instance. Key rotation is performed by ratcheting the current key forward
// using an HKDF invocation with the cipherState's salt as the salt, and the
// current key as the input.
func (c *cipherState) rotateKey() {
        var (
                info    []byte
                nextKey [32]byte
        )

        oldKey := c.secretKey
        h := hkdf.New(sha256.New, oldKey[:], c.salt[:], info)

        // hkdf(ck, k, zero)
        // |
        // | \
        // |  \
        // ck  k'
        h.Read(c.salt[:])
        h.Read(nextKey[:])

        c.InitializeKey(nextKey)
}

// symmetricState encapsulates a cipherState object and houses the ephemeral
// handshake digest state. This struct is used during the handshake to derive
// new shared secrets based off of the result of ECDH operations. Ultimately,
// the final key yielded by this struct is the result of an incremental
// Triple-DH operation.
type symmetricState struct {
        cipherState

        // chainingKey is used as the salt to the HKDF function to derive a new
        // chaining key as well as a new tempKey which is used for
        // encryption/decryption.
        chainingKey [32]byte

        // tempKey is the latter 32 bytes resulted from the latest HKDF
        // iteration. This key is used to encrypt/decrypt any handshake
        // messages or payloads sent until the next DH operation is executed.
        tempKey [32]byte

        // handshakeDigest is the cumulative hash digest of all handshake
        // messages sent from start to finish. This value is never transmitted
        // to the other side, but will be used as the AD when
        // encrypting/decrypting messages using our AEAD construction.
        handshakeDigest [32]byte
}

// mixKey implements a basic HKDF-based key ratchet. This method is called
// with the result of each DH output generated during the handshake process.
// The first 32 bytes extract from the HKDF reader is the next chaining key,
// then latter 32 bytes become the temp secret key using within any future AEAD
// operations until another DH operation is performed.
func (s *symmetricState) mixKey(input []byte) {
        var info []byte

        secret := input
        salt := s.chainingKey
        h := hkdf.New(sha256.New, secret, salt[:], info)

        // hkdf(ck, input, zero)
        // |
        // | \
        // |  \
        // ck  k
        h.Read(s.chainingKey[:])
        h.Read(s.tempKey[:])

        // cipher.k = temp_key
        s.InitializeKey(s.tempKey)
}

// mixHash hashes the passed input data into the cumulative handshake digest.
// The running result of this value (h) is used as the associated data in all
// decryption/encryption operations.
func (s *symmetricState) mixHash(data []byte) {
        h := sha256.New()
        h.Write(s.handshakeDigest[:])
        h.Write(data)

        copy(s.handshakeDigest[:], h.Sum(nil))
}

// EncryptAndHash returns the authenticated encryption of the passed plaintext.
// When encrypting the handshake digest (h) is used as the associated data to
// the AEAD cipher.
func (s *symmetricState) EncryptAndHash(plaintext []byte) []byte {
        ciphertext := s.Encrypt(s.handshakeDigest[:], nil, plaintext)

        s.mixHash(ciphertext)

        return ciphertext
}

// DecryptAndHash returns the authenticated decryption of the passed
// ciphertext. When encrypting the handshake digest (h) is used as the
// associated data to the AEAD cipher.
func (s *symmetricState) DecryptAndHash(ciphertext []byte) ([]byte, error) {
        plaintext, err := s.Decrypt(s.handshakeDigest[:], nil, ciphertext)
        if err != nil {
                return nil, err
        }

        s.mixHash(ciphertext)

        return plaintext, nil
}

// InitializeSymmetric initializes the symmetric state by setting the handshake
// digest (h) and the chaining key (ck) to protocol name.
func (s *symmetricState) InitializeSymmetric(protocolName []byte) {
        var empty [32]byte

        s.handshakeDigest = sha256.Sum256(protocolName)
        s.chainingKey = s.handshakeDigest
        s.InitializeKey(empty)
}

// handshakeState encapsulates the symmetricState and keeps track of all the
// public keys (static and ephemeral) for both sides during the handshake
// transcript. If the handshake completes successfully, then two instances of a
// cipherState are emitted: one to encrypt messages from initiator to
// responder, and the other for the opposite direction.
type handshakeState struct {
        symmetricState

        initiator bool

        localStatic    keychain.SingleKeyECDH
        localEphemeral keychain.SingleKeyECDH // nolint (false positive)

        remoteStatic    *btcec.PublicKey
        remoteEphemeral *btcec.PublicKey
}

// newHandshakeState returns a new instance of the handshake state initialized
// with the prologue and protocol name. If this is the responder's handshake
// state, then the remotePub can be nil.
func newHandshakeState(initiator bool, prologue []byte,
        localKey keychain.SingleKeyECDH,
        remotePub *btcec.PublicKey) handshakeState {

        h := handshakeState{
                initiator:    initiator,
                localStatic:  localKey,
                remoteStatic: remotePub,
        }

        // Set the current chaining key and handshake digest to the hash of the
        // protocol name, and additionally mix in the prologue. If either sides
        // disagree about the prologue or protocol name, then the handshake
        // will fail.
        h.InitializeSymmetric([]byte(protocolName))
        h.mixHash(prologue)

        // In Noise_XK, the initiator should know the responder's static
        // public key, therefore we include the responder's static key in the
        // handshake digest. If the initiator gets this value wrong, then the
        // handshake will fail.
        if initiator {
                h.mixHash(remotePub.SerializeCompressed())
        } else {
                h.mixHash(localKey.PubKey().SerializeCompressed())
        }

        return h
}

// EphemeralGenerator is a functional option that allows callers to substitute
// a custom function for use when generating ephemeral keys for ActOne or
// ActTwo. The function closure returned by this function can be passed into
// NewBrontideMachine as a function option parameter.
func EphemeralGenerator(gen func() (*btcec.PrivateKey, error)) func(*Machine) {
        return func(m *Machine) {
                m.ephemeralGen = gen
        }
}

// Machine is a state-machine which implements Brontide: an
// Authenticated-key Exchange in Three Acts. Brontide is derived from the Noise
// framework, specifically implementing the Noise_XK handshake. Once the
// initial 3-act handshake has completed all messages are encrypted with a
// chacha20 AEAD cipher. On the wire, all messages are prefixed with an
// authenticated+encrypted length field. Additionally, the encrypted+auth'd
// length prefix is used as the AD when encrypting+decryption messages. This
// construction provides confidentiality of packet length, avoids introducing
// a padding-oracle, and binds the encrypted packet length to the packet
// itself.
//
// The acts proceeds the following order (initiator on the left):
//
//        GenActOne()   ->
//                          RecvActOne()
//                      <-  GenActTwo()
//        RecvActTwo()
//        GenActThree() ->
//                          RecvActThree()
//
// This exchange corresponds to the following Noise handshake:
//
//        <- s
//        ...
//        -> e, es
//        <- e, ee
//        -> s, se
type Machine struct {
        sendCipher cipherState
        recvCipher cipherState

        ephemeralGen func() (*btcec.PrivateKey, error)

        handshakeState

        // nextCipherHeader is a static buffer that we'll use to read in the
        // next ciphertext header from the wire. The header is a 2 byte length
        // (of the next ciphertext), followed by a 16 byte MAC.
        nextCipherHeader [encHeaderSize]byte

        // pktLenBuffer is a reusable buffer for encoding the packet length.
        pktLenBuffer [lengthHeaderSize]byte

        // nextHeaderSend holds a reference to the remaining header bytes to
        // write out for a pending message. This allows us to tolerate timeout
        // errors that cause partial writes.
        nextHeaderSend []byte

        // nextBodySend holds a reference to the remaining body bytes to write
        // out for a pending message. This allows us to tolerate timeout errors
        // that cause partial writes.
        nextBodySend []byte

        // pooledHeaderBuf is the pooled buffer used for the header, which we
        // need to track so we can return it to the pool when done.
        pooledHeaderBuf *[]byte

        // pooledBodyBuf is the pooled buffer used for the body, which we need
        // to track so we can return it to the pool when done.
        pooledBodyBuf *[]byte
}

// NewBrontideMachine creates a new instance of the brontide state-machine. If
// the responder (listener) is creating the object, then the remotePub should
// be nil. The handshake state within brontide is initialized using the ascii
// string "lightning" as the prologue. The last parameter is a set of variadic
// arguments for adding additional options to the brontide Machine
// initialization.
func NewBrontideMachine(initiator bool, localKey keychain.SingleKeyECDH,
        remotePub *btcec.PublicKey, options ...func(*Machine)) *Machine {

        handshake := newHandshakeState(
                initiator, lightningPrologue, localKey, remotePub,
        )

        m := &Machine{
                handshakeState: handshake,
                ephemeralGen:   ephemeralGen,
        }

        // With the default options established, we'll now process all the
        // options passed in as parameters.
        for _, option := range options {
                option(m)
        }

        return m
}

const (
        // HandshakeVersion is the expected version of the brontide handshake.
        // Any messages that carry a different version will cause the handshake
        // to abort immediately.
        HandshakeVersion = byte(0)

        // ActOneSize is the size of the packet sent from initiator to
        // responder in ActOne. The packet consists of a handshake version, an
        // ephemeral key in compressed format, and a 16-byte poly1305 tag.
        //
        // 1 + 33 + 16
        ActOneSize = 50

        // ActTwoSize is the size the packet sent from responder to initiator
        // in ActTwo. The packet consists of a handshake version, an ephemeral
        // key in compressed format and a 16-byte poly1305 tag.
        //
        // 1 + 33 + 16
        ActTwoSize = 50

        // ActThreeSize is the size of the packet sent from initiator to
        // responder in ActThree. The packet consists of a handshake version,
        // the initiators static key encrypted with strong forward secrecy and
        // a 16-byte poly1035 tag.
        //
        // 1 + 33 + 16 + 16
        ActThreeSize = 66
)

// GenActOne generates the initial packet (act one) to be sent from initiator
// to responder. During act one the initiator generates a fresh ephemeral key,
// hashes it into the handshake digest, and performs an ECDH between this key
// and the responder's static key. Future payloads are encrypted with a key
// derived from this result.
//
//        -> e, es
func (b *Machine) GenActOne() ([ActOneSize]byte, error) {
        var actOne [ActOneSize]byte

        // e
        localEphemeral, err := b.ephemeralGen()
        if err != nil {
                return actOne, err
        }
        b.localEphemeral = &keychain.PrivKeyECDH{
                PrivKey: localEphemeral,
        }

        ephemeral := localEphemeral.PubKey().SerializeCompressed()
        b.mixHash(ephemeral)

        // es
        s, err := ecdh(b.remoteStatic, b.localEphemeral)
        if err != nil {
                return actOne, err
        }
        b.mixKey(s[:])

        authPayload := b.EncryptAndHash([]byte{})

        actOne[0] = HandshakeVersion
        copy(actOne[1:34], ephemeral)
        copy(actOne[34:], authPayload)

        return actOne, nil
}

// RecvActOne processes the act one packet sent by the initiator. The responder
// executes the mirrored actions to that of the initiator extending the
// handshake digest and deriving a new shared secret based on an ECDH with the
// initiator's ephemeral key and responder's static key.
func (b *Machine) RecvActOne(actOne [ActOneSize]byte) error {
        var (
                err error
                e   [33]byte
                p   [16]byte
        )

        // If the handshake version is unknown, then the handshake fails
        // immediately.
        if actOne[0] != HandshakeVersion {
                return fmt.Errorf("act one: invalid handshake version: %v, "+
                        "only %v is valid, msg=%x", actOne[0], HandshakeVersion,
                        actOne[:])
        }

        copy(e[:], actOne[1:34])
        copy(p[:], actOne[34:])

        // e
        b.remoteEphemeral, err = btcec.ParsePubKey(e[:])
        if err != nil {
                return err
        }
        b.mixHash(b.remoteEphemeral.SerializeCompressed())

        // es
        s, err := ecdh(b.remoteEphemeral, b.localStatic)
        if err != nil {
                return err
        }
        b.mixKey(s)

        // If the initiator doesn't know our static key, then this operation
        // will fail.
        _, err = b.DecryptAndHash(p[:])
        return err
}

// GenActTwo generates the second packet (act two) to be sent from the
// responder to the initiator. The packet for act two is identical to that of
// act one, but then results in a different ECDH operation between the
// initiator's and responder's ephemeral keys.
//
//        <- e, ee
func (b *Machine) GenActTwo() ([ActTwoSize]byte, error) {
        var actTwo [ActTwoSize]byte

        // e
        localEphemeral, err := b.ephemeralGen()
        if err != nil {
                return actTwo, err
        }
        b.localEphemeral = &keychain.PrivKeyECDH{
                PrivKey: localEphemeral,
        }

        ephemeral := localEphemeral.PubKey().SerializeCompressed()
        b.mixHash(localEphemeral.PubKey().SerializeCompressed())

        // ee
        s, err := ecdh(b.remoteEphemeral, b.localEphemeral)
        if err != nil {
                return actTwo, err
        }
        b.mixKey(s)

        authPayload := b.EncryptAndHash([]byte{})

        actTwo[0] = HandshakeVersion
        copy(actTwo[1:34], ephemeral)
        copy(actTwo[34:], authPayload)

        return actTwo, nil
}

// RecvActTwo processes the second packet (act two) sent from the responder to
// the initiator. A successful processing of this packet authenticates the
// initiator to the responder.
func (b *Machine) RecvActTwo(actTwo [ActTwoSize]byte) error {
        var (
                err error
                e   [33]byte
                p   [16]byte
        )

        // If the handshake version is unknown, then the handshake fails
        // immediately.
        if actTwo[0] != HandshakeVersion {
                return fmt.Errorf("act two: invalid handshake version: %v, "+
                        "only %v is valid, msg=%x", actTwo[0], HandshakeVersion,
                        actTwo[:])
        }

        copy(e[:], actTwo[1:34])
        copy(p[:], actTwo[34:])

        // e
        b.remoteEphemeral, err = btcec.ParsePubKey(e[:])
        if err != nil {
                return err
        }
        b.mixHash(b.remoteEphemeral.SerializeCompressed())

        // ee
        s, err := ecdh(b.remoteEphemeral, b.localEphemeral)
        if err != nil {
                return err
        }
        b.mixKey(s)

        _, err = b.DecryptAndHash(p[:])
        return err
}

// GenActThree creates the final (act three) packet of the handshake. Act three
// is to be sent from the initiator to the responder. The purpose of act three
// is to transmit the initiator's public key under strong forward secrecy to
// the responder. This act also includes the final ECDH operation which yields
// the final session.
//
//        -> s, se
func (b *Machine) GenActThree() ([ActThreeSize]byte, error) {
        var actThree [ActThreeSize]byte

        ourPubkey := b.localStatic.PubKey().SerializeCompressed()
        ciphertext := b.EncryptAndHash(ourPubkey)

        s, err := ecdh(b.remoteEphemeral, b.localStatic)
        if err != nil {
                return actThree, err
        }
        b.mixKey(s)

        authPayload := b.EncryptAndHash([]byte{})

        actThree[0] = HandshakeVersion
        copy(actThree[1:50], ciphertext)
        copy(actThree[50:], authPayload)

        // With the final ECDH operation complete, derive the session sending
        // and receiving keys.
        b.split()

        return actThree, nil
}

// RecvActThree processes the final act (act three) sent from the initiator to
// the responder. After processing this act, the responder learns of the
// initiator's static public key. Decryption of the static key serves to
// authenticate the initiator to the responder.
func (b *Machine) RecvActThree(actThree [ActThreeSize]byte) error {
        var (
                err error
                s   [33 + 16]byte
                p   [16]byte
        )

        // If the handshake version is unknown, then the handshake fails
        // immediately.
        if actThree[0] != HandshakeVersion {
                return fmt.Errorf("act three: invalid handshake version: %v, "+
                        "only %v is valid, msg=%x", actThree[0], HandshakeVersion,
                        actThree[:])
        }

        copy(s[:], actThree[1:33+16+1])
        copy(p[:], actThree[33+16+1:])

        // s
        remotePub, err := b.DecryptAndHash(s[:])
        if err != nil {
                return err
        }
        b.remoteStatic, err = btcec.ParsePubKey(remotePub)
        if err != nil {
                return err
        }

        // se
        se, err := ecdh(b.remoteStatic, b.localEphemeral)
        if err != nil {
                return err
        }
        b.mixKey(se)

        if _, err := b.DecryptAndHash(p[:]); err != nil {
                return err
        }

        // With the final ECDH operation complete, derive the session sending
        // and receiving keys.
        b.split()

        return nil
}

// split is the final wrap-up act to be executed at the end of a successful
// three act handshake. This function creates two internal cipherState
// instances: one which is used to encrypt messages from the initiator to the
// responder, and another which is used to encrypt message for the opposite
// direction.
func (b *Machine) split() {
        var (
                empty   []byte
                sendKey [32]byte
                recvKey [32]byte
        )

        h := hkdf.New(sha256.New, empty, b.chainingKey[:], empty)

        // If we're the initiator the first 32 bytes are used to encrypt our
        // messages and the second 32-bytes to decrypt their messages. For the
        // responder the opposite is true.
        if b.initiator {
                h.Read(sendKey[:])
                b.sendCipher = cipherState{}
                b.sendCipher.InitializeKeyWithSalt(b.chainingKey, sendKey)

                h.Read(recvKey[:])
                b.recvCipher = cipherState{}
                b.recvCipher.InitializeKeyWithSalt(b.chainingKey, recvKey)
        } else {
                h.Read(recvKey[:])
                b.recvCipher = cipherState{}
                b.recvCipher.InitializeKeyWithSalt(b.chainingKey, recvKey)

                h.Read(sendKey[:])
                b.sendCipher = cipherState{}
                b.sendCipher.InitializeKeyWithSalt(b.chainingKey, sendKey)
        }
}

// WriteMessage encrypts and buffers the next message p. The ciphertext of the
// message is prepended with an encrypt+auth'd length which must be used as the
// AD to the AEAD construction when being decrypted by the other side.
//
// NOTE: This DOES NOT write the message to the wire, it should be followed by a
// call to Flush to ensure the message is written.
func (b *Machine) WriteMessage(p []byte) error {
        // The total length of each message payload including the MAC size
        // payload exceed the largest number encodable within a 16-bit unsigned
        // integer.
        if len(p) > math.MaxUint16 {
                return ErrMaxMessageLengthExceeded
        }

        // If a prior message was written but it hasn't been fully flushed,
        // return an error as we only support buffering of one message at a
        // time.
        if len(b.nextHeaderSend) > 0 || len(b.nextBodySend) > 0 {
                return ErrMessageNotFlushed
        }

        // The full length of the packet is only the packet length, and does
        // NOT include the MAC.
        fullLength := uint16(len(p))

        binary.BigEndian.PutUint16(b.pktLenBuffer[:], fullLength)

        headerBufInterface := headerBufferPool.Get()
        headerBuf, ok := headerBufInterface.(*[]byte)
        if !ok {
                b.releaseBuffers()
                return fmt.Errorf("headerBufferPool returned unexpected "+
                        "type: %T", headerBufInterface)
        }
        b.pooledHeaderBuf = headerBuf

        bodyBufInterface := bodyBufferPool.Get()
        bodyBuf, ok := bodyBufInterface.(*[]byte)
        if !ok {
                b.releaseBuffers()
                return fmt.Errorf("bodyBufferPool returned unexpected "+
                        "type: %T", bodyBufInterface)
        }
        b.pooledBodyBuf = bodyBuf

        // First, generate the encrypted+MAC'd length prefix for the packet. We
        // pass our pooled buffer as the cipherText (dst) parameter.
        b.nextHeaderSend = b.sendCipher.Encrypt(
                nil, *b.pooledHeaderBuf, b.pktLenBuffer[:],
        )

        // Finally, generate the encrypted packet itself. We pass our pooled
        // buffer as the cipherText (dst) parameter.
        b.nextBodySend = b.sendCipher.Encrypt(nil, *b.pooledBodyBuf, p)

        return nil
}

// Flush attempts to write a message buffered using WriteMessage to the provided
// io.Writer. If no buffered message exists, this will result in a NOP.
// Otherwise, it will continue to write the remaining bytes, picking up where
// the byte stream left off in the event of a partial write. The number of bytes
// returned reflects the number of plaintext bytes in the payload, and does not
// account for the overhead of the header or MACs.
//
// NOTE: It is safe to call this method again iff a timeout error is returned.
func (b *Machine) Flush(w io.Writer) (int, error) {
        // First, write out the pending header bytes, if any exist. Any header
        // bytes written will not count towards the total amount flushed.
        if len(b.nextHeaderSend) > 0 {
                // Write any remaining header bytes and shift the slice to point
                // to the next segment of unwritten bytes. If an error is
                // encountered, we can continue to write the header from where
                // we left off on a subsequent call to Flush.
                n, err := w.Write(b.nextHeaderSend)
                b.nextHeaderSend = b.nextHeaderSend[n:]
                if err != nil {
                        return 0, err
                }
        }

        // Next, write the pending body bytes, if any exist. Only the number of
        // bytes written that correspond to the ciphertext will be included in
        // the total bytes written, bytes written as part of the MAC will not be
        // counted.
        var nn int
        if len(b.nextBodySend) > 0 {
                // Write out all bytes excluding the mac and shift the body
                // slice depending on the number of actual bytes written.
                n, err := w.Write(b.nextBodySend)
                b.nextBodySend = b.nextBodySend[n:]

                // If we partially or fully wrote any of the body's MAC, we'll
                // subtract that contribution from the total amount flushed to
                // preserve the abstraction of returning the number of plaintext
                // bytes written by the connection.
                //
                // There are three possible scenarios we must handle to ensure
                // the returned value is correct. In the first case, the write
                // straddles both payload and MAC bytes, and we must subtract
                // the number of MAC bytes written from n. In the second, only
                // payload bytes are written, thus we can return n unmodified.
                // The final scenario pertains to the case where only MAC bytes
                // are written, none of which count towards the total.
                //
                //                 |-----------Payload------------|----MAC----|
                // Straddle:       S---------------------------------E--------0
                // Payload-only:   S------------------------E-----------------0
                // MAC-only:                                        S-------E-0
                start, end := n+len(b.nextBodySend), len(b.nextBodySend)
                switch {

                // Straddles payload and MAC bytes, subtract number of MAC bytes
                // written from the actual number written.
                case start > macSize && end <= macSize:
                        nn = n - (macSize - end)

                // Only payload bytes are written, return n directly.
                case start > macSize && end > macSize:
                        nn = n

                // Only MAC bytes are written, return 0 bytes written.
                default:
                }

                if err != nil {
                        return nn, err
                }
        }

        // If both header and body have been fully flushed, release the pooled
        // buffers back to their pools.
        if len(b.nextHeaderSend) == 0 && len(b.nextBodySend) == 0 {
                b.releaseBuffers()
        }

        return nn, nil
}

// releaseBuffers returns the pooled buffers back to their respective pools
// and clears the references.
func (b *Machine) releaseBuffers() {
        if b.pooledHeaderBuf != nil {
                *b.pooledHeaderBuf = (*b.pooledHeaderBuf)[:0]
                headerBufferPool.Put(b.pooledHeaderBuf)
                b.pooledHeaderBuf = nil
        }

        if b.pooledBodyBuf != nil {
                *b.pooledBodyBuf = (*b.pooledBodyBuf)[:0]
                bodyBufferPool.Put(b.pooledBodyBuf)
                b.pooledBodyBuf = nil
        }

        b.nextHeaderSend = nil
        b.nextBodySend = nil
}

// ReadMessage attempts to read the next message from the passed io.Reader. In
// the case of an authentication error, a non-nil error is returned.
func (b *Machine) ReadMessage(r io.Reader) ([]byte, error) {
        pktLen, err := b.ReadHeader(r)
        if err != nil {
                return nil, err
        }

        buf := make([]byte, pktLen)
        return b.ReadBody(r, buf)
}

// ReadHeader attempts to read the next message header from the passed
// io.Reader. The header contains the length of the next body including
// additional overhead of the MAC. In the case of an authentication error, a
// non-nil error is returned.
//
// NOTE: This method SHOULD NOT be used in the case that the io.Reader may be
// adversarial and induce long delays. If the caller needs to set read deadlines
// appropriately, it is preferred that they use the split ReadHeader and
// ReadBody methods so that the deadlines can be set appropriately on each.
func (b *Machine) ReadHeader(r io.Reader) (uint32, error) {
        _, err := io.ReadFull(r, b.nextCipherHeader[:])
        if err != nil {
                return 0, err
        }

        // Attempt to decrypt+auth the packet length present in the stream.
        //
        // By passing in `nextCipherHeader` as the destination, we avoid making
        // the library allocate a new buffer to decode the plaintext.
        pktLenBytes, err := b.recvCipher.Decrypt(
                nil, b.nextCipherHeader[:0], b.nextCipherHeader[:],
        )
        if err != nil {
                return 0, err
        }

        // Compute the packet length that we will need to read off the wire.
        pktLen := uint32(binary.BigEndian.Uint16(pktLenBytes)) + macSize

        return pktLen, nil
}

// ReadBody attempts to ready the next message body from the passed io.Reader.
// The provided buffer MUST be the length indicated by the packet length
// returned by the preceding call to ReadHeader. In the case of an
// authentication error, a non-nil error is returned.
func (b *Machine) ReadBody(r io.Reader, buf []byte) ([]byte, error) {
        // Next, using the length read from the packet header, read the
        // encrypted packet itself into the buffer allocated by the read
        // pool.
        _, err := io.ReadFull(r, buf)
        if err != nil {
                return nil, err
        }

        // Finally, decrypt the message held in the buffer, and return a new
        // byte slice containing the plaintext.
        //
        // By passing in the buf (the ciphertext) as the first argument, we end
        // up re-using it as we don't force the library to allocate a new
        // buffer to decode the plaintext.
        return b.recvCipher.Decrypt(nil, buf[:0], buf)
}

lightningnetwork / lnd / 17991712134

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous