snapshot of https://github.com/grpc/grpc-go.git (master, 2026-09-15; blobs > 131072B omitted)

credentials/alts/internal/conn/record.go

@@ -0,0 +1,380 @@
+/*
+ *
+ * Copyright 2018 gRPC authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ */
+
+// Package conn contains an implementation of a secure channel created by gRPC
+// handshakers.
+package conn
+
+import (
+	"encoding/binary"
+	"fmt"
+	"math"
+	"net"
+
+	core "google.golang.org/grpc/credentials/alts/internal"
+	imem "google.golang.org/grpc/internal/mem"
+	"google.golang.org/grpc/internal/transport/readyreader"
+	"google.golang.org/grpc/mem"
+)
+
+// ALTSRecordCrypto is the interface for gRPC ALTS record protocol.
+type ALTSRecordCrypto interface {
+	// Encrypt encrypts the plaintext, computes the tag (if any) of dst and
+	// plaintext, and appends the result to dst, returning the updated slice.
+	// dst and plaintext may fully overlap or not at all.
+	Encrypt(dst, plaintext []byte) ([]byte, error)
+	// EncryptionOverhead returns the tag size (if any) in bytes.
+	EncryptionOverhead() int
+	// Decrypt decrypts ciphertext and verifies the tag (if any). If successful,
+	// this function appends the resulting plaintext to dst, returning the
+	// updated slice. dst and ciphertext may alias exactly or not at all. To
+	// reuse ciphertext's storage for the decrypted output, use ciphertext[:0]
+	// as dst. Even if the function fails, the contents of dst, up to its
+	// capacity, may be overwritten.
+	Decrypt(dst, ciphertext []byte) ([]byte, error)
+}
+
+// ALTSRecordFunc is a function type for factory functions that create
+// ALTSRecordCrypto instances.
+type ALTSRecordFunc func(s core.Side, keyData []byte) (ALTSRecordCrypto, error)
+
+const (
+	// MsgLenFieldSize is the byte size of the frame length field of a
+	// framed message.
+	MsgLenFieldSize = 4
+	// The byte size of the message type field of a framed message.
+	msgTypeFieldSize = 4
+	// The bytes size limit for a ALTS record message.
+	altsRecordLengthLimit = 1024 * 1024 // 1 MiB
+	// The default bytes size of a ALTS record message.
+	altsRecordDefaultLength = 4 * 1024 // 4KiB
+	// Message type value included in ALTS record framing.
+	altsRecordMsgType = uint32(0x06)
+	// The maximum write buffer size, which can hold multiple ALTS frames.
+	altsWriteBufferMaxSize = 512 * 1024 // 512KiB
+	// The initial buffer used to read from the network.
+	// It includes an additional 512 Bytes to hold two 16KiB records plus
+	// small framing overheads.
+	altsReadBufferInitialSize = 32*1024 + 512 // 32.5KiB
+)
+
+var (
+	protocols    = make(map[string]ALTSRecordFunc)
+	writeBufPool *imem.BinaryTieredBufferPool
+	// readBufPool pools buffers of at least `altsReadBufferInitialSize` size.
+	// Since the read buffer size is slightly larger than 32KB, using a regular
+	// BinaryTieredBufferPool results in allocating buffers of almost double the
+	// required length.
+	readBufPool = imem.NewDirtySimplePool()
+
+	// Compile-time check to ensure conn implements ReadyReader.
+	_ readyreader.Reader = &conn{}
+)
+
+func init() {
+	pool, err := imem.NewDirtyBinaryTieredBufferPool(
+		8,
+		12, // Go page size, 4KB
+		14, // 16KB (max HTTP/2 frame size used by gRPC)
+		15, // 32KB (default buffer size for gRPC)
+		16, // 64KB
+		17, // 128KB
+		19, // 512KB, max write buffer size
+	)
+	if err != nil {
+		panic(fmt.Sprintf("Failed to create write buffer pool: %v", err))
+	}
+	writeBufPool = pool
+}
+
+// RegisterProtocol register a ALTS record encryption protocol.
+func RegisterProtocol(protocol string, f ALTSRecordFunc) error {
+	if _, ok := protocols[protocol]; ok {
+		return fmt.Errorf("protocol %v is already registered", protocol)
+	}
+	protocols[protocol] = f
+	return nil
+}
+
+// conn represents a secured connection. It implements the net.Conn interface.
+type conn struct {
+	net.Conn
+	reader readyreader.Reader
+	crypto ALTSRecordCrypto
+	// buf holds data that has been read from the connection and decrypted,
+	// but has not yet been returned by Read. It is a sub-slice of protected.
+	buf                []byte
+	payloadLengthLimit int
+	// protectedHandle buffer holds data read from the network but have not yet
+	// been decrypted. This data might not compose a complete frame.
+	//
+	// The buffer pointer to points to a buffer from the readBufPool. The handle
+	// should only be returned to the pool once nextFrame and buf are empty.
+	protectedHandle *[]byte
+	// nextFrame stores the next frame (in protected buffer) info.
+	nextFrame []byte
+	// overhead is the calculated overhead of each frame.
+	overhead  int
+	constPool constBufferPool // stored as a field to avoid heap allocations.
+}
+
+// NewConn creates a new secure channel instance given the other party role and
+// handshaking result.
+func NewConn(c net.Conn, side core.Side, recordProtocol string, key []byte, protected []byte) (net.Conn, error) {
+	return NewConnWithMaxFrameSize(c, side, recordProtocol, key, protected, 0)
+}
+
+// NewConnWithMaxFrameSize creates a new secure channel instance given the
+// other party role, handshaking result, and negotiated maximum frame size.
+func NewConnWithMaxFrameSize(c net.Conn, side core.Side, recordProtocol string, key []byte, protected []byte, negotiatedMaxFrameSize int) (net.Conn, error) {
+	newCrypto := protocols[recordProtocol]
+	if newCrypto == nil {
+		return nil, fmt.Errorf("negotiated unknown next_protocol %q", recordProtocol)
+	}
+	crypto, err := newCrypto(side, key)
+	if err != nil {
+		return nil, fmt.Errorf("protocol %q: %v", recordProtocol, err)
+	}
+	overhead := MsgLenFieldSize + msgTypeFieldSize + crypto.EncryptionOverhead()
+
+	// Clamp maxRecordLen to be at least altsRecordDefaultLength.
+	maxRecordLen := max(altsRecordDefaultLength, negotiatedMaxFrameSize)
+	payloadLengthLimit := maxRecordLen - overhead
+	// We pre-allocate protected to be of size 32KB during initialization.
+	// We increase the size of the buffer by the required amount if it can't
+	// hold a complete encrypted record.
+	protectedHandle := readBufPool.Get(max(altsReadBufferInitialSize, len(protected)))
+	protectedBuf := *protectedHandle
+	// Copy additional data from hanshaker service.
+	copy(protectedBuf, protected)
+	protectedBuf = protectedBuf[:len(protected)]
+
+	altsConn := &conn{
+		Conn:               c,
+		reader:             readyreader.New(c),
+		crypto:             crypto,
+		payloadLengthLimit: payloadLengthLimit,
+		protectedHandle:    protectedHandle,
+		nextFrame:          protectedBuf,
+		overhead:           overhead,
+	}
+	return altsConn, nil
+}
+
+type constBufferPool struct {
+	buffer []byte
+}
+
+func (p *constBufferPool) Get(int) *[]byte {
+	return &p.buffer
+}
+
+func (p *constBufferPool) Put(*[]byte) {}
+
+// Read reads and decrypts a frame from the underlying connection, and copies the
+// decrypted payload into b. If the size of the payload is greater than len(b),
+// Read retains the remaining bytes in an internal buffer, and subsequent calls
+// to Read will read from this buffer until it is exhausted.
+func (p *conn) Read(b []byte) (n int, err error) {
+	p.constPool.buffer = b
+	_, n, err = p.ReadOnReady(len(b), &p.constPool)
+	return n, err
+}
+
+func (p *conn) ReadOnReady(bufSize int, pool mem.BufferPool) (*[]byte, int, error) {
+	if len(p.buf) == 0 {
+		var framedMsg []byte
+		var protected []byte
+		if p.protectedHandle != nil {
+			protected = *p.protectedHandle
+			protected = protected[:cap(protected)]
+		}
+		var err error
+		framedMsg, p.nextFrame, err = ParseFramedMsg(p.nextFrame, altsRecordLengthLimit)
+		if err != nil {
+			return nil, 0, err
+		}
+		// Check whether the next frame to be decrypted has been
+		// completely received yet.
+		if len(framedMsg) == 0 {
+			copy(protected, p.nextFrame)
+			protected = protected[:len(p.nextFrame)]
+			// Always copy next incomplete frame to the beginning of
+			// the protected buffer and reset nextFrame to it.
+			p.nextFrame = protected
+		}
+		// Check whether a complete frame has been received yet.
+		for len(framedMsg) == 0 {
+			if p.protectedHandle != nil && len(protected) == cap(protected) {
+				// We can parse the length header to know exactly how large
+				// the buffer needs to be to hold the entire frame.
+				length, didParse := parseMessageLength(protected)
+				if !didParse {
+					// The protected buffer is initialized with a capacity of
+					// larger than 4B. It should always hold the message length
+					// header.
+					panic(fmt.Sprintf("protected buffer length shorter than expected: %d vs %d", len(protected), MsgLenFieldSize))
+				}
+				oldProtectedBuf := protected
+				oldBufHandle := p.protectedHandle
+				// The new buffer must be able to hold the message length header
+				// and the entire message.
+				requiredCapacity := int(length) + MsgLenFieldSize
+				p.protectedHandle = readBufPool.Get(requiredCapacity)
+				protected = *p.protectedHandle
+				// Copy the contents of the old buffer and set the length of the
+				// new buffer to the number of bytes already read.
+				copy(protected, oldProtectedBuf)
+				protected = protected[:len(oldProtectedBuf)]
+				readBufPool.Put(oldBufHandle)
+			}
+			if p.protectedHandle == nil {
+				// Connection was idle, need to re-allocate the read buffer.
+				newBuf, nRead, err := p.reader.ReadOnReady(altsReadBufferInitialSize, readBufPool)
+				if err != nil {
+					return nil, 0, err
+				}
+				p.protectedHandle = newBuf
+				protected = (*newBuf)[:nRead]
+			} else {
+				nRead, err := p.Conn.Read(protected[len(protected):cap(protected)])
+				if err != nil {
+					return nil, 0, err
+				}
+				protected = protected[:len(protected)+nRead]
+			}
+			framedMsg, p.nextFrame, err = ParseFramedMsg(protected, altsRecordLengthLimit)
+			if err != nil {
+				return nil, 0, err
+			}
+		}
+		// Now we have a complete frame, decrypted it.
+		msg := framedMsg[MsgLenFieldSize:]
+		if len(msg) < msgTypeFieldSize {
+			return nil, 0, fmt.Errorf("received frame with size %v which is shorter than message type field size %v", len(msg), msgTypeFieldSize)
+		}
+		msgType := binary.LittleEndian.Uint32(msg[:msgTypeFieldSize])
+		if msgType&0xff != altsRecordMsgType {
+			return nil, 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 from p.buf if
+		// possible.
+		if bufSize >= len(ciphertext) {
+			allocatedBuf := pool.Get(bufSize)
+			dec, err := p.crypto.Decrypt((*allocatedBuf)[:0], ciphertext)
+			if err != nil {
+				pool.Put(allocatedBuf)
+				return nil, 0, err
+			}
+			p.dropProtectedIfEmtpy()
+			return allocatedBuf, 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 nil, 0, err
+		}
+	}
+
+	allocatedBuf := pool.Get(bufSize)
+	n := copy(*allocatedBuf, p.buf)
+	p.buf = p.buf[n:]
+	p.dropProtectedIfEmtpy()
+	return allocatedBuf, n, nil
+}
+
+func (p *conn) dropProtectedIfEmtpy() {
+	if len(p.buf) > 0 || len(p.nextFrame) > 0 {
+		return
+	}
+	// Potentially idle connection, release the read buffer.
+	p.nextFrame = nil
+	p.buf = nil
+	if p.protectedHandle != nil {
+		readBufPool.Put(p.protectedHandle)
+		p.protectedHandle = 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
+	partialBSize := len(b)
+	if size > altsWriteBufferMaxSize {
+		size = altsWriteBufferMaxSize
+		frameLen := p.payloadLengthLimit + p.overhead
+		numOfFramesInMaxWriteBuf := altsWriteBufferMaxSize / frameLen
+		partialBSize = numOfFramesInMaxWriteBuf * p.payloadLengthLimit
+	}
+	// Get a writeBuf of the required length.
+	bufHandle := writeBufPool.Get(size)
+	defer writeBufPool.Put(bufHandle)
+	writeBuf := *bufHandle
+
+	for partialBStart := 0; partialBStart < len(b); partialBStart += partialBSize {
+		partialBEnd := min(partialBStart+partialBSize, len(b))
+		partialB := b[partialBStart:partialBEnd]
+		writeBufIndex := 0
+		for len(partialB) > 0 {
+			payloadLen := min(len(partialB), p.payloadLengthLimit)
+			buf := partialB[:payloadLen]
+			partialB = partialB[payloadLen:]
+
+			// Write buffer contains: length, type, payload, and tag
+			// if any.
+
+			// 1. Fill in type field.
+			msg := writeBuf[writeBufIndex+MsgLenFieldSize:]
+			binary.LittleEndian.PutUint32(msg, altsRecordMsgType)
+
+			// 2. Encrypt the payload and create a tag if any.
+			msg, err = p.crypto.Encrypt(msg[:msgTypeFieldSize], buf)
+			if err != nil {
+				return n, err
+			}
+
+			// 3. Fill in the size field.
+			binary.LittleEndian.PutUint32(writeBuf[writeBufIndex:], uint32(len(msg)))
+
+			// 4. Increase writeBufIndex.
+			writeBufIndex += len(buf) + p.overhead
+		}
+		nn, err := p.Conn.Write(writeBuf[:writeBufIndex])
+		if err != nil {
+			// We need to calculate the actual data size that was
+			// written. This means we need to remove header,
+			// encryption overheads, and any partially-written
+			// frame data.
+			numOfWrittenFrames := int(math.Floor(float64(nn) / float64(p.payloadLengthLimit+p.overhead)))
+			return partialBStart + numOfWrittenFrames*p.payloadLengthLimit, err
+		}
+	}
+	return n, nil
+}