refactor(simplequicping): use step-by-step (#852)

See https://github.com/ooni/probe/issues/2159 and https://github.com/ooni/spec/pull/254

internal/measurexlite/conn.go

@@ -69,6 +69,67 @@ func (c *connTrace) Write(b []byte) (int, error) {
 	return count, err
 }
 
+// MaybeUDPLikeClose is a convenience function for closing a conn only when such a conn isn't nil.
+func MaybeCloseUDPLikeConn(conn model.UDPLikeConn) (err error) {
+	if conn != nil {
+		err = conn.Close()
+	}
+	return
+}
+
+// WrapUDPLikeConn returns a wrapped conn that saves network events into this trace.
+func (tx *Trace) WrapUDPLikeConn(conn model.UDPLikeConn) model.UDPLikeConn {
+	return &udpLikeConnTrace{
+		UDPLikeConn: conn,
+		tx:          tx,
+	}
+}
+
+// udpLikeConnTrace is a trace-aware model.UDPLikeConn.
+type udpLikeConnTrace struct {
+	// Implementation note: it seems ~safe to use embedding here because model.UDPLikeConn
+	// contains fields deriving from how lucas-clemente/quic-go uses the standard library
+	model.UDPLikeConn
+	tx *Trace
+}
+
+// Read implements model.UDPLikeConn.ReadFrom and saves network events.
+func (c *udpLikeConnTrace) ReadFrom(b []byte) (int, net.Addr, error) {
+	started := c.tx.TimeSince(c.tx.ZeroTime)
+	count, addr, err := c.UDPLikeConn.ReadFrom(b)
+	finished := c.tx.TimeSince(c.tx.ZeroTime)
+	address := addrStringIfNotNil(addr)
+	select {
+	case c.tx.NetworkEvent <- NewArchivalNetworkEvent(
+		c.tx.Index, started, netxlite.ReadFromOperation, "udp", address, count, err, finished):
+	default: // buffer is full
+	}
+	return count, addr, err
+}
+
+// Write implements model.UDPLikeConn.WriteTo and saves network events.
+func (c *udpLikeConnTrace) WriteTo(b []byte, addr net.Addr) (int, error) {
+	started := c.tx.TimeSince(c.tx.ZeroTime)
+	address := addr.String()
+	count, err := c.UDPLikeConn.WriteTo(b, addr)
+	finished := c.tx.TimeSince(c.tx.ZeroTime)
+	select {
+	case c.tx.NetworkEvent <- NewArchivalNetworkEvent(
+		c.tx.Index, started, netxlite.WriteToOperation, "udp", address, count, err, finished):
+	default: // buffer is full
+	}
+	return count, err
+}
+
+// addrStringIfNotNil returns the string of the given addr
+// unless the addr is nil, in which case it returns an empty string.
+func addrStringIfNotNil(addr net.Addr) (out string) {
+	if addr != nil {
+		out = addr.String()
+	}
+	return
+}
+
 // NewArchivalNetworkEvent creates a new model.ArchivalNetworkEvent.
 func NewArchivalNetworkEvent(index int64, started time.Duration, operation string, network string,
 	address string, count int, err error, finished time.Duration) *model.ArchivalNetworkEvent {