github.com/sagernet/quic-go@v0.43.1-beta.1/ech/crypto_stream_manager.go (about)

     1  package quic
     2  
     3  import (
     4  	"fmt"
     5  
     6  	"github.com/sagernet/quic-go/internal/handshake_ech"
     7  	"github.com/sagernet/quic-go/internal/protocol"
     8  	"github.com/sagernet/quic-go/internal/wire"
     9  )
    10  
    11  type cryptoDataHandler interface {
    12  	HandleMessage([]byte, protocol.EncryptionLevel) error
    13  	NextEvent() handshake.Event
    14  }
    15  
    16  type cryptoStreamManager struct {
    17  	cryptoHandler cryptoDataHandler
    18  
    19  	initialStream   cryptoStream
    20  	handshakeStream cryptoStream
    21  	oneRTTStream    cryptoStream
    22  }
    23  
    24  func newCryptoStreamManager(
    25  	cryptoHandler cryptoDataHandler,
    26  	initialStream cryptoStream,
    27  	handshakeStream cryptoStream,
    28  	oneRTTStream cryptoStream,
    29  ) *cryptoStreamManager {
    30  	return &cryptoStreamManager{
    31  		cryptoHandler:   cryptoHandler,
    32  		initialStream:   initialStream,
    33  		handshakeStream: handshakeStream,
    34  		oneRTTStream:    oneRTTStream,
    35  	}
    36  }
    37  
    38  func (m *cryptoStreamManager) HandleCryptoFrame(frame *wire.CryptoFrame, encLevel protocol.EncryptionLevel) error {
    39  	var str cryptoStream
    40  	//nolint:exhaustive // CRYPTO frames cannot be sent in 0-RTT packets.
    41  	switch encLevel {
    42  	case protocol.EncryptionInitial:
    43  		str = m.initialStream
    44  	case protocol.EncryptionHandshake:
    45  		str = m.handshakeStream
    46  	case protocol.Encryption1RTT:
    47  		str = m.oneRTTStream
    48  	default:
    49  		return fmt.Errorf("received CRYPTO frame with unexpected encryption level: %s", encLevel)
    50  	}
    51  	if err := str.HandleCryptoFrame(frame); err != nil {
    52  		return err
    53  	}
    54  	for {
    55  		data := str.GetCryptoData()
    56  		if data == nil {
    57  			return nil
    58  		}
    59  		if err := m.cryptoHandler.HandleMessage(data, encLevel); err != nil {
    60  			return err
    61  		}
    62  	}
    63  }
    64  
    65  func (m *cryptoStreamManager) GetPostHandshakeData(maxSize protocol.ByteCount) *wire.CryptoFrame {
    66  	if !m.oneRTTStream.HasData() {
    67  		return nil
    68  	}
    69  	return m.oneRTTStream.PopCryptoFrame(maxSize)
    70  }
    71  
    72  func (m *cryptoStreamManager) Drop(encLevel protocol.EncryptionLevel) error {
    73  	//nolint:exhaustive // 1-RTT keys should never get dropped.
    74  	switch encLevel {
    75  	case protocol.EncryptionInitial:
    76  		return m.initialStream.Finish()
    77  	case protocol.EncryptionHandshake:
    78  		return m.handshakeStream.Finish()
    79  	default:
    80  		panic(fmt.Sprintf("dropped unexpected encryption level: %s", encLevel))
    81  	}
    82  }