github.com/mikelsr/quic-go@v0.36.1-0.20230701132136-1d9415b66898/internal/testutils/testutils.go (about)

     1  package testutils
     2  
     3  import (
     4  	"fmt"
     5  
     6  	"github.com/mikelsr/quic-go/internal/handshake"
     7  	"github.com/mikelsr/quic-go/internal/protocol"
     8  	"github.com/mikelsr/quic-go/internal/wire"
     9  )
    10  
    11  // Utilities for simulating packet injection and man-in-the-middle (MITM) attacker tests.
    12  // Do not use for non-testing purposes.
    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.VersionNumber, 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
    37  // (the original destination connection ID) containing specified frames
    38  func ComposeInitialPacket(srcConnID protocol.ConnectionID, destConnID protocol.ConnectionID, version protocol.VersionNumber, key protocol.ConnectionID, frames []wire.Frame) []byte {
    39  	sealer, _ := handshake.NewInitialAEAD(key, protocol.PerspectiveServer, version)
    40  
    41  	// compose payload
    42  	var payload []byte
    43  	if len(frames) == 0 {
    44  		payload = make([]byte, protocol.MinInitialPacketSize)
    45  	} else {
    46  		payload = packRawPayload(version, frames)
    47  	}
    48  
    49  	// compose Initial header
    50  	payloadSize := len(payload)
    51  	pnLength := protocol.PacketNumberLen4
    52  	length := payloadSize + int(pnLength) + sealer.Overhead()
    53  	hdr := &wire.ExtendedHeader{
    54  		Header: wire.Header{
    55  			Type:             protocol.PacketTypeInitial,
    56  			SrcConnectionID:  srcConnID,
    57  			DestConnectionID: destConnID,
    58  			Length:           protocol.ByteCount(length),
    59  			Version:          version,
    60  		},
    61  		PacketNumberLen: pnLength,
    62  		PacketNumber:    0x0,
    63  	}
    64  
    65  	raw := writePacket(hdr, payload)
    66  
    67  	// encrypt payload and header
    68  	payloadOffset := len(raw) - payloadSize
    69  	var encrypted []byte
    70  	encrypted = sealer.Seal(encrypted, payload, hdr.PacketNumber, raw[:payloadOffset])
    71  	hdrBytes := raw[0:payloadOffset]
    72  	encrypted = append(hdrBytes, encrypted...)
    73  	pnOffset := payloadOffset - int(pnLength) // packet number offset
    74  	sealer.EncryptHeader(
    75  		encrypted[payloadOffset:payloadOffset+16], // first 16 bytes of payload (sample)
    76  		&encrypted[0],                     // first byte of header
    77  		encrypted[pnOffset:payloadOffset], // packet number bytes
    78  	)
    79  	return encrypted
    80  }
    81  
    82  // ComposeRetryPacket returns a new raw Retry Packet
    83  func ComposeRetryPacket(
    84  	srcConnID protocol.ConnectionID,
    85  	destConnID protocol.ConnectionID,
    86  	origDestConnID protocol.ConnectionID,
    87  	token []byte,
    88  	version protocol.VersionNumber,
    89  ) []byte {
    90  	hdr := &wire.ExtendedHeader{
    91  		Header: wire.Header{
    92  			Type:             protocol.PacketTypeRetry,
    93  			SrcConnectionID:  srcConnID,
    94  			DestConnectionID: destConnID,
    95  			Token:            token,
    96  			Version:          version,
    97  		},
    98  	}
    99  	data := writePacket(hdr, nil)
   100  	return append(data, handshake.GetRetryIntegrityTag(data, origDestConnID, version)[:]...)
   101  }