github.com/danielpfeifer02/quic-go-prio-packs@v0.41.0-28/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/danielpfeifer02/quic-go-prio-packs/internal/handshake"
    10  	"github.com/danielpfeifer02/quic-go-prio-packs/internal/protocol"
    11  	"github.com/danielpfeifer02/quic-go-prio-packs/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  	frames []wire.Frame,
    41  	sentBy protocol.Perspective,
    42  	version protocol.Version,
    43  ) []byte {
    44  	sealer, _ := handshake.NewInitialAEAD(key, sentBy, version)
    45  
    46  	// compose payload
    47  	var payload []byte
    48  	if len(frames) == 0 {
    49  		payload = make([]byte, protocol.MinInitialPacketSize)
    50  	} else {
    51  		payload = packRawPayload(version, frames)
    52  	}
    53  
    54  	// compose Initial header
    55  	payloadSize := len(payload)
    56  	const pnLength = protocol.PacketNumberLen4
    57  	length := payloadSize + int(pnLength) + sealer.Overhead()
    58  	hdr := &wire.ExtendedHeader{
    59  		Header: wire.Header{
    60  			Type:             protocol.PacketTypeInitial,
    61  			SrcConnectionID:  srcConnID,
    62  			DestConnectionID: destConnID,
    63  			Length:           protocol.ByteCount(length),
    64  			Version:          version,
    65  		},
    66  		PacketNumberLen: pnLength,
    67  		PacketNumber:    0x0,
    68  	}
    69  
    70  	raw := writePacket(hdr, payload)
    71  
    72  	// encrypt payload and header
    73  	payloadOffset := len(raw) - payloadSize
    74  	var encrypted []byte
    75  	encrypted = sealer.Seal(encrypted, payload, hdr.PacketNumber, raw[:payloadOffset])
    76  	hdrBytes := raw[0:payloadOffset]
    77  	encrypted = append(hdrBytes, encrypted...)
    78  	pnOffset := payloadOffset - int(pnLength) // packet number offset
    79  	sealer.EncryptHeader(
    80  		encrypted[payloadOffset:payloadOffset+16], // first 16 bytes of payload (sample)
    81  		&encrypted[0],                     // first byte of header
    82  		encrypted[pnOffset:payloadOffset], // packet number bytes
    83  	)
    84  	return encrypted
    85  }
    86  
    87  // ComposeRetryPacket returns a new raw Retry Packet
    88  func ComposeRetryPacket(
    89  	srcConnID protocol.ConnectionID,
    90  	destConnID protocol.ConnectionID,
    91  	origDestConnID protocol.ConnectionID,
    92  	token []byte,
    93  	version protocol.Version,
    94  ) []byte {
    95  	hdr := &wire.ExtendedHeader{
    96  		Header: wire.Header{
    97  			Type:             protocol.PacketTypeRetry,
    98  			SrcConnectionID:  srcConnID,
    99  			DestConnectionID: destConnID,
   100  			Token:            token,
   101  			Version:          version,
   102  		},
   103  	}
   104  	data := writePacket(hdr, nil)
   105  	return append(data, handshake.GetRetryIntegrityTag(data, origDestConnID, version)[:]...)
   106  }