Repository: grpc/grpc-go
File: credentials/alts/internal/conn/record.go (Go)

Code before:
				tmp := mem.DefaultBufferPool().Get(int(length))
				copy(*tmp, p.protected)
				p.protected = (*tmp)[:len(p.protected)]
				mem.DefaultBufferPool().Put(p.protectedPointer)
				p.protectedPointer = tmp
			}
			n, err = p.Conn.Read(p.protected[len(p.protected):cap(p.protected)])
			if err != nil {
				return 0, err
			}
			p.protected = p.protected[:len(p.protected)+n]
			framedMsg, p.nextFrame, err = ParseFramedMsg(p.protected, altsRecordLengthLimit)
			if err != nil {
				return 0, err
			}
		}
		// Now we have a complete frame, decrypted it.
		msg := framedMsg[MsgLenFieldSize:]
		msgType := binary.LittleEndian.Uint32(msg[:msgTypeFieldSize])
		if msgType&0xff != altsRecordMsgType {
			return 0, fmt.Errorf("received frame with incorrect message type %v, expected lower byte %v",
				msgType, altsRecordMsgType)
		}
		ciphertext := msg[msgTypeFieldSize:]

		// Decrypt directly into the buffer, avoiding a copy to p.buf if
		// possible.
		if cap(b) >= len(msg) {
			dec, err := p.crypto.Decrypt(b[:0], ciphertext)
			if err != nil {
				return 0, err
			}
			return len(dec), nil
		}
		// Resize the read buffer if needed to hold the entire decrypted
		// frame.
		if cap(*p.bufPointer) < len(ciphertext) {
			mem.DefaultBufferPool().Put(p.bufPointer)
			p.bufPointer = mem.DefaultBufferPool().Get(len(ciphertext))
		}
		p.buf = *p.bufPointer
		dec, err := p.crypto.Decrypt(p.buf[:0], ciphertext)
		if err != nil {
			return 0, err
		}
		p.buf = p.buf[:len(dec)]
	}

	n = copy(b, p.buf)
	p.buf = p.buf[n:]
	return n, nil

Review comment (suggestion):
Is there any potential issue here where we are missing the update of `p.buf` that is guaranteed on the non-error return in the old structure? Specifically `p.buf = p.buf[n:]`

Code after:
		ciphertext := msg[msgTypeFieldSize:]

		// Decrypt directly into the buffer, avoiding a copy from p.buf if
		// possible.
		if len(b) >= len(ciphertext) {
			dec, err := p.crypto.Decrypt(b[:0], ciphertext)
			if err != nil {
				return 0, err
			}
			return len(dec), nil
		}
		// Decrypt requires that if the dst and ciphertext alias, they
		// must alias exactly. Code here used to use msg[:0], but msg
		// starts MsgLenFieldSize+msgTypeFieldSize bytes earlier than
		// ciphertext, so they alias inexactly. Using ciphertext[:0]
		// arranges the appropriate aliasing without needing to copy
		// ciphertext or use a separate destination buffer. For more info
		// check: https://golang.org/pkg/crypto/cipher/#AEAD.
		p.buf, err = p.crypto.Decrypt(ciphertext[:0], ciphertext)
		if err != nil {
			return 0, err
		}
	}

	n = copy(b, p.buf)
	p.buf = p.buf[n:]
	return n, nil
}

// Write encrypts, frames, and writes bytes from b to the underlying connection.
func (p *conn) Write(b []byte) (n int, err error) {
	n = len(b)
	// Calculate the output buffer size with framing and encryption overhead.
	numOfFrames := int(math.Ceil(float64(len(b)) / float64(p.payloadLengthLimit)))
	size := len(b) + numOfFrames*p.overhead
	// If writeBuf is too small, increase its size up to the maximum size.
	partialBSize := len(b)
	if size > altsWriteBufferMaxSize {
		size = altsWriteBufferMaxSize
		const numOfFramesInMaxWriteBuf = altsWriteBufferMaxSize / altsRecordDefaultLength
		partialBSize = numOfFramesInMaxWriteBuf * p.payloadLengthLimit
	}
	if len(p.writeBuf) < size {
		p.writeBuf = make([]byte, size)
	}

	for partialBStart := 0; partialBStart < len(b); partialBStart += partialBSize {
		partialBEnd := partialBStart + partialBSize
		if partialBEnd > len(b) {
			partialBEnd = len(b)
		}