Fixes PN protection edge case for very short packets

agents/parse_agent.go

@@ -51,10 +51,10 @@ func (a *ParsingAgent) Run(conn *Connection) {
 						if cryptoState != nil && cryptoState.PacketRead != nil && cryptoState.Read != nil {
 							a.Logger.Printf("Decrypting packet number of %s packet of length %d bytes", header.PacketType().String(), len(ciphertext))
 
-							sample, sampleOffset := GetPacketSample(header, ciphertext)
-							pn := cryptoState.PacketRead.Encrypt(sample, ciphertext[sampleOffset-4:sampleOffset])
+							sample, pnOffset := GetPacketSample(header, ciphertext)
+							pn := cryptoState.PacketRead.Encrypt(sample, ciphertext[pnOffset:pnOffset+4])
 							pnLength = ReadTruncatedPN(bytes.NewReader(pn)).Length
-							copy(ciphertext[sampleOffset-4:sampleOffset], pn[:pnLength])
+							copy(ciphertext[pnOffset:pnOffset+4], pn[:pnLength])
 							header = ReadHeader(bytes.NewReader(ciphertext), a.conn) // Update PN
 						} else {
 							a.Logger.Printf("Packet number of %s packet of length %d bytes could not be decrypted, putting it back in waiting buffer\n", header.PacketType().String(), len(ciphertext))
@@ -74,7 +74,7 @@ func (a *ParsingAgent) Run(conn *Connection) {
 						pLen := int(lHeader.Length.Value) - pnLength
 
 						if hLen+pLen > len(ciphertext) {
-							a.Logger.Printf("Payload length is past the received bytes, has PN decryption failed ? Aborting")
+							a.Logger.Printf("Payload length %d is past the %d received bytes, has PN decryption failed ? Aborting", hLen+pLen, len(ciphertext))
 							break packetSelect
 						}
 

connection.go

@@ -80,9 +80,9 @@ func (c *Connection) SendPacket(packet Packet, level EncryptionLevel) {
 		protectedPayload := cryptoState.Write.Seal(nil, EncodeArgs(packet.Header().PacketNumber()), payload, header)
 		packetBytes := append(header, protectedPayload...)
 
-		sample, sampleOffset := GetPacketSample(packet.Header(), packetBytes)
+		sample, pnOffset := GetPacketSample(packet.Header(), packetBytes)
 
-		copy(packetBytes[sampleOffset-4:sampleOffset], cryptoState.PacketWrite.Encrypt(sample, packetBytes[sampleOffset-4:sampleOffset])[:packet.Header().TruncatedPN().Length])
+		copy(packetBytes[pnOffset:pnOffset+4], cryptoState.PacketWrite.Encrypt(sample, packetBytes[pnOffset:pnOffset+4])[:packet.Header().TruncatedPN().Length])
 
 		c.UdpConnection.Write(packetBytes)
 

crypto.go

@@ -108,20 +108,22 @@ func newProtectedAead(tls *pigotls.Connection, secret []byte) cipher.AEAD {
 }
 
 func GetPacketSample(header Header, packetBytes []byte) ([]byte, int) {
-	var sampleOffset int
+	var pnOffset int
 	sampleLength := 16
 	switch h := header.(type) {
 	case *LongHeader:
-		sampleOffset = h.HeaderLength() - h.TruncatedPN().Length + 4
+		pnOffset = h.HeaderLength() - h.TruncatedPN().Length
 	case *ShortHeader:
-		sampleOffset = 1 + len(h.DestinationCID) + 4
+		pnOffset = 1 + len(h.DestinationCID)
+	}
 
-		if sampleOffset + sampleLength > len(packetBytes) {
-			sampleOffset = len(packetBytes) - sampleLength
-		}
+	sampleOffset := pnOffset + 4
+	if sampleOffset + sampleLength > len(packetBytes) {
+		sampleOffset = len(packetBytes) - sampleLength
 	}
+
 	if sampleOffset <= 0 || sampleOffset+sampleLength > len(packetBytes) {
-		sampleOffset = 4
+		sampleOffset = 4  // TODO: Is this working ?
 	}
-	return packetBytes[sampleOffset:sampleOffset+sampleLength], sampleOffset
+	return packetBytes[sampleOffset:sampleOffset+sampleLength], pnOffset
 }