warmbly/internal/app/token/gen.go

package token

import (
	"context"
	"net/netip"
	"time"

	"github.com/getsentry/sentry-go"
	"github.com/golang-jwt/jwt/v5"
	"github.com/google/uuid"
	"github.com/mileusna/useragent"
	"github.com/warmbly/warmbly/internal/errx"
	"github.com/warmbly/warmbly/internal/infrastructure/db"
	"github.com/warmbly/warmbly/internal/models"
	"github.com/warmbly/warmbly/internal/pkg/crypt"
)

type TokenClaims struct {
	UserID    uuid.UUID `json:"sub"`
	SessionID uuid.UUID `json:"sid"`
	Email     string    `json:"email"`
	Nonce     string    `json:"nonce"`
	jwt.RegisteredClaims
}

func (s *tokenService) GenerateToken(userID, sessionID uuid.UUID, email, nonce string, issuedAt, expiresAt time.Time) (string, error) {
	claims := TokenClaims{
		UserID:    userID,
		SessionID: sessionID,
		Email:     email,
		Nonce:     nonce,
		RegisteredClaims: jwt.RegisteredClaims{
			ExpiresAt: jwt.NewNumericDate(expiresAt),
			IssuedAt:  jwt.NewNumericDate(issuedAt),
		},
	}
	token := jwt.NewWithClaims(jwt.SigningMethodHS256, claims)
	return token.SignedString([]byte(s.AuthSecret))
}

func (s *tokenService) VerifyToken(tokenStr string) (*TokenClaims, *errx.Error) {
	token, err := jwt.ParseWithClaims(tokenStr, &TokenClaims{}, func(token *jwt.Token) (any, error) {
		if _, ok := token.Method.(*jwt.SigningMethodHMAC); !ok {
			return nil, errx.ErrToken
		}
		return []byte(s.AuthSecret), nil
	})

	if err != nil {
		return nil, errx.ErrToken
	}

	claims, ok := token.Claims.(*TokenClaims)
	if !ok || !token.Valid {
		return nil, errx.ErrToken
	}

	if claims.ExpiresAt == nil || time.Now().After(claims.ExpiresAt.Time) {
		return nil, errx.ErrToken
	}

	return claims, nil
}

func (s *tokenService) GenerateSession(ctx context.Context, userID uuid.UUID, email, ipaddr, userAgent, authProvider string) (*models.Token, *errx.Error) {
	return s.GenerateSessionWithOrg(ctx, userID, email, ipaddr, userAgent, authProvider, nil)
}

func (s *tokenService) GenerateSessionWithOrg(ctx context.Context, userID uuid.UUID, email, ipaddr, userAgent, authProvider string, orgID *uuid.UUID) (*models.Token, *errx.Error) {
	ip, err := netip.ParseAddr(ipaddr)
	if err != nil {
		sentry.CaptureException(err)
		return nil, errx.InternalError()
	}

	ipinfo, err := s.geo.Lookup(ip)
	if err != nil {
		sentry.CaptureException(err)
		return nil, errx.InternalError()
	}

	userAgentInfo := useragent.Parse(userAgent)

	session := &models.Session{
		ID:                    uuid.New(),
		UserID:                userID,
		CurrentOrganizationID: orgID,

		LocationCity:        ipinfo.City,
		LocationRegion:      ipinfo.Region,
		LocationCountry:     ipinfo.Country,
		LocationCountryCode: ipinfo.CountryCode,
		LocationPostalCode:  ipinfo.PostalCode,

		BrowserName: userAgentInfo.Name,
		OSName:      userAgentInfo.OS,

		AuthProvider: authProvider,
	}

	tx, err := s.db.Begin(ctx)
	if err != nil {
		db.CaptureError(err, "", nil, "begin")
		return nil, errx.InternalError()
	}
	defer tx.Rollback(ctx)

	issuedAt := time.Now()
	session.LastRefreshedAt = issuedAt
	session.CreatedAt = issuedAt

	accessTokenExpiresAt := issuedAt.Add(AccessTokenLifeTime)
	accessNonce, err := crypt.Nonce()
	if err != nil {
		sentry.CaptureException(err)
		return nil, errx.InternalError()
	}
	session.AccessNonce = accessNonce

	accessToken, err := s.GenerateToken(userID, session.ID, email, accessNonce, issuedAt, accessTokenExpiresAt)
	if err != nil {
		sentry.CaptureException(err)
		return nil, errx.InternalError()
	}

	refreshTokenExpiresAt := issuedAt.Add(RefreshTokenLifeTime)
	session.ExpiresAt = &refreshTokenExpiresAt
	refreshNonce, err := crypt.Nonce()
	if err != nil {
		sentry.CaptureException(err)
		return nil, errx.InternalError()
	}
	session.RefreshNonce = refreshNonce

	refreshToken, err := s.GenerateToken(userID, session.ID, email, refreshNonce, issuedAt, refreshTokenExpiresAt)
	if err != nil {
		sentry.CaptureException(err)
		return nil, errx.InternalError()
	}

	if err := s.tokenRepository.GenerateSession(ctx, tx, session); err != nil {
		return nil, err
	}

	if err := tx.Commit(ctx); err != nil {
		db.CaptureError(err, "", nil, "commit")
		return nil, errx.InternalError()
	}

	return &models.Token{
		AccessToken:           accessToken,
		AccessTokenExpiresAt:  accessTokenExpiresAt,
		RefreshToken:          refreshToken,
		RefreshTokenExpiresAt: refreshTokenExpiresAt,
	}, nil
}

// SwitchOrganization updates the current organization for a session
func (s *tokenService) SwitchOrganization(ctx context.Context, sessionID uuid.UUID, orgID *uuid.UUID) *errx.Error {
	// Update in database
	if err := s.tokenRepository.UpdateCurrentOrganization(ctx, sessionID, orgID); err != nil {
		return err
	}

	// Invalidate cache for this session
	s.deleteSession(ctx, sessionID)

	return nil
}

// GetCurrentOrganization retrieves the current organization for a session
func (s *tokenService) GetCurrentOrganization(ctx context.Context, sessionID uuid.UUID) (*uuid.UUID, *errx.Error) {
	session, err := s.GetSession(ctx, sessionID)
	if err != nil {
		return nil, err
	}
	return session.CurrentOrganizationID, nil
}