probo/pkg/iam/session_service.go

// Copyright (c) 2025 Probo Inc <hello@getprobo.com>.
//
// Permission to use, copy, modify, and/or distribute this software for any
// purpose with or without fee is hereby granted, provided that the above
// copyright notice and this permission notice appear in all copies.
//
// THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES WITH
// REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY
// AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY SPECIAL, DIRECT,
// INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM
// LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE OR
// OTHER TORTIOUS ACTION, ARISING OUT OF OR IN CONNECTION WITH THE USE OR
// PERFORMANCE OF THIS SOFTWARE.

package iam

import (
	"context"
	"fmt"
	"net"
	"time"

	"go.gearno.de/kit/pg"
	"go.gearno.de/x/ref"
	"go.probo.inc/probo/pkg/coredata"
	"go.probo.inc/probo/pkg/gid"
	"go.probo.inc/probo/pkg/validator"
)

type (
	SessionService struct {
		*Service
	}
)

func NewSessionService(svc *Service) *SessionService {
	return &SessionService{Service: svc}
}

type (
	RevokeAllSessionsRequest struct {
		CurrentSessionID gid.GID
	}
)

func (req RevokeAllSessionsRequest) Validate() error {
	v := validator.New()
	v.Check(req.CurrentSessionID, "current_session_id", validator.GID(coredata.SessionEntityType))
	return v.Error()
}

func (s SessionService) GetSession(ctx context.Context, sessionID gid.GID) (*coredata.Session, error) {
	var (
		session = &coredata.Session{}
		now     = time.Now()
	)

	err := s.pg.WithTx(
		ctx,
		func(tx pg.Conn) error {
			if err := session.LoadByID(ctx, tx, sessionID); err != nil {
				if err == coredata.ErrResourceNotFound {
					return NewSessionNotFoundError(sessionID)
				}
			}

			if session.ExpireReason != nil {
				return NewSessionExpiredError(sessionID)
			}

			if now.After(session.ExpiredAt) {
				session.ExpireReason = ref.Ref(coredata.ExpireReasonIdleTimeout)
				session.ExpiredAt = now
				session.UpdatedAt = now
				if err := session.Update(ctx, tx); err != nil {
					return fmt.Errorf("cannot update session: %w", err)
				}

				return NewSessionExpiredError(sessionID)
			}

			return nil
		},
	)

	if err != nil {
		return nil, err
	}

	return session, nil
}

func (s SessionService) CloseSession(ctx context.Context, sessionID gid.GID) error {
	return s.pg.WithTx(
		ctx,
		func(conn pg.Conn) error {
			session := &coredata.Session{}
			if err := session.LoadByID(ctx, conn, sessionID); err != nil {
				if err == coredata.ErrResourceNotFound {
					return NewSessionNotFoundError(sessionID)
				}

				return fmt.Errorf("cannot load session: %w", err)
			}

			if session.ExpireReason != nil {
				return NewSessionExpiredError(sessionID)
			}

			session.ExpireReason = ref.Ref(coredata.ExpireReasonClosed)
			session.ExpiredAt = time.Now()
			session.UpdatedAt = time.Now()
			if err := session.Update(ctx, conn); err != nil {
				if err == coredata.ErrResourceNotFound {
					return NewSessionNotFoundError(sessionID)
				}

				return fmt.Errorf("cannot update session: %w", err)
			}

			return nil
		},
	)
}

func (s SessionService) RevokeSession(ctx context.Context, identityID gid.GID, sessionID gid.GID) error {
	now := time.Now()

	return s.pg.WithTx(
		ctx,
		func(tx pg.Conn) error {
			user := &coredata.User{}
			err := user.LoadByID(ctx, tx, identityID)
			if err != nil {
				if err == coredata.ErrResourceNotFound {
					return NewUserNotFoundError(identityID)
				}

				return fmt.Errorf("cannot load user: %w", err)
			}

			session := &coredata.Session{}
			err = session.LoadByID(ctx, tx, sessionID)
			if err != nil {
				if err == coredata.ErrResourceNotFound {
					return NewSessionNotFoundError(sessionID)
				}

				return fmt.Errorf("cannot load session: %w", err)
			}

			// TODO: move to dedicated query instead of LoadByID
			if session.UserID != identityID {
				return NewSessionNotFoundError(sessionID)
			}

			if session.ExpireReason != nil {
				return NewSessionExpiredError(sessionID)
			}

			session.ExpireReason = ref.Ref(coredata.ExpireReasonRevoked)
			session.ExpiredAt = now
			session.UpdatedAt = now
			if err := session.Update(ctx, tx); err != nil {
				if err == coredata.ErrResourceNotFound {
					return NewSessionNotFoundError(sessionID)
				}

				return fmt.Errorf("cannot update session: %w", err)
			}

			return nil

		},
	)
}

func (s SessionService) RevokeAllSessions(ctx context.Context, currentSessionID gid.GID) (int64, error) {
	var count int64

	err := s.pg.WithTx(
		ctx,
		func(tx pg.Conn) error {
			session := coredata.Session{}
			err := session.LoadByID(ctx, tx, currentSessionID)
			if err != nil {
				if err == coredata.ErrResourceNotFound {
					return NewSessionNotFoundError(currentSessionID)
				}

				return fmt.Errorf("cannot load session: %w", err)
			}

			sessions := coredata.Sessions{}
			count, err = sessions.ExpireAllForUserExceptOneSession(ctx, tx, session.UserID, session.ID)
			if err != nil {
				return fmt.Errorf("cannot expire all sessions: %w", err)
			}

			return nil
		},
	)

	return count, err
}

func (s SessionService) UpdateSessionInfo(ctx context.Context, sessionID gid.GID, userAgent string, ipAddress net.IP) error {
	return s.pg.WithTx(
		ctx,
		func(tx pg.Conn) error {
			session := &coredata.Session{}
			err := session.LoadByID(ctx, tx, sessionID)
			if err != nil {
				if err == coredata.ErrResourceNotFound {
					return NewSessionNotFoundError(sessionID)
				}

				return fmt.Errorf("cannot load session: %w", err)
			}

			session.UserAgent = userAgent
			session.IPAddress = ipAddress
			session.UpdatedAt = time.Now()

			if err := session.Update(ctx, tx); err != nil {
				if err == coredata.ErrResourceNotFound {
					return NewSessionNotFoundError(sessionID)
				}

				return fmt.Errorf("cannot update session: %w", err)
			}

			return nil
		},
	)
}

func (s SessionService) UpdateSessionData(ctx context.Context, sessionID gid.GID, data coredata.SessionData) error {
	return s.pg.WithTx(
		ctx,
		func(tx pg.Conn) error {
			session := &coredata.Session{}
			err := session.LoadByID(ctx, tx, sessionID)
			if err != nil {
				if err == coredata.ErrResourceNotFound {
					return NewSessionNotFoundError(sessionID)
				}

				return fmt.Errorf("cannot load session: %w", err)
			}

			session.Data = data
			session.UpdatedAt = time.Now()

			if err := session.Update(ctx, tx); err != nil {
				if err == coredata.ErrResourceNotFound {
					return NewSessionNotFoundError(sessionID)
				}

				return fmt.Errorf("cannot update session: %w", err)
			}

			return nil
		},
	)
}