9da1afc fix(remote): address pr review transport issues · vev

M internal/adapters/dgram/ssh_dialer.go +1 -1


@@ -50,7 +50,7 @@ func (d RemoteDialer) Dial(ctx context.Context) (ports.Transport, error) {
50	50		return nil, ctx.Err()
51	51		default:
52	52		}
53		-	return sshstdio.Dial(d.Target, d.Session)
	53	+	return sshstdio.DialContext(ctx, d.Target, d.Session)
54	54		}
55	55
56	56		func sshTargetHost(target string) string {


@@ -50,7 +50,7 @@ func (d RemoteDialer) Dial(ctx context.Context) (ports.Transport, error) {
50		return nil, ctx.Err()	50		return nil, ctx.Err()
51		default:	51		default:
52		}	52		}
53	-	return sshstdio.Dial(d.Target, d.Session)	53	+	return sshstdio.DialContext(ctx, d.Target, d.Session)
54		}	54		}
55			55
56		func sshTargetHost(target string) string {	56		func sshTargetHost(target string) string {

M internal/adapters/dgram/transport.go +10 -0


@@ -23,6 +23,7 @@ const (
23	23		defaultHeartbeat = 3 * time.Second
24	24		defaultSilenceTimeout = 10 * time.Second
25	25		maxPendingReliable = 1024
	26	+	maxRecvBuffer = 1024
26	27		)
27	28
28	29		var (

@@ -304,6 +305,15 @@ func (t *Transport) enqueueReliable(seq uint64, f ports.Frame) {
304	305		if _, exists := t.recvBuf[seq]; exists {
305	306		return
306	307		}
	308	+	if len(t.recvBuf) >= maxRecvBuffer {
	309	+	if seq != t.nextRecvSeq {
	310	+	return
	311	+	}
	312	+	for bufferedSeq := range t.recvBuf {
	313	+	delete(t.recvBuf, bufferedSeq)
	314	+	break
	315	+	}
	316	+	}
307	317		t.recvBuf[seq] = f
308	318		t.deliverCond.Signal()
309	319		}


@@ -23,6 +23,7 @@ const (
23	defaultHeartbeat = 3 * time.Second	23		defaultHeartbeat = 3 * time.Second
24	defaultSilenceTimeout = 10 * time.Second	24		defaultSilenceTimeout = 10 * time.Second
25	maxPendingReliable = 1024	25		maxPendingReliable = 1024
		26	+	maxRecvBuffer = 1024
26	)	27		)
27		28
28	var (	29		var (

@@ -304,6 +305,15 @@ func (t *Transport) enqueueReliable(seq uint64, f ports.Frame) {
304	if _, exists := t.recvBuf[seq]; exists {	305		if _, exists := t.recvBuf[seq]; exists {
305	return	306		return
306	}	307		}
		308	+	if len(t.recvBuf) >= maxRecvBuffer {
		309	+	if seq != t.nextRecvSeq {
		310	+	return
		311	+	}
		312	+	for bufferedSeq := range t.recvBuf {
		313	+	delete(t.recvBuf, bufferedSeq)
		314	+	break
		315	+	}
		316	+	}
307	t.recvBuf[seq] = f	317		t.recvBuf[seq] = f
308	t.deliverCond.Signal()	318		t.deliverCond.Signal()
309	}	319		}

M internal/adapters/dgram/transport_test.go +18 -0


@@ -319,6 +319,24 @@ func TestPendingReliableQueueBounded(t *testing.T) {
319	319		}
320	320		}
321	321
	322	+	func TestReliableRecvBufferBoundedForFarFutureSequences(t *testing.T) {
	323	+	tr := &Transport{nextRecvSeq: 1, recvBuf: make(map[uint64]ports.Frame), done: make(chan struct{})}
	324	+	tr.deliverCond = sync.NewCond(&tr.deliverMu)
	325	+	for i := 0; i < maxRecvBuffer+100; i++ {
	326	+	tr.enqueueReliable(uint64(1000+i), ports.Frame{Type: ports.MsgOutput, Payload: []byte{byte(i)}})
	327	+	}
	328	+	if got := len(tr.recvBuf); got != maxRecvBuffer {
	329	+	t.Fatalf("recvBuf len=%d, want %d", got, maxRecvBuffer)
	330	+	}
	331	+	tr.enqueueReliable(1, ports.Frame{Type: ports.MsgOutput, Payload: []byte("next")})
	332	+	if _, ok := tr.recvBuf[1]; !ok {
	333	+	t.Fatalf("next expected sequence was dropped when far-future buffer was full")
	334	+	}
	335	+	if got := len(tr.recvBuf); got > maxRecvBuffer {
	336	+	t.Fatalf("recvBuf len=%d, want <= %d", got, maxRecvBuffer)
	337	+	}
	338	+	}
	339	+
322	340		func recvMaybe(tr *Transport, d time.Duration) (ports.Frame, bool) {
323	341		ch := make(chan ports.Frame, 1)
324	342		go func() { f, _ := tr.Recv(); ch <- f }()


@@ -319,6 +319,24 @@ func TestPendingReliableQueueBounded(t *testing.T) {
319	}	319		}
320	}	320		}
321		321
		322	+	func TestReliableRecvBufferBoundedForFarFutureSequences(t *testing.T) {
		323	+	tr := &Transport{nextRecvSeq: 1, recvBuf: make(map[uint64]ports.Frame), done: make(chan struct{})}
		324	+	tr.deliverCond = sync.NewCond(&tr.deliverMu)
		325	+	for i := 0; i < maxRecvBuffer+100; i++ {
		326	+	tr.enqueueReliable(uint64(1000+i), ports.Frame{Type: ports.MsgOutput, Payload: []byte{byte(i)}})
		327	+	}
		328	+	if got := len(tr.recvBuf); got != maxRecvBuffer {
		329	+	t.Fatalf("recvBuf len=%d, want %d", got, maxRecvBuffer)
		330	+	}
		331	+	tr.enqueueReliable(1, ports.Frame{Type: ports.MsgOutput, Payload: []byte("next")})
		332	+	if _, ok := tr.recvBuf[1]; !ok {
		333	+	t.Fatalf("next expected sequence was dropped when far-future buffer was full")
		334	+	}
		335	+	if got := len(tr.recvBuf); got > maxRecvBuffer {
		336	+	t.Fatalf("recvBuf len=%d, want <= %d", got, maxRecvBuffer)
		337	+	}
		338	+	}
		339	+
