Add TODO for onFLightState13 callback

Author: theodorsm

conn.go

@@ -1505,11 +1505,10 @@ func (c *Conn) handshake(
 			closed:             make(chan struct{}),
 		}
 		c.handshakeConfig.onFlightState13 = func(_ flightVal13, s handshakeState) {
-			currentFlight := c.fsm.(*handshakeFSM13).currentFlight //nolint:forcetypeassert
-			shouldClose := currentFlight.isLastSendFlight() || currentFlight.isLastRecvFlight()
-			//nolint:godox
+			// The ACK for the last flights has been received and we are in a Finished state.
+			// nolint:godox
 			// TODO: should be moved to FSM.
-			if shouldClose && c.setHandshakeCompletedSuccessfully() {
+			if s == handshakeFinished && c.setHandshakeCompletedSuccessfully() {
 				close(done)
 			}
 		}

handshaker_13.go

@@ -76,6 +76,8 @@ func (s *handshakeFSM13) Run(ctx context.Context, conn flightConn, initialState
 	}()
 	for {
 		s.cfg.log.Tracef("[handshake13:%s] %s: %s", srvCliStr(s.state.isClient), s.currentFlight.String(), state.String())
+		// nolint:godox
+		// TODO:: refactor callback, see discussion in https://github.com/pion/dtls/pull/738#discussion_r3131501159
 		if s.cfg.onFlightState13 != nil {
 			s.cfg.onFlightState13(s.currentFlight, state)
 		}