Test ACKing and reassembly of post-handshake messages

When it comes to receiving post-handshake, we're mostly already OK. Go
ahead and test that.

I say mostly because we probably should ACK as soon as we've gotten a
complete message. (Though we'd need to be careful not to ACK too many
times we've already got a bunch of messages queued up.) On the other
hand, outside of unit tests driving KeyUpdate, there isn't really a huge
issue to delaying the ACK and thus KeyUpdate by 100ms, so maybe it's not
that big of a deal.

Bug: 42290594
Change-Id: I86d1b81ddb5ebb98022b12659fafa9f8d4fd26d0
Reviewed-on: https://boringssl-review.googlesource.com/c/boringssl/+/73147
Reviewed-by: Nick Harper <nharper@chromium.org>
Commit-Queue: David Benjamin <davidben@google.com>
diff --git a/ssl/test/runner/dtls.go b/ssl/test/runner/dtls.go
index 5011815..3b08f79 100644
--- a/ssl/test/runner/dtls.go
+++ b/ssl/test/runner/dtls.go
@@ -61,6 +61,15 @@
 	}
 }
 
+// Split returns two fragments for the message.
+func (m *DTLSMessage) Split(offset int) (DTLSFragment, DTLSFragment) {
+	if m.IsChangeCipherSpec {
+		panic("tls: cannot split ChangeCipherSpec")
+	}
+
+	return m.Fragment(0, offset), m.Fragment(offset, len(m.Data)-offset)
+}
+
 // A DTLSFragment is a DTLS handshake fragment or ChangeCipherSpec, along with
 // the epoch that it is to be sent under.
 type DTLSFragment struct {
diff --git a/ssl/test/runner/runner.go b/ssl/test/runner/runner.go
index 72bf0a8..59d381b 100644
--- a/ssl/test/runner/runner.go
+++ b/ssl/test/runner/runner.go
@@ -12356,7 +12356,7 @@
 							MaxPacketLength:             512,
 							WriteFlightDTLS: func(c *DTLSController, prev, received, next []DTLSMessage, records []DTLSRecordNumberInfo) {
 								if len(received) == 0 || received[0].Type != typeClientHello {
-									// Leave post-handshake flights alone.
+									// We test post-handshake flights separately.
 									c.WriteFlight(next)
 									return
 								}
@@ -12550,6 +12550,100 @@
 					},
 					flags: flags,
 				})
+
+				testCases = append(testCases, testCase{
+					protocol: dtls,
+					name:     "DTLS-Retransmit-Client-ACKPostHandshake" + suffix,
+					config: Config{
+						MaxVersion: vers.version,
+						Bugs: ProtocolBugs{
+							WriteFlightDTLS: func(c *DTLSController, prev, received, next []DTLSMessage, records []DTLSRecordNumberInfo) {
+								if next[0].Type != typeNewSessionTicket {
+									c.WriteFlight(next)
+									return
+								}
+
+								// The test should try to send two NewSessionTickets in a row.
+								if len(next) != 2 {
+									panic("unexpected message count")
+								}
+
+								// Send part of first ticket post-handshake message.
+								first0, second0 := next[0].Split(len(next[0].Data) / 2)
+								first1, second1 := next[1].Split(len(next[1].Data) / 2)
+								c.WriteFragments([]DTLSFragment{first0})
+
+								// The shim should ACK on a timer.
+								c.ExpectNextTimeout(useTimeouts[0] / 4)
+								c.AdvanceClock(useTimeouts[0] / 4)
+								c.ReadACK(c.InEpoch())
+
+								// The shim is just waiting for us to retransmit.
+								c.ExpectNoNextTimeout()
+
+								// Send some more fragments.
+								c.WriteFragments([]DTLSFragment{first0, second1})
+
+								// The shim should ACK, again on a timer.
+								c.ExpectNextTimeout(useTimeouts[0] / 4)
+								c.AdvanceClock(useTimeouts[0] / 4)
+								c.ReadACK(c.InEpoch())
+								c.ExpectNoNextTimeout()
+
+								// Finish up both messages. We implicitly test if shim
+								// processed these messages by checking that it returned a new
+								// session.
+								c.WriteFragments([]DTLSFragment{first1, second0})
+
+								// The shim should ACK again, once the timer expires.
+								//
+								// TODO(crbug.com/42290594): Should the shim ACK immediately?
+								// Otherwise KeyUpdates are delayed, which will complicated
+								// downstream testing.
+								c.ExpectNextTimeout(useTimeouts[0] / 4)
+								c.AdvanceClock(useTimeouts[0] / 4)
+								c.ReadACK(c.InEpoch())
+								c.ExpectNoNextTimeout()
+							},
+						},
+					},
+					flags: flags,
+				})
+
+				testCases = append(testCases, testCase{
+					protocol: dtls,
+					name:     "DTLS-Retransmit-Client-ACKPostHandshakeTwice" + suffix,
+					config: Config{
+						MaxVersion: vers.version,
+						Bugs: ProtocolBugs{
+							WriteFlightDTLS: func(c *DTLSController, prev, received, next []DTLSMessage, records []DTLSRecordNumberInfo) {
+								if next[0].Type != typeNewSessionTicket {
+									c.WriteFlight(next)
+									return
+								}
+
+								// The test should try to send two NewSessionTickets in a row.
+								if len(next) != 2 {
+									panic("unexpected message count")
+								}
+
+								// Send the flight. The shim should ACK it.
+								c.WriteFlight(next)
+								c.AdvanceClock(useTimeouts[0] / 4)
+								c.ReadACK(c.InEpoch())
+								c.ExpectNoNextTimeout()
+
+								// Retransmit the flight, as if we lost the ACK. The shim should
+								// ACK again.
+								c.WriteFlight(next)
+								c.AdvanceClock(useTimeouts[0] / 4)
+								c.ReadACK(c.InEpoch())
+								c.ExpectNoNextTimeout()
+							},
+						},
+					},
+					flags: flags,
+				})
 			}
 		}
 	}