diff --git a/conn.go b/conn.go index fdf7475c4..3d7423de7 100644 --- a/conn.go +++ b/conn.go @@ -583,7 +583,7 @@ func (c *Conn) Write(payload []byte) (int, error) { return 0, err } - ctx, cancel := c.contextWithClose(c.writeDeadline) + ctx, cancel := c.contextWithClose() defer cancel() err := c.writeApplicationData(ctx, []*dtlsflight.Outbound{ @@ -855,62 +855,109 @@ func (c *Conn) cacheHandshake(outbound *dtlsflight.Outbound, dtlsHandshake *hand return nil } -func (c *Conn) contextWithClose(ctx context.Context) (context.Context, context.CancelFunc) { - closeCtx, cancel := context.WithCancelCause(context.WithoutCancel(ctx)) - go func() { - select { - case <-c.closed.Done(): - cancel(context.Canceled) - case <-ctx.Done(): - err := ctx.Err() - if err == nil { - err = context.DeadlineExceeded - } - cancel(err) - case <-closeCtx.Done(): +func (c *Conn) contextWithClose() (context.Context, context.CancelFunc) { + closeCtx, cancel := context.WithCancelCause(context.Background()) + + detachLifetime := context.AfterFunc(c.closed, func() { + err := c.closed.Err() + if err == nil { + err = context.Canceled } - }() + cancel(err) + }) + + detachDeadline := context.AfterFunc(c.writeDeadline.Context(), func() { + err := c.writeDeadline.Err() + if err == nil { + err = context.DeadlineExceeded + } + cancel(err) + }) return closeCtx, func() { + detachLifetime() + detachDeadline() cancel(context.Canceled) } } func (c *Conn) contextWithCloseAndWriteDeadline(ctx context.Context) (context.Context, context.CancelFunc) { operationCtx, cancel := context.WithCancelCause(context.Background()) - go func() { - select { - case <-c.closed.Done(): - cancel(context.Canceled) - case <-c.writeDeadline.Done(): - cancel(context.DeadlineExceeded) - case <-ctx.Done(): - cancel(ctx.Err()) - case <-operationCtx.Done(): + + detachLifetime := context.AfterFunc(c.closed, func() { + err := c.closed.Err() + if err == nil { + err = context.Canceled } - }() + cancel(err) + }) + + detachDeadline := context.AfterFunc(c.writeDeadline.Context(), func() { + err := c.writeDeadline.Err() + if err == nil { + err = context.DeadlineExceeded + } + cancel(err) + }) + + detachCtx := context.AfterFunc(ctx, func() { + err := ctx.Err() + if err == nil { + err = context.Canceled + } + cancel(err) + }) return operationCtx, func() { + detachLifetime() + detachDeadline() + detachCtx() cancel(context.Canceled) } } func (c *Conn) compactPreparedRecords(records []preparedRecord) []preparedDatagram { - datagrams := make([]preparedDatagram, 0, len(records)) - current := preparedDatagram{} + if len(records) == 0 { + return []preparedDatagram{} + } + + totalSize := 0 + for _, record := range records { + totalSize += len(record.raw) + } + + datagrams := make([]preparedDatagram, len(records)) + flatRaw := make([]byte, 0, totalSize) + + datagramIndex := 0 + currentSize := 0 + offset := 0 + for _, record := range records { - if len(current.raw) > 0 && len(current.raw)+len(record.raw) >= c.maximumTransmissionUnit { - datagrams = append(datagrams, current) - current = preparedDatagram{} + recordSize := len(record.raw) + + flatRaw = append(flatRaw, record.raw...) + + if currentSize > 0 && currentSize+recordSize >= c.maximumTransmissionUnit { + datagrams[datagramIndex].raw = flatRaw[offset : offset+currentSize] + datagramIndex++ + offset += currentSize + currentSize = 0 } - current.raw = append(current.raw, record.raw...) + + currentSize += recordSize + if record.tracked != nil { - current.tracked = append(current.tracked, *record.tracked) + datagrams[datagramIndex].tracked = append( + datagrams[datagramIndex].tracked, + *record.tracked, + ) } } - datagrams = append(datagrams, current) - return datagrams + datagrams[datagramIndex].raw = flatRaw[offset : offset+currentSize] + + return datagrams[:datagramIndex+1] } func (c *Conn) prepareRecord(outbound *dtlsflight.Outbound) ([]byte, error) { @@ -1240,26 +1287,57 @@ func selectHandshakeFragment(offsets map[uint32]uint32, raw []byte) (bool, error return ok && length == header.FragmentLength, nil } +var noContentFragments = [][]byte{ //nolint:gochecknoglobals + {}, +} + func (c *Conn) fragmentHandshake(dtlsHandshake *handshake.Handshake) ([][]byte, error) { + messageSize := dtlsHandshake.Message.MarshalSize() + numFragments := (messageSize-1)/c.maximumTransmissionUnit + 1 + + fragmentedHandshakes := make([][]byte, numFragments) + + if numFragments == 1 { + fragmentedHandshake := make([]byte, handshake.HeaderLength+messageSize) + + headerFragment := handshake.Header{ + Type: dtlsHandshake.Header.Type, + Length: dtlsHandshake.Header.Length, + MessageSequence: dtlsHandshake.Header.MessageSequence, + FragmentOffset: uint32(0), + FragmentLength: uint32(messageSize), //nolint:gosec // G115 + } + + _, err := headerFragment.MarshalTo(fragmentedHandshake) + if err != nil { + return nil, err + } + + _, err = dtlsHandshake.Message.MarshalTo(fragmentedHandshake[handshake.HeaderLength:]) + if err != nil { + return nil, err + } + + fragmentedHandshakes[0] = fragmentedHandshake + + return fragmentedHandshakes, nil + } + content, err := dtlsHandshake.Message.Marshal() if err != nil { return nil, err } - fragmentedHandshakes := make([][]byte, 0) - contentFragments := util.SplitBytes(content, c.maximumTransmissionUnit) if len(contentFragments) == 0 { - contentFragments = [][]byte{ - {}, - } + contentFragments = noContentFragments } offset := 0 - for _, contentFragment := range contentFragments { + for i, contentFragment := range contentFragments { contentFragmentLen := len(contentFragment) - headerFragment := &handshake.Header{ + headerFragment := handshake.Header{ Type: dtlsHandshake.Header.Type, Length: dtlsHandshake.Header.Length, MessageSequence: dtlsHandshake.Header.MessageSequence, @@ -1267,15 +1345,18 @@ func (c *Conn) fragmentHandshake(dtlsHandshake *handshake.Handshake) ([][]byte, FragmentLength: uint32(contentFragmentLen), //nolint:gosec // G115 } + fragmentedHandshake := make([]byte, handshake.HeaderLength+contentFragmentLen) + offset += contentFragmentLen - fragmentedHandshake, err := headerFragment.Marshal() + _, err := headerFragment.MarshalTo(fragmentedHandshake) if err != nil { return nil, err } - fragmentedHandshake = append(fragmentedHandshake, contentFragment...) - fragmentedHandshakes = append(fragmentedHandshakes, fragmentedHandshake) + copy(fragmentedHandshake[handshake.HeaderLength:], contentFragment) + + fragmentedHandshakes[i] = fragmentedHandshake } return fragmentedHandshakes, nil @@ -2828,12 +2909,7 @@ func (c *Conn) close(byUser bool) error { } func (c *Conn) isConnectionClosed() bool { - select { - case <-c.closed.Done(): - return true - default: - return false - } + return c.closed.Err() != nil } func (c *Conn) setLocalEpoch(epoch uint16) { diff --git a/internal/closer/closer.go b/internal/closer/closer.go index 41109485d..073b3df15 100644 --- a/internal/closer/closer.go +++ b/internal/closer/closer.go @@ -6,6 +6,7 @@ package closer import ( "context" + "time" ) // Closer allows for each signaling a channel for shutdown. @@ -48,3 +49,11 @@ func (c *Closer) Err() error { func (c *Closer) Close() { c.closeFunc() } + +func (c *Closer) Deadline() (deadline time.Time, ok bool) { + return c.ctx.Deadline() +} + +func (c *Closer) Value(key any) any { + return c.ctx.Value(key) +} diff --git a/internal/flight/cache.go b/internal/flight/cache.go index ac322fc03..06718980c 100644 --- a/internal/flight/cache.go +++ b/internal/flight/cache.go @@ -5,7 +5,6 @@ package flight import ( - "bytes" "fmt" "sync" @@ -45,7 +44,7 @@ func (h *Cache) Push(data []byte, epoch, messageSequence uint16, typ handshake.T h.mu.Lock() defer h.mu.Unlock() - h.cache = append(h.cache, &HandshakeCacheItem{Data: bytes.Clone(data), Epoch: epoch, MessageSequence: messageSequence, Typ: typ, IsClient: isClient}) + h.cache = append(h.cache, &HandshakeCacheItem{Data: data, Epoch: epoch, MessageSequence: messageSequence, Typ: typ, IsClient: isClient}) } // PullExact returns the handshake message with the requested sequence and diff --git a/pkg/protocol/handshake/handshake.go b/pkg/protocol/handshake/handshake.go index 4319e6ec4..d1af6ffca 100644 --- a/pkg/protocol/handshake/handshake.go +++ b/pkg/protocol/handshake/handshake.go @@ -131,19 +131,18 @@ func (h *Handshake) Marshal() ([]byte, error) { return nil, dtlserrors.ErrUnableToMarshalFragmented } - message, err := h.Message.Marshal() + out := make([]byte, HeaderLength+h.Message.MarshalSize()) + n, err := h.Message.MarshalTo(out[HeaderLength:]) if err != nil { return nil, err } - out := make([]byte, HeaderLength+len(message)) - h.Header.Length = uint32(len(message)) //nolint:gosec // handshake messages are bounded to uint24 on the wire. + h.Header.Length = uint32(n) //nolint:gosec // handshake messages are bounded to uint24 on the wire. h.Header.FragmentLength = h.Header.Length h.Header.Type = h.Message.Type() if _, err = h.Header.MarshalTo(out); err != nil { return nil, err } - copy(out[HeaderLength:], message) return out, nil }