322	func recvMaybe(tr *Transport, d time.Duration) (ports.Frame, bool) {	340		func recvMaybe(tr *Transport, d time.Duration) (ports.Frame, bool) {
323	ch := make(chan ports.Frame, 1)	341		ch := make(chan ports.Frame, 1)
324	go func() { f, _ := tr.Recv(); ch <- f }()	342		go func() { f, _ := tr.Recv(); ch <- f }()

M internal/adapters/sshstdio/transport.go +11 -1


@@ -6,6 +6,7 @@ package sshstdio
6	6
7	7		import (
8	8		"bytes"
	9	+	"context"
9	10		"encoding/binary"
10	11		"errors"
11	12		"fmt"

@@ -152,8 +153,17 @@ func shellQuote(s string) string {
152	153		// child process' stdio. The subprocess is started with exec.Command argv, never
153	154		// through a shell.
154	155		func Dial(target, session string) (ports.Transport, error) {
	156	+	return DialContext(context.Background(), target, session)
	157	+	}
	158	+
	159	+	// DialContext is like Dial, but the context is propagated to ssh startup so a
	160	+	// canceled attach attempt interrupts the local ssh process.
	161	+	func DialContext(ctx context.Context, target, session string) (ports.Transport, error) {
	162	+	if err := ctx.Err(); err != nil {
	163	+	return nil, err
	164	+	}
155	165		spec := BuildCommand(target, session)
156		-	cmd := exec.Command(spec.Path, spec.Args...)
	166	+	cmd := exec.CommandContext(ctx, spec.Path, spec.Args...)
157	167		stdin, err := cmd.StdinPipe()
158	168		if err != nil {
159	169		return nil, fmt.Errorf("sshstdio: stdin pipe: %w", err)

 import (
 	"bytes"
+	"context"
 	"encoding/binary"
 	"errors"
 	"fmt"
 // child process' stdio. The subprocess is started with exec.Command argv, never
 // through a shell.
 func Dial(target, session string) (ports.Transport, error) {
+	return DialContext(context.Background(), target, session)
+}
+// DialContext is like Dial, but the context is propagated to ssh startup so a
+// canceled attach attempt interrupts the local ssh process.
+func DialContext(ctx context.Context, target, session string) (ports.Transport, error) {
+	if err := ctx.Err(); err != nil {
+		return nil, err
+	}
 	spec := BuildCommand(target, session)
-	cmd := exec.Command(spec.Path, spec.Args...)
+	cmd := exec.CommandContext(ctx, spec.Path, spec.Args...)
 	stdin, err := cmd.StdinPipe()
 	if err != nil {
 		return nil, fmt.Errorf("sshstdio: stdin pipe: %w", err)

M internal/adapters/sshstdio/transport_test.go +10 -0


@@ -2,6 +2,7 @@ package sshstdio
2	2
3	3		import (
4	4		"bytes"
	5	+	"context"
5	6		"errors"
6	7		"io"
7	8		"os/exec"

@@ -115,6 +116,15 @@ func TestTransportRejectsZeroLengthFrame(t *testing.T) {
115	116		}
116	117		}
117	118
	119	+	func TestDialContextCanceledBeforeStart(t *testing.T) {
	120	+	ctx, cancel := context.WithCancel(context.Background())
	121	+	cancel()
	122	+	_, err := DialContext(ctx, "example.com", "work")
	123	+	if !errors.Is(err, context.Canceled) {
	124	+	t.Fatalf("DialContext error = %v, want context.Canceled", err)
	125	+	}
	126	+	}
	127	+
118	128		func TestProcessCloser(t *testing.T) {
119	129		tests := []struct {
120	130		name string


@@ -2,6 +2,7 @@ package sshstdio
2		2
3	import (	3		import (
4	"bytes"	4		"bytes"
		5	+	"context"
5	"errors"	6		"errors"
6	"io"	7		"io"
7	"os/exec"	8		"os/exec"

@@ -115,6 +116,15 @@ func TestTransportRejectsZeroLengthFrame(t *testing.T) {
115	}	116		}
116	}	117		}
117		118
		119	+	func TestDialContextCanceledBeforeStart(t *testing.T) {
		120	+	ctx, cancel := context.WithCancel(context.Background())
		121	+	cancel()
		122	+	_, err := DialContext(ctx, "example.com", "work")
		123	+	if !errors.Is(err, context.Canceled) {
		124	+	t.Fatalf("DialContext error = %v, want context.Canceled", err)
		125	+	}
		126	+	}
		127	+
118	func TestProcessCloser(t *testing.T) {	128		func TestProcessCloser(t *testing.T) {
119	tests := []struct {	129		tests := []struct {
120	name string	130		name string

M internal/usecase/daemon/client.go +16 -1


@@ -231,15 +231,30 @@ func (d Daemon) boundedSendErr(ac attachedClient, f ports.Frame) error {
231	231		}
232	232
233	233		func (d Daemon) boundedSendOutputErr(ac attachedClient, b []byte) error {
234		-	return d.boundedSendWith(ac, func() error {
	234	+	_, err := d.boundedSendOutputErrTransport(ac, b)
	235	+	return err
	236	+	}
	237	+
	238	+	func (d Daemon) boundedSendOutputErrTransport(ac attachedClient, b []byte) (ports.Transport, error) {
	239	+	var (
	240	+	usedMu sync.Mutex
	241	+	used ports.Transport
	242	+	)
	243	+	err := d.boundedSendWith(ac, func() error {
235	244		ac.sendMu.Lock()
236	245		defer ac.sendMu.Unlock()
237	246		tr := ac.transport()
	247	+	usedMu.Lock()
	248	+	used = tr
	249	+	usedMu.Unlock()
238	250		if tr == nil {
239	251		return errors.New("client transport is nil")
240	252		}
241	253		return tr.Send(ac.nextOutputFrameLocked(b))
242	254		})
	255	+	usedMu.Lock()
	256	+	defer usedMu.Unlock()
	257	+	return used, err
243	258		}
244	259
245	260		func (d Daemon) boundedSendWith(ac attachedClient, send func() error) error {

 }
 func (d *Daemon) boundedSendOutputErr(ac *attachedClient, b []byte) error {
-	return d.boundedSendWith(ac, func() error {
+	_, err := d.boundedSendOutputErrTransport(ac, b)
+	return err
+}
+func (d *Daemon) boundedSendOutputErrTransport(ac *attachedClient, b []byte) (ports.Transport, error) {
+	var (
+		usedMu sync.Mutex
+		used   ports.Transport
+	)
+	err := d.boundedSendWith(ac, func() error {
 		ac.sendMu.Lock()
 		defer ac.sendMu.Unlock()
 		tr := ac.transport()
+		usedMu.Lock()
+		used = tr
+		usedMu.Unlock()
 		if tr == nil {
 			return errors.New("client transport is nil")
 		}
 		return tr.Send(ac.nextOutputFrameLocked(b))
 	})
+	usedMu.Lock()
+	defer usedMu.Unlock()
+	return used, err
 }
 func (d *Daemon) boundedSendWith(ac *attachedClient, send func() error) error {

M internal/usecase/daemon/copymode.go +2 -2


@@ -252,8 +252,8 @@ func (d Daemon) handleCopyInput(ac attachedClient, data []byte) {
252	252		if copyOut && text != "" {
253	253		chunks := scopy.OSC52(text)
254	254		for _, chunk := range chunks {
255		-	failed := ac.transport()
256		-	if err := d.boundedSendOutputErr(ac, chunk); err != nil {
	255	+	failed, err := d.boundedSendOutputErrTransport(ac, chunk)
	256	+	if err != nil {
257	257		d.detachOnSendError(sess, ac, failed)
258	258		return
259	259		}


@@ -252,8 +252,8 @@ func (d Daemon) handleCopyInput(ac attachedClient, data []byte) {
252		if copyOut && text != "" {	252		if copyOut && text != "" {
253		chunks := scopy.OSC52(text)	253		chunks := scopy.OSC52(text)
254		for _, chunk := range chunks {	254		for _, chunk := range chunks {
255	-	failed := ac.transport()	255	+	failed, err := d.boundedSendOutputErrTransport(ac, chunk)
256	-	if err := d.boundedSendOutputErr(ac, chunk); err != nil {	256	+	if err != nil {
257		d.detachOnSendError(sess, ac, failed)	257		d.detachOnSendError(sess, ac, failed)
258		return	258		return
259		}	259		}

M internal/usecase/daemon/copymode_test.go +29 -0


@@ -23,6 +23,35 @@ import (
23	23
24	24		// --- test doubles -----------------------------------------------------------
25	25
	26	+	func TestBoundedSendOutputErrTransportReturnsTransportUsedBySend(t *testing.T) {
	27	+	d := newTestDaemon(t, nil, stubClock{})
	28	+	ac := &attachedClient{}
	29	+	replacement := &closeTrackingTransport{}
	30	+	sendErr := errors.New("send failed")
	31	+	failed := &swapErrorTransport{ac: ac, replacement: replacement, err: sendErr}
	32	+	ac.replaceTransport(failed)
	33	+
	34	+	used, err := d.boundedSendOutputErrTransport(ac, []byte("copy"))
	35	+
	36	+	require.ErrorIs(t, err, sendErr)
	37	+	require.Same(t, failed, used)
	38	+	require.Same(t, replacement, ac.transport())
	39	+	}
	40	+
	41	+	type swapErrorTransport struct {
	42	+	ac *attachedClient
	43	+	replacement ports.Transport
	44	+	err error
	45	+	}
	46	+
	47	+	func (t *swapErrorTransport) Send(ports.Frame) error {
	48	+	t.ac.replaceTransport(t.replacement)
	49	+	return t.err
	50	+	}
	51	+
	52	+	func (t *swapErrorTransport) Recv() (ports.Frame, error) { return ports.Frame{}, io.EOF }
	53	+	func (t *swapErrorTransport) Close() error { return nil }
	54	+
26	55		// stubClock returns timers whose channel never fires, so a scheduler under it
27	56		// blocks in its debounce loop until the session context is cancelled. Used by
28	57


@@ -23,6 +23,35 @@ import (
23		23
24	// --- test doubles -----------------------------------------------------------	24		// --- test doubles -----------------------------------------------------------
25		25
		26	+	func TestBoundedSendOutputErrTransportReturnsTransportUsedBySend(t *testing.T) {
		27	+	d := newTestDaemon(t, nil, stubClock{})
		28	+	ac := &attachedClient{}
		29	+	replacement := &closeTrackingTransport{}
		30	+	sendErr := errors.New("send failed")
		31	+	failed := &swapErrorTransport{ac: ac, replacement: replacement, err: sendErr}
		32	+	ac.replaceTransport(failed)
		33	+
		34	+	used, err := d.boundedSendOutputErrTransport(ac, []byte("copy"))
		35	+
		36	+	require.ErrorIs(t, err, sendErr)
		37	+	require.Same(t, failed, used)
		38	+	require.Same(t, replacement, ac.transport())
		39	+	}
		40	+
		41	+	type swapErrorTransport struct {
		42	+	ac *attachedClient
		43	+	replacement ports.Transport
		44	+	err error
		45	+	}
		46	+
		47	+	func (t *swapErrorTransport) Send(ports.Frame) error {
		48	+	t.ac.replaceTransport(t.replacement)
		49	+	return t.err
		50	+	}
		51	+
		52	+	func (t *swapErrorTransport) Recv() (ports.Frame, error) { return ports.Frame{}, io.EOF }
		53	+	func (t *swapErrorTransport) Close() error { return nil }
		54	+
26	// stubClock returns timers whose channel never fires, so a scheduler under it	55		// stubClock returns timers whose channel never fires, so a scheduler under it
27	// blocks in its debounce loop until the session context is cancelled. Used by	56		// blocks in its debounce loop until the session context is cancelled. Used by
28		57

M internal/usecase/daemon/resume.go +5 -6


@@ -35,14 +35,13 @@ func (d Daemon) prepareParkAttachment(sess session, ac *attachedClient) bool {
35	35		if ephemeral \|\| !ac.resumeCapable {
36	36		return false
37	37		}
	38	+	d.mu.Lock()
	39	+	defer d.mu.Unlock()
	40	+	if d.closing {
	41	+	return false
	42	+	}
38	43		if ac.resumeToken == 0 {
39		-	d.mu.Lock()
40		-	if d.closing {
41		-	d.mu.Unlock()
42		-	return false
43		-	}
44	44		ac.resumeToken = d.nextResumeTokenLocked()
45		-	d.mu.Unlock()
46	45		}
47	46		return true
48	47		}

 	if ephemeral || !ac.resumeCapable {
 		return false
 	}
+	d.mu.Lock()
+	defer d.mu.Unlock()
+	if d.closing {
+		return false
+	}
 	if ac.resumeToken == 0 {
-		d.mu.Lock()
-		if d.closing {
-			d.mu.Unlock()
-			return false
-		}
 		ac.resumeToken = d.nextResumeTokenLocked()
-		d.mu.Unlock()
 	}
 	return true
 }

M pkg/dgram/dgram.go +25 -23


@@ -79,9 +79,11 @@ func (c Codec) Open(packet []byte, wantDirection uint32, aad []byte, replay Re
79	79		return ctr, pt, nil
80	80		}
81	81
	82	+	const replayWindowSize = maxFragmentCount
	83	+
82	84		type ReplayWindow struct {
83	85		max uint64
84		-	seen uint64
	86	+	seen map[uint64]struct{}
85	87		init bool
86	88		}
87	89

@@ -94,31 +96,28 @@ func (w *ReplayWindow) Check(counter uint64) bool {
94	96		if counter > w.max {
95	97		return true
96	98		}
97		-	delta := w.max - counter
98		-	if delta >= 64 {
	99	+	if w.max-counter >= uint64(replayWindowSize) {
99	100		return false
100	101		}
101		-	return (w.seen & (uint64(1) << delta)) == 0
	102	+	_, ok := w.seen[counter]
	103	+	return !ok
102	104		}
103	105
104	106		func (w *ReplayWindow) Accept(counter uint64) {
	107	+	if w.seen == nil {
	108	+	w.seen = make(map[uint64]struct{})
	109	+	}
105	110		if !w.init {
106		-	w.max, w.seen, w.init = counter, 1, true
107		-	return
	111	+	w.max, w.init = counter, true
108	112		}
109	113		if counter > w.max {
110		-	shift := counter - w.max
111		-	if shift >= 64 {
112		-	w.seen = 1
113		-	} else {
114		-	w.seen = (w.seen << shift) \| 1
115		-	}
116	114		w.max = counter
117		-	return
118	115		}
119		-	delta := w.max - counter
120		-	if delta < 64 {
121		-	w.seen \|= uint64(1) << delta
	116	+	w.seen[counter] = struct{}{}
	117	+	for seen := range w.seen {
	118	+	if w.max-seen >= uint64(replayWindowSize) {
	119	+	delete(w.seen, seen)
	120	+	}
122	121		}
123	122		}
124	123

@@ -133,18 +132,18 @@ func FragmentPayload(seq uint64, payload []byte, mtu int) ([]Fragment, error) {
133	132		if mtu <= fragmentHdr {
134	133		return nil, ErrFragment
135	134		}
136		-	max := mtu - fragmentHdr
137		-	count := (len(payload) + max - 1) / max
	135	+	maxData := mtu - fragmentHdr
	136	+	count := (len(payload) + maxData - 1) / maxData
138	137		if count == 0 {
139	138		count = 1
140	139		}
141		-	if count > 0xffff {
	140	+	if count > maxFragmentCount {
142	141		return nil, ErrFragment
143	142		}
144	143		out := make([]Fragment, count)
145	144		for i := 0; i < count; i++ {
146		-	start := i * max
147		-	end := min(len(payload), start+max)
	145	+	start := i * maxData
	146	+	end := min(len(payload), start+maxData)
148	147		out[i] = Fragment{Seq: seq, Index: uint16(i), Count: uint16(count), Data: append([]byte(nil), payload[start:end]...)}
149	148		}
150	149		return out, nil

@@ -176,7 +175,10 @@ func UnmarshalFragment(b []byte) (Fragment, error) {
176	175		return f, nil
177	176		}
178	177
179		-	const maxReassemblyInflight = 1024
	178	+	const (
	179	+	maxFragmentCount = 1024
	180	+	maxReassemblyInflight = 1024
	181	+	)
180	182
181	183		type Reassembler struct {
182	184		inflight map[uint64]*assembly

@@ -190,7 +192,7 @@ type assembly struct {
190	192
191	193		func NewReassembler() Reassembler { return &Reassembler{inflight: make(map[uint64]assembly)} }
192	194		func (r *Reassembler) Add(f Fragment) ([]byte, bool, error) {
193		-	if f.Count == 0 \|\| f.Index >= f.Count {
	195	+	if f.Count == 0 \|\| int(f.Count) > maxFragmentCount \|\| f.Index >= f.Count {
194	196		return nil, false, ErrFragment
195	197		}
196	198		a := r.inflight[f.Seq]

 	return ctr, pt, nil
 }
+const replayWindowSize = maxFragmentCount
 type ReplayWindow struct {
 	max  uint64
-	seen uint64
+	seen map[uint64]struct{}
 	init bool
 }
 	if counter > w.max {
 		return true
 	}
-	delta := w.max - counter
+	if w.max-counter >= uint64(replayWindowSize) {
-	if delta >= 64 {
 		return false
 	}
-	return (w.seen & (uint64(1) << delta)) == 0
+	_, ok := w.seen[counter]
+	return !ok
 }
 func (w *ReplayWindow) Accept(counter uint64) {
+	if w.seen == nil {
+		w.seen = make(map[uint64]struct{})
+	}
 	if !w.init {
-		w.max, w.seen, w.init = counter, 1, true
+		w.max, w.init = counter, true
-		return
 	}
 	if counter > w.max {
-		shift := counter - w.max
-		if shift >= 64 {
-			w.seen = 1
-		} else {
-			w.seen = (w.seen << shift) | 1
-		}
 		w.max = counter
-		return
 	}
-	delta := w.max - counter
+	w.seen[counter] = struct{}{}
-	if delta < 64 {
+	for seen := range w.seen {
-		w.seen |= uint64(1) << delta
+		if w.max-seen >= uint64(replayWindowSize) {
+			delete(w.seen, seen)
+		}
 	}
 }
 	if mtu <= fragmentHdr {
 		return nil, ErrFragment
 	}
-	max := mtu - fragmentHdr
+	maxData := mtu - fragmentHdr
-	count := (len(payload) + max - 1) / max
+	count := (len(payload) + maxData - 1) / maxData
 	if count == 0 {
 		count = 1
 	}
-	if count > 0xffff {
+	if count > maxFragmentCount {
 		return nil, ErrFragment
 	}
 	out := make([]Fragment, count)
 	for i := 0; i < count; i++ {
-		start := i * max
+		start := i * maxData
-		end := min(len(payload), start+max)
+		end := min(len(payload), start+maxData)
 		out[i] = Fragment{Seq: seq, Index: uint16(i), Count: uint16(count), Data: append([]byte(nil), payload[start:end]...)}
 	}
 	return out, nil
 	return f, nil
 }
-const maxReassemblyInflight = 1024
+const (
+	maxFragmentCount      = 1024
+	maxReassemblyInflight = 1024
+)
 type Reassembler struct {
 	inflight map[uint64]*assembly
 func NewReassembler() *Reassembler { return &Reassembler{inflight: make(map[uint64]*assembly)} }
 func (r *Reassembler) Add(f Fragment) ([]byte, bool, error) {
-	if f.Count == 0 || f.Index >= f.Count {
+	if f.Count == 0 || int(f.Count) > maxFragmentCount || f.Index >= f.Count {
 		return nil, false, ErrFragment
 	}
 	a := r.inflight[f.Seq]

M pkg/dgram/dgram_test.go +36 -0


@@ -56,6 +56,25 @@ func TestReplayReject(t *testing.T) {
56	56		}
57	57		}
58	58
	59	+	func TestReplayWindowAllowsLargeFragmentReordering(t *testing.T) {
	60	+	c := testCodec(t)
	61	+	rw := NewReplayWindow()
	62	+	last := uint64(maxFragmentCount)
	63	+	if _, _, err := c.Open(c.Seal(1, last, []byte("last"), nil), 1, nil, rw); err != nil {
	64	+	t.Fatal(err)
	65	+	}
	66	+	if _, _, err := c.Open(c.Seal(1, 1, []byte("first"), nil), 1, nil, rw); err != nil {
	67	+	t.Fatalf("fragment within max fragment window rejected: %v", err)
	68	+	}
	69	+	if _, _, err := c.Open(c.Seal(1, last+2, []byte("advance"), nil), 1, nil, rw); err != nil {
	70	+	t.Fatal(err)
	71	+	}
	72	+	tooOld := c.Seal(1, 0, []byte("old"), nil)
	73	+	if _, _, err := c.Open(tooOld, 1, nil, rw); !errors.Is(err, ErrReplay) {
	74	+	t.Fatalf("err=%v, want ErrReplay", err)
	75	+	}
	76	+	}
	77	+
59	78		func TestFragmentation(t *testing.T) {
60	79		payload := bytes.Repeat([]byte("abcdef"), 400)
61	80		frags, err := FragmentPayload(9, payload, 100)

@@ -104,6 +123,23 @@ func TestReassemblyDisorderDuplicates(t *testing.T) {
104	123		}
105	124		}
106	125
	126	+	func TestReassemblyRejectsOversizedCountBeforeInsert(t *testing.T) {
	127	+	r := NewReassembler()
	128	+	if _, _, err := r.Add(Fragment{Seq: 99, Index: 0, Count: uint16(maxFragmentCount + 1), Data: []byte("x")}); !errors.Is(err, ErrFragment) {
	129	+	t.Fatalf("err=%v, want ErrFragment", err)
	130	+	}
	131	+	if len(r.inflight) != 0 \|\| len(r.order) != 0 {
	132	+	t.Fatalf("oversized fragment was inserted: inflight=%d order=%d", len(r.inflight), len(r.order))
	133	+	}
	134	+	}
	135	+
	136	+	func TestFragmentPayloadRejectsOverMaxFragmentCount(t *testing.T) {
	137	+	payload := bytes.Repeat([]byte("x"), maxFragmentCount+1)
	138	+	if _, err := FragmentPayload(1, payload, fragmentHdr+1); !errors.Is(err, ErrFragment) {
	139	+	t.Fatalf("err=%v, want ErrFragment", err)
	140	+	}
	141	+	}
	142	+
107	143		func TestReassemblyCountMismatchRemovesStaleAssembly(t *testing.T) {
108	144		r := NewReassembler()
109	145		if _, _, err := r.Add(Fragment{Seq: 7, Index: 0, Count: 2, Data: []byte("stale")}); err != nil {


@@ -56,6 +56,25 @@ func TestReplayReject(t *testing.T) {
56	}	56		}
57	}	57		}
58		58
		59	+	func TestReplayWindowAllowsLargeFragmentReordering(t *testing.T) {
		60	+	c := testCodec(t)
		61	+	rw := NewReplayWindow()
		62	+	last := uint64(maxFragmentCount)
		63	+	if _, _, err := c.Open(c.Seal(1, last, []byte("last"), nil), 1, nil, rw); err != nil {
		64	+	t.Fatal(err)
		65	+	}
		66	+	if _, _, err := c.Open(c.Seal(1, 1, []byte("first"), nil), 1, nil, rw); err != nil {
		67	+	t.Fatalf("fragment within max fragment window rejected: %v", err)
		68	+	}
		69	+	if _, _, err := c.Open(c.Seal(1, last+2, []byte("advance"), nil), 1, nil, rw); err != nil {
		70	+	t.Fatal(err)
		71	+	}
		72	+	tooOld := c.Seal(1, 0, []byte("old"), nil)
		73	+	if _, _, err := c.Open(tooOld, 1, nil, rw); !errors.Is(err, ErrReplay) {
		74	+	t.Fatalf("err=%v, want ErrReplay", err)
		75	+	}
		76	+	}
		77	+
59	func TestFragmentation(t *testing.T) {	78		func TestFragmentation(t *testing.T) {
60	payload := bytes.Repeat([]byte("abcdef"), 400)	79		payload := bytes.Repeat([]byte("abcdef"), 400)
61	frags, err := FragmentPayload(9, payload, 100)	80		frags, err := FragmentPayload(9, payload, 100)

@@ -104,6 +123,23 @@ func TestReassemblyDisorderDuplicates(t *testing.T) {
104	}	123		}
105	}	124		}
106		125
		126	+	func TestReassemblyRejectsOversizedCountBeforeInsert(t *testing.T) {
		127	+	r := NewReassembler()
		128	+	if _, _, err := r.Add(Fragment{Seq: 99, Index: 0, Count: uint16(maxFragmentCount + 1), Data: []byte("x")}); !errors.Is(err, ErrFragment) {
		129	+	t.Fatalf("err=%v, want ErrFragment", err)
		130	+	}
		131	+	if len(r.inflight) != 0 \|\| len(r.order) != 0 {
		132	+	t.Fatalf("oversized fragment was inserted: inflight=%d order=%d", len(r.inflight), len(r.order))
		133	+	}
		134	+	}
		135	+
		136	+	func TestFragmentPayloadRejectsOverMaxFragmentCount(t *testing.T) {
		137	+	payload := bytes.Repeat([]byte("x"), maxFragmentCount+1)
		138	+	if _, err := FragmentPayload(1, payload, fragmentHdr+1); !errors.Is(err, ErrFragment) {
		139	+	t.Fatalf("err=%v, want ErrFragment", err)
		140	+	}
		141	+	}
		142	+
107	func TestReassemblyCountMismatchRemovesStaleAssembly(t *testing.T) {	143		func TestReassemblyCountMismatchRemovesStaleAssembly(t *testing.T) {
108	r := NewReassembler()	144		r := NewReassembler()
109	if _, _, err := r.Add(Fragment{Seq: 7, Index: 0, Count: 2, Data: []byte("stale")}); err != nil {	145		if _, _, err := r.Add(Fragment{Seq: 7, Index: 0, Count: 2, Data: []byte("stale")}); err != nil {