Fix header protection for very small packets

agents/parse_agent.go

@@ -56,6 +56,10 @@ func (a *ParsingAgent) Run(conn *Connection) {
 							}
 
 							sample, pnOffset := GetPacketSample(header, ciphertext)
+							if sample == nil {
+								a.Logger.Printf("Packet is too short to for header protection, dropping it\n")
+								break packetSelect
+							}
 							mask := cryptoState.HeaderRead.Encrypt(sample, make([]byte, 5, 5))
 							ciphertext[0] ^= mask[0] & firstByteMask
 

connection.go

@@ -119,6 +119,11 @@ func (c *Connection) EncodeAndEncrypt(packet Packet, level EncryptionLevel) []by
 			firstByteMask = 0x0F
 		}
 		sample, pnOffset := GetPacketSample(packet.Header(), packetBytes)
+		if sample == nil {
+			paddedSize := pnOffset + 4
+			packet.(Framer).PadTo(paddedSize)
+			return c.EncodeAndEncrypt(packet, level)
+		}
 		mask := cryptoState.HeaderWrite.Encrypt(sample, make([]byte, 5, 5))
 		packetBytes[0] ^= mask[0] & firstByteMask
 

crypto.go

@@ -149,9 +149,8 @@ func GetPacketSample(header Header, packetBytes []byte) ([]byte, int) {
 	sampleOffset := pnOffset + 4
 
 	if sampleOffset+sampleLength > len(packetBytes) {
-		paddedBytes := make([]byte, sampleOffset+sampleLength)
-		copy(paddedBytes, packetBytes)
-		packetBytes = paddedBytes
+		// Packet is too short for sampling header protection, it must be padded first
+		return nil, pnOffset
 	}
 
 	return packetBytes[sampleOffset:sampleOffset+sampleLength], pnOffset