TUN-5494: Send a RPC with terminate reason to edge if the session is closed locally

datagramsession/event.go

@@ -1,6 +1,7 @@
 package datagramsession
 
 import (
+	"fmt"
 	"io"
 
 	"github.com/google/uuid"
@@ -24,6 +25,22 @@ func newRegisterSessionEvent(sessionID uuid.UUID, originProxy io.ReadWriteCloser
 // unregisterSessionEvent is an event to stop tracking and terminate the session.
 type unregisterSessionEvent struct {
 	sessionID uuid.UUID
+	err       *errClosedSession
+}
+
+// ClosedSessionError represent a condition that closes the session other than I/O
+// I/O error is not included, because the side that closes the session is ambiguous.
+type errClosedSession struct {
+	message  string
+	byRemote bool
+}
+
+func (sc *errClosedSession) Error() string {
+	if sc.byRemote {
+		return fmt.Sprintf("session closed by remote due to %s", sc.message)
+	} else {
+		return fmt.Sprintf("session closed by local due to %s", sc.message)
+	}
 }
 
 // newDatagram is an event when transport receives new datagram

datagramsession/manager.go

@@ -20,7 +20,7 @@ type Manager interface {
 	// RegisterSession starts tracking a session. Caller is responsible for starting the session
 	RegisterSession(ctx context.Context, sessionID uuid.UUID, dstConn io.ReadWriteCloser) (*Session, error)
 	// UnregisterSession stops tracking the session and terminates it
-	UnregisterSession(ctx context.Context, sessionID uuid.UUID) error
+	UnregisterSession(ctx context.Context, sessionID uuid.UUID, message string, byRemote bool) error
 }
 
 type manager struct {
@@ -100,8 +100,14 @@ func (m *manager) registerSession(ctx context.Context, registration *registerSes
 	registration.resultChan <- session
 }
 
-func (m *manager) UnregisterSession(ctx context.Context, sessionID uuid.UUID) error {
-	event := &unregisterSessionEvent{sessionID: sessionID}
+func (m *manager) UnregisterSession(ctx context.Context, sessionID uuid.UUID, message string, byRemote bool) error {
+	event := &unregisterSessionEvent{
+		sessionID: sessionID,
+		err: &errClosedSession{
+			message:  message,
+			byRemote: byRemote,
+		},
+	}
 	select {
 	case <-ctx.Done():
 		return ctx.Err()
@@ -114,7 +120,7 @@ func (m *manager) unregisterSession(unregistration *unregisterSessionEvent) {
 	session, ok := m.sessions[unregistration.sessionID]
 	if ok {
 		delete(m.sessions, unregistration.sessionID)
-		session.close()
+		session.close(unregistration.err)
 	}
 }
 

datagramsession/manager_test.go

@@ -17,8 +17,9 @@ import (
 
 func TestManagerServe(t *testing.T) {
 	const (
-		sessions = 20
-		msgs     = 50
+		sessions            = 20
+		msgs                = 50
+		remoteUnregisterMsg = "eyeball closed connection"
 	)
 	log := zerolog.Nop()
 	transport := &mockQUICTransport{
@@ -89,7 +90,13 @@ func TestManagerServe(t *testing.T) {
 
 			sessionDone := make(chan struct{})
 			go func() {
-				session.Serve(ctx, time.Minute*2)
+				closedByRemote, err := session.Serve(ctx, time.Minute*2)
+				closeSession := &errClosedSession{
+					message:  remoteUnregisterMsg,
+					byRemote: true,
+				}
+				require.Equal(t, closeSession, err)
+				require.True(t, closedByRemote)
 				close(sessionDone)
 			}()
 
@@ -100,7 +107,7 @@ func TestManagerServe(t *testing.T) {
 			// Make sure eyeball and origin have received all messages before unregistering the session
 			require.NoError(t, reqErrGroup.Wait())
 
-			require.NoError(t, mg.UnregisterSession(ctx, sessionID))
+			require.NoError(t, mg.UnregisterSession(ctx, sessionID, remoteUnregisterMsg, true))
 			<-sessionDone
 
 			return nil

datagramsession/session.go

@@ -2,6 +2,7 @@ package datagramsession
 
 import (
 	"context"
+	"fmt"
 	"io"
 	"time"
 
@@ -12,6 +13,10 @@ const (
 	defaultCloseIdleAfter = time.Second * 210
 )
 
+func SessionIdleErr(timeout time.Duration) error {
+	return fmt.Errorf("session idle for %v", timeout)
+}
+
 // Each Session is a bidirectional pipe of datagrams between transport and dstConn
 // Currently the only implementation of transport is quic DatagramMuxer
 // Destination can be a connection with origin or with eyeball
@@ -24,47 +29,53 @@ const (
 // - Datagrams from cloudflared are read by Manager from the transport. Manager finds the corresponding Session and calls the
 //   write method of the Session to send to eyeball
 type Session struct {
-	id        uuid.UUID
+	ID        uuid.UUID
 	transport transport
 	dstConn   io.ReadWriteCloser
 	// activeAtChan is used to communicate the last read/write time
 	activeAtChan chan time.Time
-	doneChan     chan struct{}
+	closeChan    chan error
 }
 
 func newSession(id uuid.UUID, transport transport, dstConn io.ReadWriteCloser) *Session {
 	return &Session{
-		id:        id,
+		ID:        id,
 		transport: transport,
 		dstConn:   dstConn,
 		// activeAtChan has low capacity. It can be full when there are many concurrent read/write. markActive() will
 		// drop instead of blocking because last active time only needs to be an approximation
 		activeAtChan: make(chan time.Time, 2),
-		doneChan:     make(chan struct{}),
+		// capacity is 2 because close() and dstToTransport routine in Serve() can write to this channel
+		closeChan: make(chan error, 2),
 	}
 }
 
-func (s *Session) Serve(ctx context.Context, closeAfterIdle time.Duration) error {
-	serveCtx, cancel := context.WithCancel(ctx)
-	defer cancel()
-	go s.waitForCloseCondition(serveCtx, closeAfterIdle)
-	// QUIC implementation copies data to another buffer before returning https://github.com/lucas-clemente/quic-go/blob/v0.24.0/session.go#L1967-L1975
-	// This makes it safe to share readBuffer between iterations
-	readBuffer := make([]byte, s.transport.MTU())
-	for {
-		if err := s.dstToTransport(readBuffer); err != nil {
-			return err
+func (s *Session) Serve(ctx context.Context, closeAfterIdle time.Duration) (closedByRemote bool, err error) {
+	go func() {
+		// QUIC implementation copies data to another buffer before returning https://github.com/lucas-clemente/quic-go/blob/v0.24.0/session.go#L1967-L1975
+		// This makes it safe to share readBuffer between iterations
+		readBuffer := make([]byte, s.transport.MTU())
+		for {
+			if err := s.dstToTransport(readBuffer); err != nil {
+				s.closeChan <- err
+				return
+			}
 		}
+	}()
+	err = s.waitForCloseCondition(ctx, closeAfterIdle)
+	if closeSession, ok := err.(*errClosedSession); ok {
+		closedByRemote = closeSession.byRemote
 	}
+	return closedByRemote, err
 }
 
-func (s *Session) waitForCloseCondition(ctx context.Context, closeAfterIdle time.Duration) {
+func (s *Session) waitForCloseCondition(ctx context.Context, closeAfterIdle time.Duration) error {
+	// Closing dstConn cancels read so dstToTransport routine in Serve() can return
+	defer s.dstConn.Close()
 	if closeAfterIdle == 0 {
 		// provide deafult is caller doesn't specify one
 		closeAfterIdle = defaultCloseIdleAfter
 	}
-	// Closing dstConn cancels read so Serve function can return
-	defer s.dstConn.Close()
 
 	checkIdleFreq := closeAfterIdle / 8
 	checkIdleTicker := time.NewTicker(checkIdleFreq)
@@ -74,14 +85,14 @@ func (s *Session) waitForCloseCondition(ctx context.Context, closeAfterIdle time
 	for {
 		select {
 		case <-ctx.Done():
-			return
-		case <-s.doneChan:
-			return
+			return ctx.Err()
+		case reason := <-s.closeChan:
+			return reason
 		// TODO: TUN-5423 evaluate if using atomic is more efficient
 		case now := <-checkIdleTicker.C:
 			// The session is considered inactive if current time is after (last active time + allowed idle time)
 			if now.After(activeAt.Add(closeAfterIdle)) {
-				return
+				return SessionIdleErr(closeAfterIdle)
 			}
 		case activeAt = <-s.activeAtChan: // Update last active time
 		}
@@ -92,7 +103,7 @@ func (s *Session) dstToTransport(buffer []byte) error {
 	n, err := s.dstConn.Read(buffer)
 	s.markActive()
 	if n > 0 {
-		if err := s.transport.SendTo(s.id, buffer[:n]); err != nil {
+		if err := s.transport.SendTo(s.ID, buffer[:n]); err != nil {
 			return err
 		}
 	}
@@ -113,6 +124,6 @@ func (s *Session) markActive() {
 	}
 }
 
-func (s *Session) close() {
-	close(s.doneChan)
+func (s *Session) close(err *errClosedSession) {
+	s.closeChan <- err
 }

datagramsession/session_test.go

@@ -31,6 +31,12 @@ func TestCloseIdle(t *testing.T) {
 }
 
 func testSessionReturns(t *testing.T, closeBy closeMethod, closeAfterIdle time.Duration) {
+	var (
+		localCloseReason = &errClosedSession{
+			message:  "connection closed by origin",
+			byRemote: false,
+		}
+	)
 	sessionID := uuid.New()
 	cfdConn, originConn := net.Pipe()
 	payload := testPayload(sessionID)
@@ -43,7 +49,18 @@ func testSessionReturns(t *testing.T, closeBy closeMethod, closeAfterIdle time.D
 	ctx, cancel := context.WithCancel(context.Background())
 	sessionDone := make(chan struct{})
 	go func() {
-		session.Serve(ctx, closeAfterIdle)
+		closedByRemote, err := session.Serve(ctx, closeAfterIdle)
+		switch closeBy {
+		case closeByContext:
+			require.Equal(t, context.Canceled, err)
+			require.False(t, closedByRemote)
+		case closeByCallingClose:
+			require.Equal(t, localCloseReason, err)
+			require.Equal(t, localCloseReason.byRemote, closedByRemote)
+		case closeByTimeout:
+			require.Equal(t, SessionIdleErr(closeAfterIdle), err)
+			require.False(t, closedByRemote)
+		}
 		close(sessionDone)
 	}()
 
@@ -64,7 +81,7 @@ func testSessionReturns(t *testing.T, closeBy closeMethod, closeAfterIdle time.D
 	case closeByContext:
 		cancel()
 	case closeByCallingClose:
-		session.close()
+		session.close(localCloseReason)
 	}
 
 	<-sessionDone