Fix DTLS handling of multiple records in a packet.

9a41d1b946c17f353c1f7f1d35713b5623dc399e broke handling of multiple records in
a single packet. If |extend| is true, not all of the previous packet should be
consumed, only up to the record length.

Add a test which stresses the DTLS stack's handling of multiple handshake
fragments in a handshake record and multiple handshake records in a packet.

Change-Id: I96571098ad9001e96440501c4730325227b155b8
Reviewed-on: https://boringssl-review.googlesource.com/4950
Reviewed-by: Adam Langley <agl@google.com>
diff --git a/ssl/test/runner/common.go b/ssl/test/runner/common.go
index 9e6cce3..864f526 100644
--- a/ssl/test/runner/common.go
+++ b/ssl/test/runner/common.go
@@ -478,7 +478,9 @@
 	// MaxHandshakeRecordLength, if non-zero, is the maximum size of a
 	// handshake record. Handshake messages will be split into multiple
 	// records at the specified size, except that the client_version will
-	// never be fragmented.
+	// never be fragmented. For DTLS, it is the maximum handshake fragment
+	// size, not record size; DTLS allows multiple handshake fragments in a
+	// single handshake record. See |PackHandshakeFragments|.
 	MaxHandshakeRecordLength int
 
 	// FragmentClientVersion will allow MaxHandshakeRecordLength to apply to
@@ -712,6 +714,15 @@
 	// DHGroupPrime, if not nil, is used to define the (finite field)
 	// Diffie-Hellman group. The generator used is always two.
 	DHGroupPrime *big.Int
+
+	// PackHandshakeFragments, if true, causes handshake fragments to be
+	// packed into individual handshake records, up to the specified record
+	// size.
+	PackHandshakeFragments int
+
+	// PackHandshakeRecords, if true, causes handshake records to be packed
+	// into individual packets, up to the specified packet size.
+	PackHandshakeRecords int
 }
 
 func (c *Config) serverInit() {
diff --git a/ssl/test/runner/dtls.go b/ssl/test/runner/dtls.go
index 85c4247..50f7786 100644
--- a/ssl/test/runner/dtls.go
+++ b/ssl/test/runner/dtls.go
@@ -196,6 +196,8 @@
 		return nil
 	}
 
+	// This is a test-only DTLS implementation, so there is no need to
+	// retain |c.pendingFragments| for a future retransmit.
 	var fragments [][]byte
 	fragments, c.pendingFragments = c.pendingFragments, fragments
 
@@ -208,38 +210,66 @@
 		fragments = tmp
 	}
 
-	// Send them all.
+	maxRecordLen := c.config.Bugs.PackHandshakeFragments
+	maxPacketLen := c.config.Bugs.PackHandshakeRecords
+
+	// Pack handshake fragments into records.
+	var records [][]byte
 	for _, fragment := range fragments {
 		if c.config.Bugs.SplitFragmentHeader {
-			if _, err := c.dtlsWriteRawRecord(recordTypeHandshake, fragment[:2]); err != nil {
-				return err
+			records = append(records, fragment[:2])
+			records = append(records, fragment[2:])
+		} else if c.config.Bugs.SplitFragmentBody {
+			if len(fragment) > 12 {
+				records = append(records, fragment[:13])
+				records = append(records, fragment[13:])
+			} else {
+				records = append(records, fragment)
 			}
-			fragment = fragment[2:]
-		} else if c.config.Bugs.SplitFragmentBody && len(fragment) > 12 {
-			if _, err := c.dtlsWriteRawRecord(recordTypeHandshake, fragment[:13]); err != nil {
-				return err
-			}
-			fragment = fragment[13:]
+		} else if i := len(records) - 1; len(records) > 0 && len(records[i])+len(fragment) <= maxRecordLen {
+			records[i] = append(records[i], fragment...)
+		} else {
+			// The fragment will be appended to, so copy it.
+			records = append(records, append([]byte{}, fragment...))
+		}
+	}
+
+	// Format them into packets.
+	var packets [][]byte
+	for _, record := range records {
+		b, err := c.dtlsSealRecord(recordTypeHandshake, record)
+		if err != nil {
+			return err
 		}
 
-		// TODO(davidben): A real DTLS implementation needs to
-		// retransmit handshake messages. For testing purposes, we don't
-		// actually care.
-		if _, err := c.dtlsWriteRawRecord(recordTypeHandshake, fragment); err != nil {
+		if i := len(packets) - 1; len(packets) > 0 && len(packets[i])+len(b.data) <= maxPacketLen {
+			packets[i] = append(packets[i], b.data...)
+		} else {
+			// The sealed record will be appended to and reused by
+			// |c.out|, so copy it.
+			packets = append(packets, append([]byte{}, b.data...))
+		}
+		c.out.freeBlock(b)
+	}
+
+	// Send all the packets.
+	for _, packet := range packets {
+		if _, err := c.conn.Write(packet); err != nil {
 			return err
 		}
 	}
 	return nil
 }
 
-func (c *Conn) dtlsWriteRawRecord(typ recordType, data []byte) (n int, err error) {
+// dtlsSealRecord seals a record into a block from |c.out|'s pool.
+func (c *Conn) dtlsSealRecord(typ recordType, data []byte) (b *block, err error) {
 	recordHeaderLen := dtlsRecordHeaderLen
 	maxLen := c.config.Bugs.MaxHandshakeRecordLength
 	if maxLen <= 0 {
 		maxLen = 1024
 	}
 
-	b := c.out.newBlock()
+	b = c.out.newBlock()
 
 	explicitIVLen := 0
 	explicitIVIsSeq := false
@@ -286,6 +316,14 @@
 	}
 	copy(b.data[recordHeaderLen+explicitIVLen:], data)
 	c.out.encrypt(b, explicitIVLen)
+	return
+}
+
+func (c *Conn) dtlsWriteRawRecord(typ recordType, data []byte) (n int, err error) {
+	b, err := c.dtlsSealRecord(typ, data)
+	if err != nil {
+		return
+	}
 
 	_, err = c.conn.Write(b.data)
 	if err != nil {
diff --git a/ssl/test/runner/runner.go b/ssl/test/runner/runner.go
index 6510e33..6c0d294 100644
--- a/ssl/test/runner/runner.go
+++ b/ssl/test/runner/runner.go
@@ -1105,6 +1105,17 @@
 		},
 		flags: []string{"-async"},
 	},
+	{
+		protocol: dtls,
+		name:     "PackDTLSHandshake",
+		config: Config{
+			Bugs: ProtocolBugs{
+				MaxHandshakeRecordLength: 2,
+				PackHandshakeFragments:   20,
+				PackHandshakeRecords:     200,
+			},
+		},
+	},
 }
 
 func doExchange(test *testCase, config *Config, conn net.Conn, messageLen int, isResume bool) error {