Skip to content

Commit 3701052

Browse files
committed
TUN-8775: Make sure the session Close can only be called once
The previous capture of the sync.OnceValue was re-initialized for each call to `Close`. This needed to be initialized during the creation of the session to ensure that the sync.OnceValue reference was held for the session's lifetime. Closes TUN-8775
1 parent f07d04d commit 3701052

File tree

2 files changed

+19
-9
lines changed

2 files changed

+19
-9
lines changed

quic/v3/session.go

Lines changed: 14 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -73,6 +73,9 @@ type session struct {
7373
contextChan chan context.Context
7474
metrics Metrics
7575
log *zerolog.Logger
76+
77+
// A special close function that we wrap with sync.Once to make sure it is only called once
78+
closeFn func() error
7679
}
7780

7881
func NewSession(
@@ -86,6 +89,7 @@ func NewSession(
8689
log *zerolog.Logger,
8790
) Session {
8891
logger := log.With().Str(logFlowID, id.String()).Logger()
92+
closeChan := make(chan error, 1)
8993
session := &session{
9094
id: id,
9195
closeAfterIdle: closeAfterIdle,
@@ -96,11 +100,19 @@ func NewSession(
96100
// activeAtChan has low capacity. It can be full when there are many concurrent read/write. markActive() will
97101
// drop instead of blocking because last active time only needs to be an approximation
98102
activeAtChan: make(chan time.Time, 1),
99-
closeChan: make(chan error, 1),
103+
closeChan: closeChan,
100104
// contextChan is an unbounded channel to help enforce one active migration of a session at a time.
101105
contextChan: make(chan context.Context),
102106
metrics: metrics,
103107
log: &logger,
108+
closeFn: sync.OnceValue(func() error {
109+
// We don't want to block on sending to the close channel if it is already full
110+
select {
111+
case closeChan <- SessionCloseErr:
112+
default:
113+
}
114+
return origin.Close()
115+
}),
104116
}
105117
session.eyeball.Store(&eyeball)
106118
return session
@@ -218,14 +230,7 @@ func (s *session) markActive() {
218230

219231
func (s *session) Close() error {
220232
// Make sure that we only close the origin connection once
221-
return sync.OnceValue(func() error {
222-
// We don't want to block on sending to the close channel if it is already full
223-
select {
224-
case s.closeChan <- SessionCloseErr:
225-
default:
226-
}
227-
return s.origin.Close()
228-
})()
233+
return s.closeFn()
229234
}
230235

231236
func (s *session) waitForCloseCondition(ctx context.Context, closeAfterIdle time.Duration) error {

quic/v3/session_test.go

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -255,11 +255,16 @@ func TestSessionClose_Multiple(t *testing.T) {
255255
if !origin.closed.Load() {
256256
t.Fatal("origin wasn't closed")
257257
}
258+
// Reset the closed status to make sure it isn't closed again
259+
origin.closed.Store(false)
258260
// subsequent closes shouldn't call close again or cause any errors
259261
err = session.Close()
260262
if err != nil {
261263
t.Fatal(err)
262264
}
265+
if origin.closed.Load() {
266+
t.Fatal("origin was incorrectly closed twice")
267+
}
263268
}
264269

265270
func TestSessionServe_IdleTimeout(t *testing.T) {

0 commit comments

Comments
 (0)