fix: register session cleanup before slot is visible (zombie session race)

foodgroup/auth.go

@@ -81,8 +81,8 @@ type AuthService struct {
 // This method does not verify that the user and chat room exist because it
 // implicitly trusts the contents of the token signed by
 // {{OServiceService.ServiceRequest}}.
-func (s AuthService) RegisterChatSession(ctx context.Context, serverCookie state.ServerCookie) (*state.SessionInstance, error) {
-	sess, err := s.chatSessionRegistry.AddSession(ctx, serverCookie.ChatCookie, serverCookie.ScreenName)
+func (s AuthService) RegisterChatSession(ctx context.Context, authCookie state.ServerCookie, sessCfg func(sess *state.Session)) (*state.SessionInstance, error) {
+	sess, err := s.chatSessionRegistry.AddSession(ctx, authCookie.ChatCookie, authCookie.ScreenName, sessCfg)
 	if err != nil {
 		return nil, fmt.Errorf("AddSession: %w", err)
 	}
@@ -108,8 +108,9 @@ func (s AuthService) CrackCookie(authCookie []byte) (state.ServerCookie, error)
 }
 
 // RegisterBOSSession adds a new session to the session registry.
-func (s AuthService) RegisterBOSSession(ctx context.Context, serverCookie state.ServerCookie) (*state.SessionInstance, error) {
-	u, err := s.userManager.User(ctx, serverCookie.ScreenName.IdentScreenName())
+func (s AuthService) RegisterBOSSession(ctx context.Context, authCookie state.ServerCookie, sessCfg func(sess *state.Session)) (*state.SessionInstance, error) {
+
+	u, err := s.userManager.User(ctx, authCookie.ScreenName.IdentScreenName())
 	if err != nil {
 		return nil, fmt.Errorf("failed to retrieve user: %w", err)
 	}
@@ -120,14 +121,20 @@ func (s AuthService) RegisterBOSSession(ctx context.Context, serverCookie state.
 	ctx, cancel := context.WithTimeout(ctx, time.Second*5)
 	defer cancel()
 
-	flag := wire.MultiConnFlag(serverCookie.MultiConnFlag)
+	flag := wire.MultiConnFlag(authCookie.MultiConnFlag)
 
 	doMultiSess := false
 	if flag == wire.MultiConnFlagsRecentClient {
 		doMultiSess = true
 	}
 
-	sess, err := s.sessionManager.AddSession(ctx, u.DisplayScreenName, doMultiSess)
+	cfg := func(sess *state.Session) {
+		sess.SetSignonTime(time.Now())
+		sess.SetRateClasses(time.Now(), s.rateLimitClasses)
+		sess.SetMemberSince(time.Now())
+	}
+
+	sess, err := s.sessionManager.AddSession(ctx, u.DisplayScreenName, doMultiSess, sessCfg, cfg)
 	if err != nil {
 		return nil, fmt.Errorf("AddSession: %w", err)
 	}
@@ -143,12 +150,10 @@ func (s AuthService) RegisterBOSSession(ctx context.Context, serverCookie state.
 		sess.SetUserInfoFlag(wire.OServiceUserFlagBot)
 	}
 
-	sess.SetKerberosAuth(serverCookie.KerberosAuth == 1)
-	sess.Session().SetSignonTime(time.Now())
-	sess.Session().SetRateClasses(time.Now(), s.rateLimitClasses)
+	sess.SetKerberosAuth(authCookie.KerberosAuth == 1)
+
 	// set string containing OSCAR client name and version
-	sess.SetClientID(serverCookie.ClientID)
-	sess.Session().SetMemberSince(time.Now())
+	sess.SetClientID(authCookie.ClientID)
 	sess.Session().SetOfflineMsgCount(u.OfflineMsgCount)
 
 	if _, alreadySet := sess.Session().BuddyIcon(); !alreadySet {
@@ -196,15 +201,18 @@ func (s AuthService) RetrieveBOSSession(ctx context.Context, serverCookie state.
 }
 
 // Signout removes this user's session.
-func (s AuthService) Signout(ctx context.Context, instance *state.SessionInstance) {
-	s.sessionManager.RemoveSession(instance)
+func (s AuthService) Signout(ctx context.Context, session *state.Session) {
+	s.sessionManager.RemoveSession(session)
 }
 
 // SignoutChat removes user from chat room and notifies remaining participants
 // of their departure.
-func (s AuthService) SignoutChat(ctx context.Context, instance *state.SessionInstance) {
-	alertUserLeft(ctx, instance, s.chatMessageRelayer)
-	s.chatSessionRegistry.RemoveSession(instance)
+func (s AuthService) SignoutChat(ctx context.Context, sess *state.Session) {
+	instances := sess.Instances()
+	for _, instance := range instances {
+		alertUserLeft(ctx, instance, s.chatMessageRelayer)
+	}
+	s.chatSessionRegistry.RemoveSession(sess)
 }
 
 // BUCPChallenge processes a BUCP authentication challenge request. It

foodgroup/auth_test.go

@@ -1745,7 +1745,7 @@ func TestAuthService_RegisterChatSession_HappyPath(t *testing.T) {
 
 	chatSessionRegistry := newMockChatSessionRegistry(t)
 	chatSessionRegistry.EXPECT().
-		AddSession(mock.Anything, serverCookie.ChatCookie, instance.DisplayScreenName()).
+		AddSession(mock.Anything, serverCookie.ChatCookie, instance.DisplayScreenName(), mock.Anything).
 		Return(instance, nil)
 
 	chatCookieBuf := &bytes.Buffer{}
@@ -1753,7 +1753,7 @@ func TestAuthService_RegisterChatSession_HappyPath(t *testing.T) {
 
 	svc := NewAuthService(config.Config{}, nil, nil, chatSessionRegistry, nil, nil, nil, nil, nil, wire.DefaultRateLimitClasses(), nil, slog.Default())
 
-	have, err := svc.RegisterChatSession(context.Background(), serverCookie)
+	have, err := svc.RegisterChatSession(context.Background(), serverCookie, nil)
 	assert.NoError(t, err)
 	assert.Equal(t, instance, have)
 }
@@ -1939,7 +1939,7 @@ func TestAuthService_RegisterBOSSession(t *testing.T) {
 			sessionRegistry := newMockSessionRegistry(t)
 			for _, params := range tc.mockParams.addSessionParams {
 				sessionRegistry.EXPECT().
-					AddSession(mock.Anything, params.screenName, params.doMultiSess).
+					AddSession(mock.Anything, params.screenName, params.doMultiSess, mock.Anything, mock.Anything).
 					Return(params.result, params.err)
 			}
 			userManager := newMockUserManager(t)
@@ -1963,7 +1963,7 @@ func TestAuthService_RegisterBOSSession(t *testing.T) {
 
 			svc := NewAuthService(config.Config{}, sessionRegistry, nil, nil, userManager, nil, nil, accountManager, bartItemManager, wire.DefaultRateLimitClasses(), nil, slog.Default())
 
-			have, err := svc.RegisterBOSSession(context.Background(), tc.cookie)
+			have, err := svc.RegisterBOSSession(context.Background(), tc.cookie, nil)
 			assert.NoError(t, err)
 
 			if tc.wantSess != nil {
@@ -2107,11 +2107,11 @@ func TestAuthService_SignoutChat(t *testing.T) {
 			sessionManager := newMockChatSessionRegistry(t)
 			for _, params := range tt.mockParams.removeSessionParams {
 				sessionManager.EXPECT().
-					RemoveSession(matchSession(params.screenName))
+					RemoveSession(matchUserSession(params.screenName))
 			}
 
 			svc := NewAuthService(config.Config{}, nil, nil, sessionManager, nil, nil, chatMessageRelayer, nil, nil, wire.DefaultRateLimitClasses(), nil, slog.Default())
-			svc.SignoutChat(context.Background(), tt.instance)
+			svc.SignoutChat(context.Background(), tt.instance.Session())
 		})
 	}
 }
@@ -2153,11 +2153,11 @@ func TestAuthService_Signout(t *testing.T) {
 		t.Run(tt.name, func(t *testing.T) {
 			sessionManager := newMockSessionRegistry(t)
 			for _, params := range tt.mockParams.removeSessionParams {
-				sessionManager.EXPECT().RemoveSession(matchSession(params.screenName))
+				sessionManager.EXPECT().RemoveSession(matchUserSession(params.screenName))
 			}
 			svc := NewAuthService(config.Config{}, sessionManager, nil, nil, nil, nil, nil, nil, nil, wire.DefaultRateLimitClasses(), nil, slog.Default())
 
-			svc.Signout(context.Background(), tt.instance)
+			svc.Signout(context.Background(), tt.instance.Session())
 		})
 	}
 }

foodgroup/helpers_test.go

@@ -953,6 +953,13 @@ func matchSession(mustMatch state.IdentScreenName) interface{} {
 	})
 }
 
+// matchUserSession matches a mock call for *state.Session by ident screen name.
+func matchUserSession(mustMatch state.IdentScreenName) interface{} {
+	return mock.MatchedBy(func(s *state.Session) bool {
+		return mustMatch == s.IdentScreenName()
+	})
+}
+
 // matchContext matches any instance of Context interface.
 func matchContext() interface{} {
 	return mock.MatchedBy(func(ctx any) bool {

foodgroup/mock_chat_session_registry_test.go

@@ -39,8 +39,16 @@ func (_m *mockChatSessionRegistry) EXPECT() *mockChatSessionRegistry_Expecter {
 }
 
 // AddSession provides a mock function for the type mockChatSessionRegistry
-func (_mock *mockChatSessionRegistry) AddSession(ctx context.Context, chatCookie string, screenName state.DisplayScreenName) (*state.SessionInstance, error) {
-	ret := _mock.Called(ctx, chatCookie, screenName)
+func (_mock *mockChatSessionRegistry) AddSession(ctx context.Context, chatCookie string, screenName state.DisplayScreenName, cfg ...func(sess *state.Session)) (*state.SessionInstance, error) {
+	// func(sess *state.Session)
+	_va := make([]interface{}, len(cfg))
+	for _i := range cfg {
+		_va[_i] = cfg[_i]
+	}
+	var _ca []interface{}
+	_ca = append(_ca, ctx, chatCookie, screenName)
+	_ca = append(_ca, _va...)
+	ret := _mock.Called(_ca...)
 
 	if len(ret) == 0 {
 		panic("no return value specified for AddSession")
@@ -48,18 +56,18 @@ func (_mock *mockChatSessionRegistry) AddSession(ctx context.Context, chatCookie
 
 	var r0 *state.SessionInstance
 	var r1 error
-	if returnFunc, ok := ret.Get(0).(func(context.Context, string, state.DisplayScreenName) (*state.SessionInstance, error)); ok {
-		return returnFunc(ctx, chatCookie, screenName)
+	if returnFunc, ok := ret.Get(0).(func(context.Context, string, state.DisplayScreenName, ...func(sess *state.Session)) (*state.SessionInstance, error)); ok {
+		return returnFunc(ctx, chatCookie, screenName, cfg...)
 	}
-	if returnFunc, ok := ret.Get(0).(func(context.Context, string, state.DisplayScreenName) *state.SessionInstance); ok {
-		r0 = returnFunc(ctx, chatCookie, screenName)
+	if returnFunc, ok := ret.Get(0).(func(context.Context, string, state.DisplayScreenName, ...func(sess *state.Session)) *state.SessionInstance); ok {
+		r0 = returnFunc(ctx, chatCookie, screenName, cfg...)
 	} else {
 		if ret.Get(0) != nil {
 			r0 = ret.Get(0).(*state.SessionInstance)
 		}
 	}
-	if returnFunc, ok := ret.Get(1).(func(context.Context, string, state.DisplayScreenName) error); ok {
-		r1 = returnFunc(ctx, chatCookie, screenName)
+	if returnFunc, ok := ret.Get(1).(func(context.Context, string, state.DisplayScreenName, ...func(sess *state.Session)) error); ok {
+		r1 = returnFunc(ctx, chatCookie, screenName, cfg...)
 	} else {
 		r1 = ret.Error(1)
 	}
@@ -75,11 +83,13 @@ type mockChatSessionRegistry_AddSession_Call struct {
 //   - ctx context.Context
 //   - chatCookie string
 //   - screenName state.DisplayScreenName
-func (_e *mockChatSessionRegistry_Expecter) AddSession(ctx interface{}, chatCookie interface{}, screenName interface{}) *mockChatSessionRegistry_AddSession_Call {
-	return &mockChatSessionRegistry_AddSession_Call{Call: _e.mock.On("AddSession", ctx, chatCookie, screenName)}
+//   - cfg ...func(sess *state.Session)
+func (_e *mockChatSessionRegistry_Expecter) AddSession(ctx interface{}, chatCookie interface{}, screenName interface{}, cfg ...interface{}) *mockChatSessionRegistry_AddSession_Call {
+	return &mockChatSessionRegistry_AddSession_Call{Call: _e.mock.On("AddSession",
+		append([]interface{}{ctx, chatCookie, screenName}, cfg...)...)}
 }
 
-func (_c *mockChatSessionRegistry_AddSession_Call) Run(run func(ctx context.Context, chatCookie string, screenName state.DisplayScreenName)) *mockChatSessionRegistry_AddSession_Call {
+func (_c *mockChatSessionRegistry_AddSession_Call) Run(run func(ctx context.Context, chatCookie string, screenName state.DisplayScreenName, cfg ...func(sess *state.Session))) *mockChatSessionRegistry_AddSession_Call {
 	_c.Call.Run(func(args mock.Arguments) {
 		var arg0 context.Context
 		if args[0] != nil {
@@ -93,10 +103,19 @@ func (_c *mockChatSessionRegistry_AddSession_Call) Run(run func(ctx context.Cont
 		if args[2] != nil {
 			arg2 = args[2].(state.DisplayScreenName)
 		}
+		var arg3 []func(sess *state.Session)
+		variadicArgs := make([]func(sess *state.Session), len(args)-3)
+		for i, a := range args[3:] {
+			if a != nil {
+				variadicArgs[i] = a.(func(sess *state.Session))
+			}
+		}
+		arg3 = variadicArgs
 		run(
 			arg0,
 			arg1,
 			arg2,
+			arg3...,
 		)
 	})
 	return _c
@@ -107,14 +126,14 @@ func (_c *mockChatSessionRegistry_AddSession_Call) Return(sessionInstance *state
 	return _c
 }
 
-func (_c *mockChatSessionRegistry_AddSession_Call) RunAndReturn(run func(ctx context.Context, chatCookie string, screenName state.DisplayScreenName) (*state.SessionInstance, error)) *mockChatSessionRegistry_AddSession_Call {
+func (_c *mockChatSessionRegistry_AddSession_Call) RunAndReturn(run func(ctx context.Context, chatCookie string, screenName state.DisplayScreenName, cfg ...func(sess *state.Session)) (*state.SessionInstance, error)) *mockChatSessionRegistry_AddSession_Call {
 	_c.Call.Return(run)
 	return _c
 }
 
 // RemoveSession provides a mock function for the type mockChatSessionRegistry
-func (_mock *mockChatSessionRegistry) RemoveSession(instance *state.SessionInstance) {
-	_mock.Called(instance)
+func (_mock *mockChatSessionRegistry) RemoveSession(sess *state.Session) {
+	_mock.Called(sess)
 	return
 }
 
@@ -124,16 +143,16 @@ type mockChatSessionRegistry_RemoveSession_Call struct {
 }
 
 // RemoveSession is a helper method to define mock.On call
-//   - instance *state.SessionInstance
-func (_e *mockChatSessionRegistry_Expecter) RemoveSession(instance interface{}) *mockChatSessionRegistry_RemoveSession_Call {
-	return &mockChatSessionRegistry_RemoveSession_Call{Call: _e.mock.On("RemoveSession", instance)}
+//   - sess *state.Session
+func (_e *mockChatSessionRegistry_Expecter) RemoveSession(sess interface{}) *mockChatSessionRegistry_RemoveSession_Call {
+	return &mockChatSessionRegistry_RemoveSession_Call{Call: _e.mock.On("RemoveSession", sess)}
 }
 
-func (_c *mockChatSessionRegistry_RemoveSession_Call) Run(run func(instance *state.SessionInstance)) *mockChatSessionRegistry_RemoveSession_Call {
+func (_c *mockChatSessionRegistry_RemoveSession_Call) Run(run func(sess *state.Session)) *mockChatSessionRegistry_RemoveSession_Call {
 	_c.Call.Run(func(args mock.Arguments) {
-		var arg0 *state.SessionInstance
+		var arg0 *state.Session
 		if args[0] != nil {
-			arg0 = args[0].(*state.SessionInstance)
+			arg0 = args[0].(*state.Session)
 		}
 		run(
 			arg0,
@@ -147,7 +166,7 @@ func (_c *mockChatSessionRegistry_RemoveSession_Call) Return() *mockChatSessionR
 	return _c
 }
 
-func (_c *mockChatSessionRegistry_RemoveSession_Call) RunAndReturn(run func(instance *state.SessionInstance)) *mockChatSessionRegistry_RemoveSession_Call {
+func (_c *mockChatSessionRegistry_RemoveSession_Call) RunAndReturn(run func(sess *state.Session)) *mockChatSessionRegistry_RemoveSession_Call {
 	_c.Run(run)
 	return _c
 }

foodgroup/mock_session_registry_test.go

@@ -39,8 +39,16 @@ func (_m *mockSessionRegistry) EXPECT() *mockSessionRegistry_Expecter {
 }
 
 // AddSession provides a mock function for the type mockSessionRegistry
-func (_mock *mockSessionRegistry) AddSession(ctx context.Context, screenName state.DisplayScreenName, doMultiSess bool) (*state.SessionInstance, error) {
-	ret := _mock.Called(ctx, screenName, doMultiSess)
+func (_mock *mockSessionRegistry) AddSession(ctx context.Context, screenName state.DisplayScreenName, doMultiSess bool, cfg ...func(sess *state.Session)) (*state.SessionInstance, error) {
+	// func(sess *state.Session)
+	_va := make([]interface{}, len(cfg))
+	for _i := range cfg {
+		_va[_i] = cfg[_i]
+	}
+	var _ca []interface{}
+	_ca = append(_ca, ctx, screenName, doMultiSess)
+	_ca = append(_ca, _va...)
+	ret := _mock.Called(_ca...)
 
 	if len(ret) == 0 {
 		panic("no return value specified for AddSession")
@@ -48,18 +56,18 @@ func (_mock *mockSessionRegistry) AddSession(ctx context.Context, screenName sta
 
 	var r0 *state.SessionInstance
 	var r1 error
-	if returnFunc, ok := ret.Get(0).(func(context.Context, state.DisplayScreenName, bool) (*state.SessionInstance, error)); ok {
-		return returnFunc(ctx, screenName, doMultiSess)
+	if returnFunc, ok := ret.Get(0).(func(context.Context, state.DisplayScreenName, bool, ...func(sess *state.Session)) (*state.SessionInstance, error)); ok {
+		return returnFunc(ctx, screenName, doMultiSess, cfg...)
 	}
-	if returnFunc, ok := ret.Get(0).(func(context.Context, state.DisplayScreenName, bool) *state.SessionInstance); ok {
-		r0 = returnFunc(ctx, screenName, doMultiSess)
+	if returnFunc, ok := ret.Get(0).(func(context.Context, state.DisplayScreenName, bool, ...func(sess *state.Session)) *state.SessionInstance); ok {
+		r0 = returnFunc(ctx, screenName, doMultiSess, cfg...)
 	} else {
 		if ret.Get(0) != nil {
 			r0 = ret.Get(0).(*state.SessionInstance)
 		}
 	}
-	if returnFunc, ok := ret.Get(1).(func(context.Context, state.DisplayScreenName, bool) error); ok {
-		r1 = returnFunc(ctx, screenName, doMultiSess)
+	if returnFunc, ok := ret.Get(1).(func(context.Context, state.DisplayScreenName, bool, ...func(sess *state.Session)) error); ok {
+		r1 = returnFunc(ctx, screenName, doMultiSess, cfg...)
 	} else {
 		r1 = ret.Error(1)
 	}
@@ -75,11 +83,13 @@ type mockSessionRegistry_AddSession_Call struct {
 //   - ctx context.Context
 //   - screenName state.DisplayScreenName
 //   - doMultiSess bool
-func (_e *mockSessionRegistry_Expecter) AddSession(ctx interface{}, screenName interface{}, doMultiSess interface{}) *mockSessionRegistry_AddSession_Call {
-	return &mockSessionRegistry_AddSession_Call{Call: _e.mock.On("AddSession", ctx, screenName, doMultiSess)}
+//   - cfg ...func(sess *state.Session)
+func (_e *mockSessionRegistry_Expecter) AddSession(ctx interface{}, screenName interface{}, doMultiSess interface{}, cfg ...interface{}) *mockSessionRegistry_AddSession_Call {
+	return &mockSessionRegistry_AddSession_Call{Call: _e.mock.On("AddSession",
+		append([]interface{}{ctx, screenName, doMultiSess}, cfg...)...)}
 }
 
-func (_c *mockSessionRegistry_AddSession_Call) Run(run func(ctx context.Context, screenName state.DisplayScreenName, doMultiSess bool)) *mockSessionRegistry_AddSession_Call {
+func (_c *mockSessionRegistry_AddSession_Call) Run(run func(ctx context.Context, screenName state.DisplayScreenName, doMultiSess bool, cfg ...func(sess *state.Session))) *mockSessionRegistry_AddSession_Call {
 	_c.Call.Run(func(args mock.Arguments) {
 		var arg0 context.Context
 		if args[0] != nil {
@@ -93,10 +103,19 @@ func (_c *mockSessionRegistry_AddSession_Call) Run(run func(ctx context.Context,
 		if args[2] != nil {
 			arg2 = args[2].(bool)
 		}
+		var arg3 []func(sess *state.Session)
+		variadicArgs := make([]func(sess *state.Session), len(args)-3)
+		for i, a := range args[3:] {
+			if a != nil {
+				variadicArgs[i] = a.(func(sess *state.Session))
+			}
+		}
+		arg3 = variadicArgs
 		run(
 			arg0,
 			arg1,
 			arg2,
+			arg3...,
 		)
 	})
 	return _c
@@ -107,14 +126,14 @@ func (_c *mockSessionRegistry_AddSession_Call) Return(sessionInstance *state.Ses
 	return _c
 }
 
-func (_c *mockSessionRegistry_AddSession_Call) RunAndReturn(run func(ctx context.Context, screenName state.DisplayScreenName, doMultiSess bool) (*state.SessionInstance, error)) *mockSessionRegistry_AddSession_Call {
+func (_c *mockSessionRegistry_AddSession_Call) RunAndReturn(run func(ctx context.Context, screenName state.DisplayScreenName, doMultiSess bool, cfg ...func(sess *state.Session)) (*state.SessionInstance, error)) *mockSessionRegistry_AddSession_Call {
 	_c.Call.Return(run)
 	return _c
 }
 
 // RemoveSession provides a mock function for the type mockSessionRegistry
-func (_mock *mockSessionRegistry) RemoveSession(instance *state.SessionInstance) {
-	_mock.Called(instance)
+func (_mock *mockSessionRegistry) RemoveSession(session *state.Session) {
+	_mock.Called(session)
 	return
 }
 
@@ -124,16 +143,16 @@ type mockSessionRegistry_RemoveSession_Call struct {
 }
 
 // RemoveSession is a helper method to define mock.On call
-//   - instance *state.SessionInstance
-func (_e *mockSessionRegistry_Expecter) RemoveSession(instance interface{}) *mockSessionRegistry_RemoveSession_Call {
-	return &mockSessionRegistry_RemoveSession_Call{Call: _e.mock.On("RemoveSession", instance)}
+//   - session *state.Session
+func (_e *mockSessionRegistry_Expecter) RemoveSession(session interface{}) *mockSessionRegistry_RemoveSession_Call {
+	return &mockSessionRegistry_RemoveSession_Call{Call: _e.mock.On("RemoveSession", session)}
 }
 
-func (_c *mockSessionRegistry_RemoveSession_Call) Run(run func(instance *state.SessionInstance)) *mockSessionRegistry_RemoveSession_Call {
+func (_c *mockSessionRegistry_RemoveSession_Call) Run(run func(session *state.Session)) *mockSessionRegistry_RemoveSession_Call {
 	_c.Call.Run(func(args mock.Arguments) {
-		var arg0 *state.SessionInstance
+		var arg0 *state.Session
 		if args[0] != nil {
-			arg0 = args[0].(*state.SessionInstance)
+			arg0 = args[0].(*state.Session)
 		}
 		run(
 			arg0,
@@ -147,7 +166,7 @@ func (_c *mockSessionRegistry_RemoveSession_Call) Return() *mockSessionRegistry_
 	return _c
 }
 
-func (_c *mockSessionRegistry_RemoveSession_Call) RunAndReturn(run func(instance *state.SessionInstance)) *mockSessionRegistry_RemoveSession_Call {
+func (_c *mockSessionRegistry_RemoveSession_Call) RunAndReturn(run func(session *state.Session)) *mockSessionRegistry_RemoveSession_Call {
 	_c.Run(run)
 	return _c
 }

foodgroup/types.go

@@ -170,10 +170,10 @@ type ChatSessionRegistry interface {
 	// param identifies the chat room to which screenName is added. It returns
 	// the newly created session instance registered in the chat session
 	// manager.
-	AddSession(ctx context.Context, chatCookie string, screenName state.DisplayScreenName) (*state.SessionInstance, error)
+	AddSession(ctx context.Context, chatCookie string, screenName state.DisplayScreenName, cfg ...func(sess *state.Session)) (*state.SessionInstance, error)
 
 	// RemoveSession removes a session from the chat session manager.
-	RemoveSession(instance *state.SessionInstance)
+	RemoveSession(sess *state.Session)
 }
 
 // ClientSideBuddyListManager defines operations for managing a user's buddy list,
@@ -364,11 +364,11 @@ type SessionRegistry interface {
 	// When multiple concurrent calls are made for the same screen name, only one will succeed;
 	// the others will return an error once the context is done.
 	// If doMultiSess is true, allows multiple sessions for the same screen name.
-	AddSession(ctx context.Context, screenName state.DisplayScreenName, doMultiSess bool) (*state.SessionInstance, error)
+	AddSession(ctx context.Context, screenName state.DisplayScreenName, doMultiSess bool, cfg ...func(sess *state.Session)) (*state.SessionInstance, error)
 
 	// RemoveSession removes the given session from the registry, allowing future sessions
 	// to be created for the same screen name.
-	RemoveSession(instance *state.SessionInstance)
+	RemoveSession(session *state.Session)
 }
 
 // SessionRetriever defines a method for retrieving an active session

server/oscar/mock_auth_test.go

@@ -389,8 +389,8 @@ func (_c *mockAuthService_KerberosLogin_Call) RunAndReturn(run func(ctx context.
 }
 
 // RegisterBOSSession provides a mock function for the type mockAuthService
-func (_mock *mockAuthService) RegisterBOSSession(ctx context.Context, authCookie state.ServerCookie) (*state.SessionInstance, error) {
-	ret := _mock.Called(ctx, authCookie)
+func (_mock *mockAuthService) RegisterBOSSession(ctx context.Context, authCookie state.ServerCookie, sessCfg func(sess *state.Session)) (*state.SessionInstance, error) {
+	ret := _mock.Called(ctx, authCookie, sessCfg)
 
 	if len(ret) == 0 {
 		panic("no return value specified for RegisterBOSSession")
@@ -398,18 +398,18 @@ func (_mock *mockAuthService) RegisterBOSSession(ctx context.Context, authCookie
 
 	var r0 *state.SessionInstance
 	var r1 error
-	if returnFunc, ok := ret.Get(0).(func(context.Context, state.ServerCookie) (*state.SessionInstance, error)); ok {
-		return returnFunc(ctx, authCookie)
+	if returnFunc, ok := ret.Get(0).(func(context.Context, state.ServerCookie, func(sess *state.Session)) (*state.SessionInstance, error)); ok {
+		return returnFunc(ctx, authCookie, sessCfg)
 	}
-	if returnFunc, ok := ret.Get(0).(func(context.Context, state.ServerCookie) *state.SessionInstance); ok {
-		r0 = returnFunc(ctx, authCookie)
+	if returnFunc, ok := ret.Get(0).(func(context.Context, state.ServerCookie, func(sess *state.Session)) *state.SessionInstance); ok {
+		r0 = returnFunc(ctx, authCookie, sessCfg)
 	} else {
 		if ret.Get(0) != nil {
 			r0 = ret.Get(0).(*state.SessionInstance)
 		}
 	}
-	if returnFunc, ok := ret.Get(1).(func(context.Context, state.ServerCookie) error); ok {
-		r1 = returnFunc(ctx, authCookie)
+	if returnFunc, ok := ret.Get(1).(func(context.Context, state.ServerCookie, func(sess *state.Session)) error); ok {
+		r1 = returnFunc(ctx, authCookie, sessCfg)
 	} else {
 		r1 = ret.Error(1)
 	}
@@ -424,11 +424,12 @@ type mockAuthService_RegisterBOSSession_Call struct {
 // RegisterBOSSession is a helper method to define mock.On call
 //   - ctx context.Context
 //   - authCookie state.ServerCookie
-func (_e *mockAuthService_Expecter) RegisterBOSSession(ctx interface{}, authCookie interface{}) *mockAuthService_RegisterBOSSession_Call {
-	return &mockAuthService_RegisterBOSSession_Call{Call: _e.mock.On("RegisterBOSSession", ctx, authCookie)}
+//   - conf func(sess *state.Session)
+func (_e *mockAuthService_Expecter) RegisterBOSSession(ctx interface{}, authCookie interface{}, conf interface{}) *mockAuthService_RegisterBOSSession_Call {
+	return &mockAuthService_RegisterBOSSession_Call{Call: _e.mock.On("RegisterBOSSession", ctx, authCookie, conf)}
 }
 
-func (_c *mockAuthService_RegisterBOSSession_Call) Run(run func(ctx context.Context, authCookie state.ServerCookie)) *mockAuthService_RegisterBOSSession_Call {
+func (_c *mockAuthService_RegisterBOSSession_Call) Run(run func(ctx context.Context, authCookie state.ServerCookie, conf func(sess *state.Session))) *mockAuthService_RegisterBOSSession_Call {
 	_c.Call.Run(func(args mock.Arguments) {
 		var arg0 context.Context
 		if args[0] != nil {
@@ -438,9 +439,14 @@ func (_c *mockAuthService_RegisterBOSSession_Call) Run(run func(ctx context.Cont
 		if args[1] != nil {
 			arg1 = args[1].(state.ServerCookie)
 		}
+		var arg2 func(sess *state.Session)
+		if args[2] != nil {
+			arg2 = args[2].(func(sess *state.Session))
+		}
 		run(
 			arg0,
 			arg1,
+			arg2,
 		)
 	})
 	return _c
@@ -451,14 +457,14 @@ func (_c *mockAuthService_RegisterBOSSession_Call) Return(sessionInstance *state
 	return _c
 }
 
-func (_c *mockAuthService_RegisterBOSSession_Call) RunAndReturn(run func(ctx context.Context, authCookie state.ServerCookie) (*state.SessionInstance, error)) *mockAuthService_RegisterBOSSession_Call {
+func (_c *mockAuthService_RegisterBOSSession_Call) RunAndReturn(run func(ctx context.Context, authCookie state.ServerCookie, conf func(sess *state.Session)) (*state.SessionInstance, error)) *mockAuthService_RegisterBOSSession_Call {
 	_c.Call.Return(run)
 	return _c
 }
 
 // RegisterChatSession provides a mock function for the type mockAuthService
-func (_mock *mockAuthService) RegisterChatSession(ctx context.Context, authCookie state.ServerCookie) (*state.SessionInstance, error) {
-	ret := _mock.Called(ctx, authCookie)
+func (_mock *mockAuthService) RegisterChatSession(ctx context.Context, authCookie state.ServerCookie, sessCfg func(sess *state.Session)) (*state.SessionInstance, error) {
+	ret := _mock.Called(ctx, authCookie, sessCfg)
 
 	if len(ret) == 0 {
 		panic("no return value specified for RegisterChatSession")
@@ -466,18 +472,18 @@ func (_mock *mockAuthService) RegisterChatSession(ctx context.Context, authCooki
 
 	var r0 *state.SessionInstance
 	var r1 error
-	if returnFunc, ok := ret.Get(0).(func(context.Context, state.ServerCookie) (*state.SessionInstance, error)); ok {
-		return returnFunc(ctx, authCookie)
+	if returnFunc, ok := ret.Get(0).(func(context.Context, state.ServerCookie, func(sess *state.Session)) (*state.SessionInstance, error)); ok {
+		return returnFunc(ctx, authCookie, sessCfg)
 	}
-	if returnFunc, ok := ret.Get(0).(func(context.Context, state.ServerCookie) *state.SessionInstance); ok {
-		r0 = returnFunc(ctx, authCookie)
+	if returnFunc, ok := ret.Get(0).(func(context.Context, state.ServerCookie, func(sess *state.Session)) *state.SessionInstance); ok {
+		r0 = returnFunc(ctx, authCookie, sessCfg)
 	} else {
 		if ret.Get(0) != nil {
 			r0 = ret.Get(0).(*state.SessionInstance)
 		}
 	}
-	if returnFunc, ok := ret.Get(1).(func(context.Context, state.ServerCookie) error); ok {
-		r1 = returnFunc(ctx, authCookie)
+	if returnFunc, ok := ret.Get(1).(func(context.Context, state.ServerCookie, func(sess *state.Session)) error); ok {
+		r1 = returnFunc(ctx, authCookie, sessCfg)
 	} else {
 		r1 = ret.Error(1)
 	}
@@ -492,11 +498,12 @@ type mockAuthService_RegisterChatSession_Call struct {
 // RegisterChatSession is a helper method to define mock.On call
 //   - ctx context.Context
 //   - authCookie state.ServerCookie
-func (_e *mockAuthService_Expecter) RegisterChatSession(ctx interface{}, authCookie interface{}) *mockAuthService_RegisterChatSession_Call {
-	return &mockAuthService_RegisterChatSession_Call{Call: _e.mock.On("RegisterChatSession", ctx, authCookie)}
+//   - cfg func(sess *state.Session)
+func (_e *mockAuthService_Expecter) RegisterChatSession(ctx interface{}, authCookie interface{}, cfg interface{}) *mockAuthService_RegisterChatSession_Call {
+	return &mockAuthService_RegisterChatSession_Call{Call: _e.mock.On("RegisterChatSession", ctx, authCookie, cfg)}
 }
 
-func (_c *mockAuthService_RegisterChatSession_Call) Run(run func(ctx context.Context, authCookie state.ServerCookie)) *mockAuthService_RegisterChatSession_Call {
+func (_c *mockAuthService_RegisterChatSession_Call) Run(run func(ctx context.Context, authCookie state.ServerCookie, cfg func(sess *state.Session))) *mockAuthService_RegisterChatSession_Call {
 	_c.Call.Run(func(args mock.Arguments) {
 		var arg0 context.Context
 		if args[0] != nil {
@@ -506,9 +513,14 @@ func (_c *mockAuthService_RegisterChatSession_Call) Run(run func(ctx context.Con
 		if args[1] != nil {
 			arg1 = args[1].(state.ServerCookie)
 		}
+		var arg2 func(sess *state.Session)
+		if args[2] != nil {
+			arg2 = args[2].(func(sess *state.Session))
+		}
 		run(
 			arg0,
 			arg1,
+			arg2,
 		)
 	})
 	return _c
@@ -519,7 +531,7 @@ func (_c *mockAuthService_RegisterChatSession_Call) Return(sessionInstance *stat
 	return _c
 }
 
-func (_c *mockAuthService_RegisterChatSession_Call) RunAndReturn(run func(ctx context.Context, authCookie state.ServerCookie) (*state.SessionInstance, error)) *mockAuthService_RegisterChatSession_Call {
+func (_c *mockAuthService_RegisterChatSession_Call) RunAndReturn(run func(ctx context.Context, authCookie state.ServerCookie, cfg func(sess *state.Session)) (*state.SessionInstance, error)) *mockAuthService_RegisterChatSession_Call {
 	_c.Call.Return(run)
 	return _c
 }
@@ -593,8 +605,8 @@ func (_c *mockAuthService_RetrieveBOSSession_Call) RunAndReturn(run func(ctx con
 }
 
 // Signout provides a mock function for the type mockAuthService
-func (_mock *mockAuthService) Signout(ctx context.Context, instance *state.SessionInstance) {
-	_mock.Called(ctx, instance)
+func (_mock *mockAuthService) Signout(ctx context.Context, session *state.Session) {
+	_mock.Called(ctx, session)
 	return
 }
 
@@ -605,20 +617,20 @@ type mockAuthService_Signout_Call struct {
 
 // Signout is a helper method to define mock.On call
 //   - ctx context.Context
-//   - instance *state.SessionInstance
-func (_e *mockAuthService_Expecter) Signout(ctx interface{}, instance interface{}) *mockAuthService_Signout_Call {
-	return &mockAuthService_Signout_Call{Call: _e.mock.On("Signout", ctx, instance)}
+//   - session *state.Session
+func (_e *mockAuthService_Expecter) Signout(ctx interface{}, session interface{}) *mockAuthService_Signout_Call {
+	return &mockAuthService_Signout_Call{Call: _e.mock.On("Signout", ctx, session)}
 }
 
-func (_c *mockAuthService_Signout_Call) Run(run func(ctx context.Context, instance *state.SessionInstance)) *mockAuthService_Signout_Call {
+func (_c *mockAuthService_Signout_Call) Run(run func(ctx context.Context, session *state.Session)) *mockAuthService_Signout_Call {
 	_c.Call.Run(func(args mock.Arguments) {
 		var arg0 context.Context
 		if args[0] != nil {
 			arg0 = args[0].(context.Context)
 		}
-		var arg1 *state.SessionInstance
+		var arg1 *state.Session
 		if args[1] != nil {
-			arg1 = args[1].(*state.SessionInstance)
+			arg1 = args[1].(*state.Session)
 		}
 		run(
 			arg0,
@@ -633,13 +645,13 @@ func (_c *mockAuthService_Signout_Call) Return() *mockAuthService_Signout_Call {
 	return _c
 }
 
-func (_c *mockAuthService_Signout_Call) RunAndReturn(run func(ctx context.Context, instance *state.SessionInstance)) *mockAuthService_Signout_Call {
+func (_c *mockAuthService_Signout_Call) RunAndReturn(run func(ctx context.Context, session *state.Session)) *mockAuthService_Signout_Call {
 	_c.Run(run)
 	return _c
 }
 
 // SignoutChat provides a mock function for the type mockAuthService
-func (_mock *mockAuthService) SignoutChat(ctx context.Context, instance *state.SessionInstance) {
+func (_mock *mockAuthService) SignoutChat(ctx context.Context, instance *state.Session) {
 	_mock.Called(ctx, instance)
 	return
 }
@@ -651,20 +663,20 @@ type mockAuthService_SignoutChat_Call struct {
 
 // SignoutChat is a helper method to define mock.On call
 //   - ctx context.Context
-//   - instance *state.SessionInstance
+//   - instance *state.Session
 func (_e *mockAuthService_Expecter) SignoutChat(ctx interface{}, instance interface{}) *mockAuthService_SignoutChat_Call {
 	return &mockAuthService_SignoutChat_Call{Call: _e.mock.On("SignoutChat", ctx, instance)}
 }
 
-func (_c *mockAuthService_SignoutChat_Call) Run(run func(ctx context.Context, instance *state.SessionInstance)) *mockAuthService_SignoutChat_Call {
+func (_c *mockAuthService_SignoutChat_Call) Run(run func(ctx context.Context, instance *state.Session)) *mockAuthService_SignoutChat_Call {
 	_c.Call.Run(func(args mock.Arguments) {
 		var arg0 context.Context
 		if args[0] != nil {
 			arg0 = args[0].(context.Context)
 		}
-		var arg1 *state.SessionInstance
+		var arg1 *state.Session
 		if args[1] != nil {
-			arg1 = args[1].(*state.SessionInstance)
+			arg1 = args[1].(*state.Session)
 		}
 		run(
 			arg0,
@@ -679,7 +691,7 @@ func (_c *mockAuthService_SignoutChat_Call) Return() *mockAuthService_SignoutCha
 	return _c
 }
 
-func (_c *mockAuthService_SignoutChat_Call) RunAndReturn(run func(ctx context.Context, instance *state.SessionInstance)) *mockAuthService_SignoutChat_Call {
+func (_c *mockAuthService_SignoutChat_Call) RunAndReturn(run func(ctx context.Context, instance *state.Session)) *mockAuthService_SignoutChat_Call {
 	_c.Run(run)
 	return _c
 }

server/oscar/server.go

@@ -249,7 +249,33 @@ func (s oscarServer) connectToOSCARService(
 	var instance *state.SessionInstance
 	switch cookie.Service {
 	case wire.BOS:
-		instance, err = s.AuthService.RegisterBOSSession(ctx, cookie)
+
+		sessCfg := func(sess *state.Session) {
+			sess.OnSessionClose(func() {
+				if !shuttingDown(ctx) {
+					instances := sess.Instances()
+					if len(instances) > 0 {
+						if err := s.DepartureNotifier.BroadcastBuddyDeparted(ctx, instances[0]); err != nil {
+							s.Logger.ErrorContext(ctx, "error sending buddy departure notifications", "err", err.Error())
+						}
+					}
+				}
+
+				ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second)
+				defer cancel()
+
+				// buddy list must be cleared before session is closed, otherwise
+				// there will be a race condition that could cause the buddy list
+				// be prematurely deleted.
+				if err := s.BuddyListRegistry.UnregisterBuddyList(ctx, instance.IdentScreenName()); err != nil {
+					s.Logger.ErrorContext(ctx, "error removing buddy list entry", "err", err.Error())
+				}
+				s.ChatSessionManager.RemoveUserFromAllChats(instance.IdentScreenName())
+				s.AuthService.Signout(ctx, sess)
+			})
+		}
+
+		instance, err = s.AuthService.RegisterBOSSession(ctx, cookie, sessCfg)
 		if err != nil {
 			if errors.Is(err, state.ErrMaxConcurrentSessionsReached) {
 				s.Logger.Debug("session registration failed", "err", err.Error())
@@ -305,26 +331,6 @@ func (s oscarServer) connectToOSCARService(
 			}
 		})
 
-		instance.Session().OnSessionClose(func() {
-			if !shuttingDown(ctx) {
-				if err := s.DepartureNotifier.BroadcastBuddyDeparted(ctx, instance); err != nil {
-					s.Logger.ErrorContext(ctx, "error sending buddy departure notifications", "err", err.Error())
-				}
-			}
-
-			ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second)
-			defer cancel()
-
-			// buddy list must be cleared before session is closed, otherwise
-			// there will be a race condition that could cause the buddy list
-			// be prematurely deleted.
-			if err := s.BuddyListRegistry.UnregisterBuddyList(ctx, instance.IdentScreenName()); err != nil {
-				s.Logger.ErrorContext(ctx, "error removing buddy list entry", "err", err.Error())
-			}
-			s.ChatSessionManager.RemoveUserFromAllChats(instance.IdentScreenName())
-			s.AuthService.Signout(ctx, instance)
-		})
-
 		if remoteAddr, ok := ctx.Value("ip").(string); ok {
 			ip, err := netip.ParseAddrPort(remoteAddr)
 			if err != nil {
@@ -335,7 +341,14 @@ func (s oscarServer) connectToOSCARService(
 
 		go s.receiveSessMessages(ctx, instance, flapc)
 	case wire.Chat:
-		instance, err = s.AuthService.RegisterChatSession(ctx, cookie)
+		sessCfg := func(sess *state.Session) {
+			sess.OnSessionClose(func() {
+				ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second)
+				defer cancel()
+				s.SignoutChat(ctx, sess)
+			})
+		}
+		instance, err = s.AuthService.RegisterChatSession(ctx, cookie, sessCfg)
 		if err != nil {
 			return err
 		}
@@ -346,12 +359,6 @@ func (s oscarServer) connectToOSCARService(
 			instance.CloseInstance()
 		}()
 
-		instance.Session().OnSessionClose(func() {
-			ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second)
-			defer cancel()
-			s.SignoutChat(ctx, instance)
-		})
-
 		go s.receiveSessMessages(ctx, instance, flapc)
 	default:
 		instance, err = s.AuthService.RetrieveBOSSession(ctx, cookie)

server/oscar/server_test.go

@@ -371,12 +371,17 @@ func TestOscarServer_RouteConnection_BOS(t *testing.T) {
 
 	authService := newMockAuthService(t)
 	authService.EXPECT().
-		RegisterBOSSession(mock.Anything, state.ServerCookie{Service: wire.BOS}).
+		RegisterBOSSession(mock.Anything, state.ServerCookie{Service: wire.BOS}, mock.Anything).
+		Run(func(ctx context.Context, cookie state.ServerCookie, conf func(*state.Session)) {
+			if conf != nil {
+				conf(instance.Session())
+			}
+		}).
 		Return(instance, nil)
 	wg.Add(1)
 	authService.EXPECT().
-		Signout(mock.Anything, instance).
-		Run(func(ctx context.Context, s *state.SessionInstance) {
+		Signout(mock.Anything, instance.Session()).
+		Run(func(ctx context.Context, s *state.Session) {
 			defer wg.Done()
 		})
 
@@ -495,7 +500,7 @@ func TestOscarServer_RouteConnection_BOS_MultiSessionSignoff(t *testing.T) {
 
 	authService := newMockAuthService(t)
 	authService.EXPECT().
-		RegisterBOSSession(mock.Anything, state.ServerCookie{Service: wire.BOS}).
+		RegisterBOSSession(mock.Anything, state.ServerCookie{Service: wire.BOS}, mock.Anything).
 		Return(instance, nil)
 
 	authService.EXPECT().
@@ -598,7 +603,7 @@ func TestOscarServer_RouteConnection_BOS_MaxConcurrentSessionsReached(t *testing
 
 	authService := newMockAuthService(t)
 	authService.EXPECT().
-		RegisterBOSSession(mock.Anything, state.ServerCookie{Service: wire.BOS}).
+		RegisterBOSSession(mock.Anything, state.ServerCookie{Service: wire.BOS}, mock.Anything).
 		Return(nil, state.ErrMaxConcurrentSessionsReached)
 
 	authService.EXPECT().
@@ -667,12 +672,17 @@ func TestOscarServer_RouteConnection_Chat(t *testing.T) {
 
 	authService := newMockAuthService(t)
 	authService.EXPECT().
-		RegisterChatSession(mock.Anything, state.ServerCookie{Service: wire.Chat}).
+		RegisterChatSession(mock.Anything, state.ServerCookie{Service: wire.Chat}, mock.Anything).
+		Run(func(_ context.Context, _ state.ServerCookie, cfg func(*state.Session)) {
+			if cfg != nil {
+				cfg(instance.Session())
+			}
+		}).
 		Return(instance, nil)
 	wg.Add(1)
 	authService.EXPECT().
-		SignoutChat(mock.Anything, instance).
-		Run(func(ctx context.Context, s *state.SessionInstance) {
+		SignoutChat(mock.Anything, instance.Session()).
+		Run(func(ctx context.Context, s *state.Session) {
 			defer wg.Done()
 		})
 
@@ -913,14 +923,19 @@ func Test_oscarServer_receiveSessMessages_BOS_integration(t *testing.T) {
 		CrackCookie(mock.Anything).
 		Return(state.ServerCookie{Service: wire.BOS}, nil)
 	authService.EXPECT().
-		RegisterBOSSession(mock.Anything, state.ServerCookie{Service: wire.BOS}).
+		RegisterBOSSession(mock.Anything, state.ServerCookie{Service: wire.BOS}, mock.Anything).
+		Run(func(ctx context.Context, cookie state.ServerCookie, conf func(*state.Session)) {
+			if conf != nil {
+				conf(instance.Session())
+			}
+		}).
 		Return(instance, nil)
 
 	var signoutWG sync.WaitGroup
 	signoutWG.Add(1)
 	authService.EXPECT().
-		Signout(mock.Anything, instance).
-		Run(func(ctx context.Context, s *state.SessionInstance) { signoutWG.Done() })
+		Signout(mock.Anything, instance.Session()).
+		Run(func(ctx context.Context, s *state.Session) { signoutWG.Done() })
 
 	onlineNotifier := newMockOnlineNotifier(t)
 	onlineNotifier.EXPECT().
@@ -1059,14 +1074,19 @@ func Test_oscarServer_receiveSessMessages_Chat_integration(t *testing.T) {
 		CrackCookie(mock.Anything).
 		Return(state.ServerCookie{Service: wire.Chat}, nil)
 	authService.EXPECT().
-		RegisterChatSession(mock.Anything, state.ServerCookie{Service: wire.Chat}).
+		RegisterChatSession(mock.Anything, state.ServerCookie{Service: wire.Chat}, mock.Anything).
+		Run(func(_ context.Context, _ state.ServerCookie, cfg func(*state.Session)) {
+			if cfg != nil {
+				cfg(instance.Session())
+			}
+		}).
 		Return(instance, nil)
 
 	var signoutWG sync.WaitGroup
 	signoutWG.Add(1)
 	authService.EXPECT().
-		SignoutChat(mock.Anything, instance).
-		Run(func(ctx context.Context, s *state.SessionInstance) { signoutWG.Done() })
+		SignoutChat(mock.Anything, instance.Session()).
+		Run(func(ctx context.Context, s *state.Session) { signoutWG.Done() })
 
 	onlineNotifier := newMockOnlineNotifier(t)
 	onlineNotifier.EXPECT().

server/oscar/types.go

@@ -51,11 +51,11 @@ type AuthService interface {
 	CrackCookie(authCookie []byte) (state.ServerCookie, error)
 	FLAPLogin(ctx context.Context, inFrame wire.FLAPSignonFrame, advertisedHost string) (wire.TLVRestBlock, error)
 	KerberosLogin(ctx context.Context, inBody wire.SNAC_0x050C_0x0002_KerberosLoginRequest, advertisedHost string) (wire.SNACMessage, error)
-	RegisterBOSSession(ctx context.Context, authCookie state.ServerCookie) (*state.SessionInstance, error)
-	RegisterChatSession(ctx context.Context, authCookie state.ServerCookie) (*state.SessionInstance, error)
+	RegisterBOSSession(ctx context.Context, authCookie state.ServerCookie, sessCfg func(sess *state.Session)) (*state.SessionInstance, error)
+	RegisterChatSession(ctx context.Context, authCookie state.ServerCookie, sessCfg func(sess *state.Session)) (*state.SessionInstance, error)
 	RetrieveBOSSession(ctx context.Context, authCookie state.ServerCookie) (*state.SessionInstance, error)
-	Signout(ctx context.Context, instance *state.SessionInstance)
-	SignoutChat(ctx context.Context, instance *state.SessionInstance)
+	Signout(ctx context.Context, session *state.Session)
+	SignoutChat(ctx context.Context, instance *state.Session)
 }
 
 type AdminService interface {

server/toc/cmd_client.go

@@ -527,17 +527,18 @@ func (s OSCARProxy) ChatAccept(
 		return 0, s.runtimeErr(ctx, fmt.Errorf("AuthService.CrackCookie: %w", err))
 	}
 
-	chatSess, err := s.AuthService.RegisterChatSession(ctx, serverCookie)
+	sessCfg := func(sess *state.Session) {
+		sess.OnSessionClose(func() {
+			ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second)
+			defer cancel()
+			s.AuthService.SignoutChat(ctx, sess)
+		})
+	}
+	chatSess, err := s.AuthService.RegisterChatSession(ctx, serverCookie, sessCfg)
 	if err != nil {
 		return 0, s.runtimeErr(ctx, fmt.Errorf("AuthService.RegisterChatSession: %w", err))
 	}
 
-	chatSess.Session().OnSessionClose(func() {
-		ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second)
-		defer cancel()
-		s.AuthService.SignoutChat(ctx, chatSess)
-	})
-
 	if msg, isLimited := s.checkRateLimit(ctx, me, wire.OService, wire.OServiceClientOnline); isLimited {
 		return 0, msg
 	}
@@ -722,17 +723,18 @@ func (s OSCARProxy) ChatJoin(
 		return 0, s.runtimeErr(ctx, fmt.Errorf("AuthService.CrackCookie: %w", err))
 	}
 
-	chatSess, err := s.AuthService.RegisterChatSession(ctx, serverCookie)
+	sessCfg := func(sess *state.Session) {
+		sess.OnSessionClose(func() {
+			ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second)
+			defer cancel()
+			s.AuthService.SignoutChat(ctx, sess)
+		})
+	}
+	chatSess, err := s.AuthService.RegisterChatSession(ctx, serverCookie, sessCfg)
 	if err != nil {
 		return 0, s.runtimeErr(ctx, fmt.Errorf("AuthService.RegisterChatSession: %w", err))
 	}
 
-	chatSess.Session().OnSessionClose(func() {
-		ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second)
-		defer cancel()
-		s.AuthService.SignoutChat(ctx, chatSess)
-	})
-
 	if msg, isLimited := s.checkRateLimit(ctx, me, wire.OService, wire.OServiceClientOnline); isLimited {
 		return 0, msg
 	}
@@ -2349,7 +2351,33 @@ func (s OSCARProxy) Signon(ctx context.Context, args []byte, recalcWarning func(
 		return nil, s.runtimeErr(ctx, fmt.Errorf("AuthService.CrackCookie: %w", err))
 	}
 
-	instance, err := s.AuthService.RegisterBOSSession(ctx, serverCookie)
+	fnCfg := func(sess *state.Session) {
+		sess.OnSessionClose(func() {
+			if !shuttingDown(ctx) {
+				instances := sess.Instances()
+				if len(instances) > 0 {
+					if err := s.BuddyService.BroadcastBuddyDeparted(ctx, instances[0]); err != nil {
+						s.Logger.ErrorContext(ctx, "error sending buddy departure notifications", "err", err.Error())
+					}
+				}
+			}
+
+			ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second)
+			defer cancel()
+
+			// buddy list must be cleared before session is closed, otherwise
+			// there will be a race condition that could cause the buddy list
+			// be prematurely deleted.
+			if err := s.BuddyListRegistry.UnregisterBuddyList(ctx, sess.IdentScreenName()); err != nil {
+				s.Logger.ErrorContext(ctx, "error removing buddy list entry", "err", err.Error())
+			}
+			s.ChatSessionManager.RemoveUserFromAllChats(sess.IdentScreenName())
+			s.AuthService.Signout(ctx, sess)
+		})
+
+	}
+
+	instance, err := s.AuthService.RegisterBOSSession(ctx, serverCookie, fnCfg)
 	if err != nil {
 		return nil, s.runtimeErr(ctx, fmt.Errorf("AuthService.RegisterBOSSession: %w", err))
 	}
@@ -2388,26 +2416,6 @@ func (s OSCARProxy) Signon(ctx context.Context, args []byte, recalcWarning func(
 		}
 	})
 
-	instance.Session().OnSessionClose(func() {
-		if !shuttingDown(ctx) {
-			if err := s.BuddyService.BroadcastBuddyDeparted(ctx, instance); err != nil {
-				s.Logger.ErrorContext(ctx, "error sending buddy departure notifications", "err", err.Error())
-			}
-		}
-
-		ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second)
-		defer cancel()
-
-		// buddy list must be cleared before session is closed, otherwise
-		// there will be a race condition that could cause the buddy list
-		// be prematurely deleted.
-		if err := s.BuddyListRegistry.UnregisterBuddyList(ctx, instance.IdentScreenName()); err != nil {
-			s.Logger.ErrorContext(ctx, "error removing buddy list entry", "err", err.Error())
-		}
-		s.ChatSessionManager.RemoveUserFromAllChats(instance.IdentScreenName())
-		s.AuthService.Signout(ctx, instance)
-	})
-
 	// set chat capability so that... tk
 	instance.SetCaps([][16]byte{wire.CapChat})
 

server/toc/cmd_client_test.go

@@ -950,7 +950,7 @@ func TestOSCARProxy_RecvClientCmd_ChatAccept(t *testing.T) {
 			authSvc := newMockAuthService(t)
 			for _, params := range tc.mockParams.authParams.registerChatSessionParams {
 				authSvc.EXPECT().
-					RegisterChatSession(ctx, params.authCookie).
+					RegisterChatSession(ctx, params.authCookie, mock.Anything).
 					Return(params.instance, params.err)
 			}
 			for _, params := range tc.mockParams.authParams.crackCookieParams {
@@ -1482,7 +1482,7 @@ func TestOSCARProxy_RecvClientCmd_ChatJoin(t *testing.T) {
 			authSvc := newMockAuthService(t)
 			for _, params := range tc.mockParams.authParams.registerChatSessionParams {
 				authSvc.EXPECT().
-					RegisterChatSession(ctx, params.authCookie).
+					RegisterChatSession(ctx, params.authCookie, mock.Anything).
 					Return(params.instance, params.err)
 			}
 			for _, params := range tc.mockParams.authParams.crackCookieParams {
@@ -6915,7 +6915,7 @@ func TestOSCARProxy_Signon(t *testing.T) {
 			}
 			for _, params := range tc.mockParams.registerBOSSessionParams {
 				authSvc.EXPECT().
-					RegisterBOSSession(matchContext(), params.authCookie).
+					RegisterBOSSession(matchContext(), params.authCookie, mock.Anything).
 					Return(params.instance, params.err)
 			}
 			buddyRegistry := newMockBuddyListRegistry(t)

server/toc/mock_auth_service_test.go

@@ -317,8 +317,8 @@ func (_c *mockAuthService_FLAPLogin_Call) RunAndReturn(run func(ctx context.Cont
 }
 
 // RegisterBOSSession provides a mock function for the type mockAuthService
-func (_mock *mockAuthService) RegisterBOSSession(ctx context.Context, authCookie state.ServerCookie) (*state.SessionInstance, error) {
-	ret := _mock.Called(ctx, authCookie)
+func (_mock *mockAuthService) RegisterBOSSession(ctx context.Context, authCookie state.ServerCookie, cfg func(*state.Session)) (*state.SessionInstance, error) {
+	ret := _mock.Called(ctx, authCookie, cfg)
 
 	if len(ret) == 0 {
 		panic("no return value specified for RegisterBOSSession")
@@ -326,18 +326,18 @@ func (_mock *mockAuthService) RegisterBOSSession(ctx context.Context, authCookie
 
 	var r0 *state.SessionInstance
 	var r1 error
-	if returnFunc, ok := ret.Get(0).(func(context.Context, state.ServerCookie) (*state.SessionInstance, error)); ok {
-		return returnFunc(ctx, authCookie)
+	if returnFunc, ok := ret.Get(0).(func(context.Context, state.ServerCookie, func(*state.Session)) (*state.SessionInstance, error)); ok {
+		return returnFunc(ctx, authCookie, cfg)
 	}
-	if returnFunc, ok := ret.Get(0).(func(context.Context, state.ServerCookie) *state.SessionInstance); ok {
-		r0 = returnFunc(ctx, authCookie)
+	if returnFunc, ok := ret.Get(0).(func(context.Context, state.ServerCookie, func(*state.Session)) *state.SessionInstance); ok {
+		r0 = returnFunc(ctx, authCookie, cfg)
 	} else {
 		if ret.Get(0) != nil {
 			r0 = ret.Get(0).(*state.SessionInstance)
 		}
 	}
-	if returnFunc, ok := ret.Get(1).(func(context.Context, state.ServerCookie) error); ok {
-		r1 = returnFunc(ctx, authCookie)
+	if returnFunc, ok := ret.Get(1).(func(context.Context, state.ServerCookie, func(*state.Session)) error); ok {
+		r1 = returnFunc(ctx, authCookie, cfg)
 	} else {
 		r1 = ret.Error(1)
 	}
@@ -352,11 +352,12 @@ type mockAuthService_RegisterBOSSession_Call struct {
 // RegisterBOSSession is a helper method to define mock.On call
 //   - ctx context.Context
 //   - authCookie state.ServerCookie
-func (_e *mockAuthService_Expecter) RegisterBOSSession(ctx interface{}, authCookie interface{}) *mockAuthService_RegisterBOSSession_Call {
-	return &mockAuthService_RegisterBOSSession_Call{Call: _e.mock.On("RegisterBOSSession", ctx, authCookie)}
+//   - cfg func(*state.Session)
+func (_e *mockAuthService_Expecter) RegisterBOSSession(ctx interface{}, authCookie interface{}, cfg interface{}) *mockAuthService_RegisterBOSSession_Call {
+	return &mockAuthService_RegisterBOSSession_Call{Call: _e.mock.On("RegisterBOSSession", ctx, authCookie, cfg)}
 }
 
-func (_c *mockAuthService_RegisterBOSSession_Call) Run(run func(ctx context.Context, authCookie state.ServerCookie)) *mockAuthService_RegisterBOSSession_Call {
+func (_c *mockAuthService_RegisterBOSSession_Call) Run(run func(ctx context.Context, authCookie state.ServerCookie, cfg func(*state.Session))) *mockAuthService_RegisterBOSSession_Call {
 	_c.Call.Run(func(args mock.Arguments) {
 		var arg0 context.Context
 		if args[0] != nil {
@@ -366,9 +367,14 @@ func (_c *mockAuthService_RegisterBOSSession_Call) Run(run func(ctx context.Cont
 		if args[1] != nil {
 			arg1 = args[1].(state.ServerCookie)
 		}
+		var arg2 func(*state.Session)
+		if args[2] != nil {
+			arg2 = args[2].(func(*state.Session))
+		}
 		run(
 			arg0,
 			arg1,
+			arg2,
 		)
 	})
 	return _c
@@ -379,14 +385,14 @@ func (_c *mockAuthService_RegisterBOSSession_Call) Return(sessionInstance *state
 	return _c
 }
 
-func (_c *mockAuthService_RegisterBOSSession_Call) RunAndReturn(run func(ctx context.Context, authCookie state.ServerCookie) (*state.SessionInstance, error)) *mockAuthService_RegisterBOSSession_Call {
+func (_c *mockAuthService_RegisterBOSSession_Call) RunAndReturn(run func(ctx context.Context, authCookie state.ServerCookie, cfg func(*state.Session)) (*state.SessionInstance, error)) *mockAuthService_RegisterBOSSession_Call {
 	_c.Call.Return(run)
 	return _c
 }
 
 // RegisterChatSession provides a mock function for the type mockAuthService
-func (_mock *mockAuthService) RegisterChatSession(ctx context.Context, authCookie state.ServerCookie) (*state.SessionInstance, error) {
-	ret := _mock.Called(ctx, authCookie)
+func (_mock *mockAuthService) RegisterChatSession(ctx context.Context, authCookie state.ServerCookie, cfg func(sess *state.Session)) (*state.SessionInstance, error) {
+	ret := _mock.Called(ctx, authCookie, cfg)
 
 	if len(ret) == 0 {
 		panic("no return value specified for RegisterChatSession")
@@ -394,18 +400,18 @@ func (_mock *mockAuthService) RegisterChatSession(ctx context.Context, authCooki
 
 	var r0 *state.SessionInstance
 	var r1 error
-	if returnFunc, ok := ret.Get(0).(func(context.Context, state.ServerCookie) (*state.SessionInstance, error)); ok {
-		return returnFunc(ctx, authCookie)
+	if returnFunc, ok := ret.Get(0).(func(context.Context, state.ServerCookie, func(sess *state.Session)) (*state.SessionInstance, error)); ok {
+		return returnFunc(ctx, authCookie, cfg)
 	}
-	if returnFunc, ok := ret.Get(0).(func(context.Context, state.ServerCookie) *state.SessionInstance); ok {
-		r0 = returnFunc(ctx, authCookie)
+	if returnFunc, ok := ret.Get(0).(func(context.Context, state.ServerCookie, func(sess *state.Session)) *state.SessionInstance); ok {
+		r0 = returnFunc(ctx, authCookie, cfg)
 	} else {
 		if ret.Get(0) != nil {
 			r0 = ret.Get(0).(*state.SessionInstance)
 		}
 	}
-	if returnFunc, ok := ret.Get(1).(func(context.Context, state.ServerCookie) error); ok {
-		r1 = returnFunc(ctx, authCookie)
+	if returnFunc, ok := ret.Get(1).(func(context.Context, state.ServerCookie, func(sess *state.Session)) error); ok {
+		r1 = returnFunc(ctx, authCookie, cfg)
 	} else {
 		r1 = ret.Error(1)
 	}
@@ -420,11 +426,12 @@ type mockAuthService_RegisterChatSession_Call struct {
 // RegisterChatSession is a helper method to define mock.On call
 //   - ctx context.Context
 //   - authCookie state.ServerCookie
-func (_e *mockAuthService_Expecter) RegisterChatSession(ctx interface{}, authCookie interface{}) *mockAuthService_RegisterChatSession_Call {
-	return &mockAuthService_RegisterChatSession_Call{Call: _e.mock.On("RegisterChatSession", ctx, authCookie)}
+//   - cfg func(sess *state.Session)
+func (_e *mockAuthService_Expecter) RegisterChatSession(ctx interface{}, authCookie interface{}, cfg interface{}) *mockAuthService_RegisterChatSession_Call {
+	return &mockAuthService_RegisterChatSession_Call{Call: _e.mock.On("RegisterChatSession", ctx, authCookie, cfg)}
 }
 
-func (_c *mockAuthService_RegisterChatSession_Call) Run(run func(ctx context.Context, authCookie state.ServerCookie)) *mockAuthService_RegisterChatSession_Call {
+func (_c *mockAuthService_RegisterChatSession_Call) Run(run func(ctx context.Context, authCookie state.ServerCookie, cfg func(sess *state.Session))) *mockAuthService_RegisterChatSession_Call {
 	_c.Call.Run(func(args mock.Arguments) {
 		var arg0 context.Context
 		if args[0] != nil {
@@ -434,9 +441,14 @@ func (_c *mockAuthService_RegisterChatSession_Call) Run(run func(ctx context.Con
 		if args[1] != nil {
 			arg1 = args[1].(state.ServerCookie)
 		}
+		var arg2 func(sess *state.Session)
+		if args[2] != nil {
+			arg2 = args[2].(func(sess *state.Session))
+		}
 		run(
 			arg0,
 			arg1,
+			arg2,
 		)
 	})
 	return _c
@@ -447,7 +459,7 @@ func (_c *mockAuthService_RegisterChatSession_Call) Return(sessionInstance *stat
 	return _c
 }
 
-func (_c *mockAuthService_RegisterChatSession_Call) RunAndReturn(run func(ctx context.Context, authCookie state.ServerCookie) (*state.SessionInstance, error)) *mockAuthService_RegisterChatSession_Call {
+func (_c *mockAuthService_RegisterChatSession_Call) RunAndReturn(run func(ctx context.Context, authCookie state.ServerCookie, cfg func(sess *state.Session)) (*state.SessionInstance, error)) *mockAuthService_RegisterChatSession_Call {
 	_c.Call.Return(run)
 	return _c
 }
@@ -521,8 +533,8 @@ func (_c *mockAuthService_RetrieveBOSSession_Call) RunAndReturn(run func(ctx con
 }
 
 // Signout provides a mock function for the type mockAuthService
-func (_mock *mockAuthService) Signout(ctx context.Context, instance *state.SessionInstance) {
-	_mock.Called(ctx, instance)
+func (_mock *mockAuthService) Signout(ctx context.Context, session *state.Session) {
+	_mock.Called(ctx, session)
 	return
 }
 
@@ -533,20 +545,20 @@ type mockAuthService_Signout_Call struct {
 
 // Signout is a helper method to define mock.On call
 //   - ctx context.Context
-//   - instance *state.SessionInstance
-func (_e *mockAuthService_Expecter) Signout(ctx interface{}, instance interface{}) *mockAuthService_Signout_Call {
-	return &mockAuthService_Signout_Call{Call: _e.mock.On("Signout", ctx, instance)}
+//   - session *state.Session
+func (_e *mockAuthService_Expecter) Signout(ctx interface{}, session interface{}) *mockAuthService_Signout_Call {
+	return &mockAuthService_Signout_Call{Call: _e.mock.On("Signout", ctx, session)}
 }
 
-func (_c *mockAuthService_Signout_Call) Run(run func(ctx context.Context, instance *state.SessionInstance)) *mockAuthService_Signout_Call {
+func (_c *mockAuthService_Signout_Call) Run(run func(ctx context.Context, session *state.Session)) *mockAuthService_Signout_Call {
 	_c.Call.Run(func(args mock.Arguments) {
 		var arg0 context.Context
 		if args[0] != nil {
 			arg0 = args[0].(context.Context)
 		}
-		var arg1 *state.SessionInstance
+		var arg1 *state.Session
 		if args[1] != nil {
-			arg1 = args[1].(*state.SessionInstance)
+			arg1 = args[1].(*state.Session)
 		}
 		run(
 			arg0,
@@ -561,14 +573,14 @@ func (_c *mockAuthService_Signout_Call) Return() *mockAuthService_Signout_Call {
 	return _c
 }
 
-func (_c *mockAuthService_Signout_Call) RunAndReturn(run func(ctx context.Context, instance *state.SessionInstance)) *mockAuthService_Signout_Call {
+func (_c *mockAuthService_Signout_Call) RunAndReturn(run func(ctx context.Context, session *state.Session)) *mockAuthService_Signout_Call {
 	_c.Run(run)
 	return _c
 }
 
 // SignoutChat provides a mock function for the type mockAuthService
-func (_mock *mockAuthService) SignoutChat(ctx context.Context, instance *state.SessionInstance) {
-	_mock.Called(ctx, instance)
+func (_mock *mockAuthService) SignoutChat(ctx context.Context, sess *state.Session) {
+	_mock.Called(ctx, sess)
 	return
 }
 
@@ -579,20 +591,20 @@ type mockAuthService_SignoutChat_Call struct {
 
 // SignoutChat is a helper method to define mock.On call
 //   - ctx context.Context
-//   - instance *state.SessionInstance
-func (_e *mockAuthService_Expecter) SignoutChat(ctx interface{}, instance interface{}) *mockAuthService_SignoutChat_Call {
-	return &mockAuthService_SignoutChat_Call{Call: _e.mock.On("SignoutChat", ctx, instance)}
+//   - sess *state.Session
+func (_e *mockAuthService_Expecter) SignoutChat(ctx interface{}, sess interface{}) *mockAuthService_SignoutChat_Call {
+	return &mockAuthService_SignoutChat_Call{Call: _e.mock.On("SignoutChat", ctx, sess)}
 }
 
-func (_c *mockAuthService_SignoutChat_Call) Run(run func(ctx context.Context, instance *state.SessionInstance)) *mockAuthService_SignoutChat_Call {
+func (_c *mockAuthService_SignoutChat_Call) Run(run func(ctx context.Context, sess *state.Session)) *mockAuthService_SignoutChat_Call {
 	_c.Call.Run(func(args mock.Arguments) {
 		var arg0 context.Context
 		if args[0] != nil {
 			arg0 = args[0].(context.Context)
 		}
-		var arg1 *state.SessionInstance
+		var arg1 *state.Session
 		if args[1] != nil {
-			arg1 = args[1].(*state.SessionInstance)
+			arg1 = args[1].(*state.Session)
 		}
 		run(
 			arg0,
@@ -607,7 +619,7 @@ func (_c *mockAuthService_SignoutChat_Call) Return() *mockAuthService_SignoutCha
 	return _c
 }
 
-func (_c *mockAuthService_SignoutChat_Call) RunAndReturn(run func(ctx context.Context, instance *state.SessionInstance)) *mockAuthService_SignoutChat_Call {
+func (_c *mockAuthService_SignoutChat_Call) RunAndReturn(run func(ctx context.Context, sess *state.Session)) *mockAuthService_SignoutChat_Call {
 	_c.Run(run)
 	return _c
 }

server/toc/types.go

@@ -49,11 +49,11 @@ type AuthService interface {
 	BUCPLogin(ctx context.Context, inBody wire.SNAC_0x17_0x02_BUCPLoginRequest, advertisedHost string) (wire.SNACMessage, error)
 	CrackCookie(authCookie []byte) (state.ServerCookie, error)
 	FLAPLogin(ctx context.Context, inFrame wire.FLAPSignonFrame, advertisedHost string) (wire.TLVRestBlock, error)
-	RegisterBOSSession(ctx context.Context, authCookie state.ServerCookie) (*state.SessionInstance, error)
-	RegisterChatSession(ctx context.Context, authCookie state.ServerCookie) (*state.SessionInstance, error)
+	RegisterBOSSession(ctx context.Context, authCookie state.ServerCookie, cfg func(*state.Session)) (*state.SessionInstance, error)
+	RegisterChatSession(ctx context.Context, authCookie state.ServerCookie, cfg func(sess *state.Session)) (*state.SessionInstance, error)
 	RetrieveBOSSession(ctx context.Context, authCookie state.ServerCookie) (*state.SessionInstance, error)
-	Signout(ctx context.Context, instance *state.SessionInstance)
-	SignoutChat(ctx context.Context, instance *state.SessionInstance)
+	Signout(ctx context.Context, session *state.Session)
+	SignoutChat(ctx context.Context, sess *state.Session)
 }
 
 type LocateService interface {

server/webapi/handlers/oscar_bridge.go

@@ -30,11 +30,11 @@ type OSCARBridgeHandler struct {
 // OSCARAuthService defines methods needed for OSCAR authentication and session management.
 type OSCARAuthService interface {
 	// RegisterBOSSession creates a new BOS (Basic OSCAR Service) session
-	RegisterBOSSession(ctx context.Context, authCookie state.ServerCookie) (*state.SessionInstance, error)
+	RegisterBOSSession(ctx context.Context, authCookie state.ServerCookie, conf func(sess *state.Session)) (*state.SessionInstance, error)
 	// RetrieveBOSSession retrieves an existing BOS session
 	RetrieveBOSSession(ctx context.Context, authCookie state.ServerCookie) (*state.SessionInstance, error)
 	// Signout ends an OSCAR session
-	Signout(ctx context.Context, instance *state.SessionInstance)
+	Signout(ctx context.Context, session *state.Session)
 }
 
 // CookieBaker issues and validates authentication cookies for OSCAR services.

server/webapi/handlers/session.go

@@ -34,7 +34,7 @@ type SessionHandler struct {
 type AuthService interface {
 	BUCPChallenge(ctx context.Context, bodyIn wire.SNAC_0x17_0x06_BUCPChallengeRequest, newUUID func() uuid.UUID) (wire.SNACMessage, error)
 	BUCPLogin(ctx context.Context, bodyIn wire.SNAC_0x17_0x02_BUCPLoginRequest, advertisedHost string) (wire.SNACMessage, error)
-	RegisterBOSSession(ctx context.Context, authCookie state.ServerCookie) (*state.SessionInstance, error)
+	RegisterBOSSession(ctx context.Context, authCookie state.ServerCookie, conf func(sess *state.Session)) (*state.SessionInstance, error)
 }
 
 // SessionManager defines methods for OSCAR session management.

server/webapi/types.go

@@ -48,11 +48,11 @@ type AuthService interface {
 	BUCPLogin(ctx context.Context, inBody wire.SNAC_0x17_0x02_BUCPLoginRequest, advertisedHost string) (wire.SNACMessage, error)
 	CrackCookie(authCookie []byte) (state.ServerCookie, error)
 	FLAPLogin(ctx context.Context, inFrame wire.FLAPSignonFrame, advertisedHost string) (wire.TLVRestBlock, error)
-	RegisterBOSSession(ctx context.Context, authCookie state.ServerCookie) (*state.SessionInstance, error)
-	RegisterChatSession(ctx context.Context, authCookie state.ServerCookie) (*state.SessionInstance, error)
+	RegisterBOSSession(ctx context.Context, authCookie state.ServerCookie, conf func(sess *state.Session)) (*state.SessionInstance, error)
+	RegisterChatSession(ctx context.Context, authCookie state.ServerCookie, cfg func(sess *state.Session)) (*state.SessionInstance, error)
 	RetrieveBOSSession(ctx context.Context, authCookie state.ServerCookie) (*state.SessionInstance, error)
-	Signout(ctx context.Context, instance *state.SessionInstance)
-	SignoutChat(ctx context.Context, instance *state.SessionInstance)
+	Signout(ctx context.Context, session *state.Session)
+	SignoutChat(ctx context.Context, sess *state.Session)
 }
 
 type LocateService interface {

state/session.go

@@ -1095,10 +1095,8 @@ func (s *SessionInstance) CloseInstance() {
 	onInstanceCloseFn := s.onInstanceCloseFn
 	s.mutex.Unlock()
 
-	s.session.RemoveInstance(s)
-
 	count := s.session.InstanceCount()
-	if count == 0 {
+	if count == 1 {
 		s.session.mutex.RLock()
 		onSessCloseFn := s.session.onSessCloseFn
 		s.session.mutex.RUnlock()
@@ -1106,6 +1104,9 @@ func (s *SessionInstance) CloseInstance() {
 	} else {
 		onInstanceCloseFn()
 	}
+
+	s.session.RemoveInstance(s)
+
 }
 
 // OnClose registers a function to be called when the instance closes,

state/session_manager.go

@@ -177,7 +177,7 @@ func (s *InMemorySessionManager) maybeRelayMessageActiveOnly(ctx context.Context
 	}
 }
 
-func (s *InMemorySessionManager) AddSession(ctx context.Context, screenName DisplayScreenName, doMultiSess bool) (*SessionInstance, error) {
+func (s *InMemorySessionManager) AddSession(ctx context.Context, screenName DisplayScreenName, doMultiSess bool, cfg ...func(sess *Session)) (*SessionInstance, error) {
 	s.lockUser(screenName.IdentScreenName())
 	defer s.unlockUser(screenName.IdentScreenName())
 
@@ -189,7 +189,7 @@ func (s *InMemorySessionManager) AddSession(ctx context.Context, screenName Disp
 		if doMultiSess {
 			if !active.multiSession {
 				active.session.CloseSession()
-				return s.newSession(screenName, doMultiSess)
+				return s.newSession(screenName, doMultiSess, cfg)
 			}
 
 			// Check if we've reached the maximum number of concurrent sessions
@@ -212,14 +212,18 @@ func (s *InMemorySessionManager) AddSession(ctx context.Context, screenName Disp
 		}
 	}
 
-	return s.newSession(screenName, doMultiSess)
+	return s.newSession(screenName, doMultiSess, cfg)
 }
 
-func (s *InMemorySessionManager) newSession(screenName DisplayScreenName, doMultiSess bool) (*SessionInstance, error) {
+func (s *InMemorySessionManager) newSession(screenName DisplayScreenName, doMultiSess bool, cfg []func(sess *Session)) (*SessionInstance, error) {
 	sess := NewSession()
 	sess.SetIdentScreenName(screenName.IdentScreenName())
 	sess.SetDisplayScreenName(screenName)
 
+	for _, f := range cfg {
+		f(sess)
+	}
+
 	// Create a new instance within the session group
 	instance := sess.AddInstance()
 
@@ -244,11 +248,11 @@ func (s *InMemorySessionManager) findRec(identScreenName IdentScreenName) *sessi
 }
 
 // RemoveSession takes a session out of the session pool.
-func (s *InMemorySessionManager) RemoveSession(instance *SessionInstance) {
+func (s *InMemorySessionManager) RemoveSession(session *Session) {
 	s.mapMutex.Lock()
 	defer s.mapMutex.Unlock()
-	if rec, ok := s.store[instance.IdentScreenName()]; ok && rec.session == instance.Session() {
-		delete(s.store, instance.IdentScreenName())
+	if rec, ok := s.store[session.IdentScreenName()]; ok && rec.session == session {
+		delete(s.store, session.IdentScreenName())
 		close(rec.removed)
 	}
 }
@@ -321,8 +325,10 @@ type InMemoryChatSessionManager struct {
 }
 
 // AddSession adds a user to a chat room. If screenName already exists, the old
-// session is replaced by a new one.
-func (s *InMemoryChatSessionManager) AddSession(ctx context.Context, chatCookie string, screenName DisplayScreenName) (*SessionInstance, error) {
+// session is replaced by a new one. Optional cfg callbacks let callers attach
+// setup (for example shutdown behavior) while the session is still being
+// created.
+func (s *InMemoryChatSessionManager) AddSession(ctx context.Context, chatCookie string, screenName DisplayScreenName, cfg ...func(sess *Session)) (*SessionInstance, error) {
 	s.mapMutex.Lock()
 	if _, ok := s.store[chatCookie]; !ok {
 		s.store[chatCookie] = NewInMemorySessionManager(s.logger)
@@ -333,7 +339,7 @@ func (s *InMemoryChatSessionManager) AddSession(ctx context.Context, chatCookie
 	ctx, cancel := context.WithTimeout(ctx, time.Second*5)
 	defer cancel()
 
-	sess, err := sessionManager.AddSession(ctx, screenName, false)
+	sess, err := sessionManager.AddSession(ctx, screenName, false, cfg...)
 	if err != nil {
 		return nil, fmt.Errorf("AddSession: %w", err)
 	}
@@ -362,18 +368,18 @@ func (s *InMemoryChatSessionManager) AddSession(ctx context.Context, chatCookie
 
 // RemoveSession removes a user session from a chat room. It panics if you
 // attempt to remove the session twice.
-func (s *InMemoryChatSessionManager) RemoveSession(instance *SessionInstance) {
+func (s *InMemoryChatSessionManager) RemoveSession(sess *Session) {
 	s.mapMutex.Lock()
 	defer s.mapMutex.Unlock()
 
-	sessionManager, ok := s.store[instance.ChatRoomCookie()]
+	sessionManager, ok := s.store[sess.ChatRoomCookie()]
 	if !ok {
 		panic("attempting to remove a session after its room has been deleted")
 	}
-	sessionManager.RemoveSession(instance)
+	sessionManager.RemoveSession(sess)
 
 	if sessionManager.Empty() {
-		delete(s.store, instance.ChatRoomCookie())
+		delete(s.store, sess.ChatRoomCookie())
 	}
 }
 

state/session_manager_test.go

@@ -25,7 +25,7 @@ func TestInMemorySessionManager_AddSession(t *testing.T) {
 
 	go func() {
 		<-sess1.Closed()
-		sm.RemoveSession(sess1)
+		sm.RemoveSession(sess1.Session())
 	}()
 
 	sess2, err := sm.AddSession(ctx, "user-screen-name", false)
@@ -36,6 +36,24 @@ func TestInMemorySessionManager_AddSession(t *testing.T) {
 	assert.Contains(t, sm.AllSessions(), sess2.Session())
 }
 
+func TestInMemorySessionManager_AddSession_AppliesCfgToSession(t *testing.T) {
+	sm := NewInMemorySessionManager(slog.Default())
+	wantUIN := uint32(424242)
+	wantCookie := "cfg-test-cookie"
+
+	instance, err := sm.AddSession(context.Background(), "user-screen-name", false,
+		func(sess *Session) {
+			sess.SetUIN(wantUIN)
+			sess.SetChatRoomCookie(wantCookie)
+		},
+	)
+	assert.NoError(t, err)
+
+	s := instance.Session()
+	assert.Equal(t, wantUIN, s.UIN(), "cfg mutates Session before AddInstance / store insert")
+	assert.Equal(t, wantCookie, s.ChatRoomCookie())
+}
+
 func TestInMemorySessionManager_AddSession_Timeout(t *testing.T) {
 	sm := NewInMemorySessionManager(slog.Default())
 
@@ -66,7 +84,7 @@ func TestInMemorySessionManager_Remove_Existing(t *testing.T) {
 	assert.Equal(t, user1Old.Session(), rec.session)
 
 	// Remove the session
-	sm.RemoveSession(user1Old)
+	sm.RemoveSession(user1Old.Session())
 
 	// Verify the session is no longer in the store
 	_, ok = sm.store[user1Old.IdentScreenName()]
@@ -89,7 +107,7 @@ func TestInMemorySessionManager_Remove_Existing(t *testing.T) {
 	user2.SetSignonComplete()
 
 	// Remove user1New and verify it's gone
-	sm.RemoveSession(user1New)
+	sm.RemoveSession(user1New.Session())
 	_, ok = sm.store[user1New.IdentScreenName()]
 	assert.False(t, ok)
 
@@ -112,7 +130,7 @@ func TestInMemorySessionManager_Remove_MissingSameScreenName(t *testing.T) {
 	assert.Equal(t, user1Old.Session(), recOld.session)
 
 	// Remove the old session
-	sm.RemoveSession(user1Old)
+	sm.RemoveSession(user1Old.Session())
 	_, ok = sm.store[user1Old.IdentScreenName()]
 	assert.False(t, ok)
 
@@ -132,7 +150,7 @@ func TestInMemorySessionManager_Remove_MissingSameScreenName(t *testing.T) {
 	user2.SetSignonComplete()
 
 	// Try to remove the old session again - should do nothing because Session doesn't match
-	sm.RemoveSession(user1Old)
+	sm.RemoveSession(user1Old.Session())
 
 	// Verify the new session is still in the store (not removed)
 	recNewAfter, ok := sm.store[user1New.IdentScreenName()]
@@ -420,7 +438,7 @@ func TestInMemorySessionManager_SessionReplacement_NoMultiSess_NoMultiSess(t *te
 		synctest.Wait()
 
 		// AddSession() is blocked waiting for the lock, now unblock it
-		sm.RemoveSession(sess1)
+		sm.RemoveSession(sess1.Session())
 
 		wg.Wait()
 
@@ -466,7 +484,7 @@ func TestInMemorySessionManager_SessionReplacement_MultiSess_NoMultiSess(t *test
 
 		// AddSession() is blocked waiting for the lock, now unblock it
 		for _, sess := range sessList {
-			sm.RemoveSession(sess)
+			sm.RemoveSession(sess.Session())
 		}
 
 		wg.Wait()
@@ -506,7 +524,7 @@ func TestInMemorySessionManager_SessionReplacement_NoMultiSess_MultiSess(t *test
 		synctest.Wait()
 
 		// AddSession() is blocked waiting for the lock, now unblock it
-		sm.RemoveSession(sess1)
+		sm.RemoveSession(sess1.Session())
 
 		wg.Wait()
 
@@ -532,13 +550,30 @@ func TestInMemorySessionManager_RemoveSession_DoubleLogin_NoMultiSess_Chaos(t *t
 			sess1, err := sm.AddSession(context.Background(), "user-screen-name-1", false)
 			assert.NoError(t, err)
 			time.Sleep(time.Duration(rand.Intn(1000)) * time.Microsecond)
-			sm.RemoveSession(sess1)
+			sm.RemoveSession(sess1.Session())
 		}()
 	}
 
 	wg.Wait()
 }
 
+func TestInMemoryChatSessionManager_AddSession_AppliesCfgToSession(t *testing.T) {
+	sm := NewInMemoryChatSessionManager(slog.Default())
+	chatCookie := "chat-room-cfg"
+	wantUIN := uint32(777001)
+
+	instance, err := sm.AddSession(context.Background(), chatCookie, "user-screen-name",
+		func(sess *Session) {
+			sess.SetUIN(wantUIN)
+		},
+	)
+	assert.NoError(t, err)
+
+	s := instance.Session()
+	assert.Equal(t, wantUIN, s.UIN(), "cfg mutates Session via inner AddSession before instance is returned")
+	assert.Equal(t, chatCookie, s.ChatRoomCookie())
+}
+
 func TestInMemoryChatSessionManager_RelayToAllExcept_HappyPath(t *testing.T) {
 	sm := NewInMemoryChatSessionManager(slog.Default())
 
@@ -635,8 +670,8 @@ func TestInMemoryChatSessionManager_RemoveSession(t *testing.T) {
 
 	assert.Len(t, sm.AllSessions("chat-room-1"), 2)
 
-	sm.RemoveSession(user1)
-	sm.RemoveSession(user2)
+	sm.RemoveSession(user1.Session())
+	sm.RemoveSession(user2.Session())
 
 	assert.Empty(t, sm.AllSessions("chat-room-1"))
 }
@@ -669,7 +704,7 @@ func TestInMemoryChatSessionManager_RemoveSession_DoubleLogin(t *testing.T) {
 			synctest.Wait()
 
 			// AddSession() is blocked waiting for the lock, now unblock it
-			sm.RemoveSession(chatSess1)
+			sm.RemoveSession(chatSess1.Session())
 
 			wg.Wait()
 		})
@@ -1089,7 +1124,7 @@ func TestInMemorySessionManager_AddSession_MaxConcurrentSessions(t *testing.T) {
 		// Close and remove the first session to allow a new one
 		go func() {
 			<-sess1.Closed()
-			sm.RemoveSession(sess1)
+			sm.RemoveSession(sess1.Session())
 		}()
 
 		sess2, err := sm.AddSession(context.Background(), "user-screen-name-1", false)