github.com/metacubex/quic-go@v0.44.1-0.20240520163451-20b689a59136/testutils/testutils.go (about)

     1  // Package testutils contains utilities for simulating packet injection and man-in-the-middle (MITM) attacker tests.
     2  // It is not supposed to be used for non-testing purposes.
     3  // The API is not guaranteed to be stable.
     4  package testutils
     5  
     6  import (
     7  	"fmt"
     8  
     9  	"github.com/metacubex/quic-go/internal/handshake"
    10  	"github.com/metacubex/quic-go/internal/protocol"
    11  	"github.com/metacubex/quic-go/internal/wire"
    12  )
    13  
    14  // writePacket returns a new raw packet with the specified header and payload
    15  func writePacket(hdr *wire.ExtendedHeader, data []byte) []byte {
    16  	b, err := hdr.Append(nil, hdr.Version)
    17  	if err != nil {
    18  		panic(fmt.Sprintf("failed to write header: %s", err))
    19  	}
    20  	return append(b, data...)
    21  }
    22  
    23  // packRawPayload returns a new raw payload containing given frames
    24  func packRawPayload(version protocol.Version, frames []wire.Frame) []byte {
    25  	var b []byte
    26  	for _, cf := range frames {
    27  		var err error
    28  		b, err = cf.Append(b, version)
    29  		if err != nil {
    30  			panic(err)
    31  		}
    32  	}
    33  	return b
    34  }
    35  
    36  // ComposeInitialPacket returns an Initial packet encrypted under key (the original destination connection ID)
    37  // containing specified frames.
    38  func ComposeInitialPacket(
    39  	srcConnID, destConnID, key protocol.ConnectionID,
    40  	token []byte,
    41  	frames []wire.Frame,
    42  	sentBy protocol.Perspective,
    43  	version protocol.Version,
    44  ) []byte {
    45  	sealer, _ := handshake.NewInitialAEAD(key, sentBy, version)
    46  
    47  	// compose payload
    48  	var payload []byte
    49  	if len(frames) == 0 {
    50  		payload = make([]byte, protocol.MinInitialPacketSize)
    51  	} else {
    52  		payload = packRawPayload(version, frames)
    53  	}
    54  
    55  	// compose Initial header
    56  	payloadSize := len(payload)
    57  	const pnLength = protocol.PacketNumberLen4
    58  	length := payloadSize + int(pnLength) + sealer.Overhead()
    59  	hdr := &wire.ExtendedHeader{
    60  		Header: wire.Header{
    61  			Type:             protocol.PacketTypeInitial,
    62  			Token:            token,
    63  			SrcConnectionID:  srcConnID,
    64  			DestConnectionID: destConnID,
    65  			Length:           protocol.ByteCount(length),
    66  			Version:          version,
    67  		},
    68  		PacketNumberLen: pnLength,
    69  		PacketNumber:    0x0,
    70  	}
    71  
    72  	raw := writePacket(hdr, payload)
    73  
    74  	// encrypt payload and header
    75  	payloadOffset := len(raw) - payloadSize
    76  	var encrypted []byte
    77  	encrypted = sealer.Seal(encrypted, payload, hdr.PacketNumber, raw[:payloadOffset])
    78  	hdrBytes := raw[0:payloadOffset]
    79  	encrypted = append(hdrBytes, encrypted...)
    80  	pnOffset := payloadOffset - int(pnLength) // packet number offset
    81  	sealer.EncryptHeader(
    82  		encrypted[payloadOffset:payloadOffset+16], // first 16 bytes of payload (sample)
    83  		&encrypted[0],                     // first byte of header
    84  		encrypted[pnOffset:payloadOffset], // packet number bytes
    85  	)
    86  	return encrypted
    87  }
    88  
    89  // ComposeRetryPacket returns a new raw Retry Packet
    90  func ComposeRetryPacket(
    91  	srcConnID protocol.ConnectionID,
    92  	destConnID protocol.ConnectionID,
    93  	origDestConnID protocol.ConnectionID,
    94  	token []byte,
    95  	version protocol.Version,
    96  ) []byte {
    97  	hdr := &wire.ExtendedHeader{
    98  		Header: wire.Header{
    99  			Type:             protocol.PacketTypeRetry,
   100  			SrcConnectionID:  srcConnID,
   101  			DestConnectionID: destConnID,
   102  			Token:            token,
   103  			Version:          version,
   104  		},
   105  	}
   106  	data := writePacket(hdr, nil)
   107  	return append(data, handshake.GetRetryIntegrityTag(data, origDestConnID, version)[:]...)
   108  }