geo/server/internal/tenant/app/desktop_client_service.go

package app

import (
	"context"
	"crypto/rand"
	"crypto/sha256"
	"encoding/hex"
	"fmt"
	"strings"
	"time"

	"github.com/google/uuid"
	"github.com/jackc/pgx/v5"
	goredis "github.com/redis/go-redis/v9"

	"github.com/geo-platform/tenant-api/internal/shared/auth"
	"github.com/geo-platform/tenant-api/internal/shared/response"
	"github.com/geo-platform/tenant-api/internal/tenant/repository"
)

const DesktopClientTokenTTL = 30 * 24 * time.Hour

type DesktopClientService struct {
	repo    repository.DesktopClientRepository
	redis   *goredis.Client
	taskSvc *DesktopTaskService
}

func NewDesktopClientService(repo repository.DesktopClientRepository, redis *goredis.Client) *DesktopClientService {
	return &DesktopClientService{repo: repo, redis: redis}
}

func (s *DesktopClientService) WithTaskService(taskSvc *DesktopTaskService) *DesktopClientService {
	s.taskSvc = taskSvc
	return s
}

type RegisterDesktopClientRequest struct {
	ClientID      string `json:"client_id" binding:"required"`
	DeviceName    string `json:"device_name" binding:"required"`
	OS            string `json:"os" binding:"required"`
	CPUArch       string `json:"cpu_arch"`
	ClientVersion string `json:"client_version" binding:"required"`
	Channel       string `json:"channel"`
}

type HeartbeatDesktopClientRequest struct {
	DeviceName    *string  `json:"device_name"`
	OS            *string  `json:"os"`
	CPUArch       *string  `json:"cpu_arch"`
	ClientVersion *string  `json:"client_version"`
	Channel       *string  `json:"channel"`
	AccountIDs    []string `json:"account_ids"`
	Startup       bool     `json:"startup"`
}

type DesktopClientView struct {
	ID            string     `json:"id"`
	TenantID      int64      `json:"tenant_id"`
	WorkspaceID   int64      `json:"workspace_id"`
	UserID        int64      `json:"user_id"`
	DeviceName    *string    `json:"device_name"`
	OS            *string    `json:"os"`
	CPUArch       *string    `json:"cpu_arch"`
	ClientVersion *string    `json:"client_version"`
	Channel       *string    `json:"channel"`
	CreatedAt     time.Time  `json:"created_at"`
	LastSeenAt    *time.Time `json:"last_seen_at"`
	LastRotatedAt *time.Time `json:"last_rotated_at"`
	RevokedAt     *time.Time `json:"revoked_at"`
}

type RegisterDesktopClientResponse struct {
	ClientID    string            `json:"client_id"`
	ClientToken string            `json:"client_token"`
	ExpiresAt   int64             `json:"expires_at"`
	Client      DesktopClientView `json:"client"`
}

type RotateDesktopClientResponse struct {
	ClientToken string            `json:"client_token"`
	ExpiresAt   int64             `json:"expires_at"`
	Client      DesktopClientView `json:"client"`
}

type HeartbeatDesktopClientResponse struct {
	Client                       DesktopClientView `json:"client"`
	ServerTime                   time.Time         `json:"server_time"`
	OnlineClientCount            int               `json:"online_client_count"`
	CurrentUserOnlineClientCount int               `json:"current_user_online_client_count"`
}

type OfflineDesktopClientResponse struct {
	ServerTime time.Time `json:"server_time"`
}

func HashDesktopClientToken(raw string) []byte {
	sum := sha256.Sum256([]byte(raw))
	return sum[:]
}

func (s *DesktopClientService) Register(ctx context.Context, actor auth.Actor, req RegisterDesktopClientRequest) (*RegisterDesktopClientResponse, error) {
	if actor.TenantID == 0 || actor.PrimaryWorkspaceID == 0 || actor.UserID == 0 {
		return nil, response.ErrUnauthorized(40107, "desktop_register_requires_actor", "desktop client registration requires authenticated workspace context")
	}

	rawToken, tokenHash, err := newDesktopClientToken()
	if err != nil {
		return nil, response.ErrInternal(50071, "desktop_client_token_failed", "failed to generate desktop client token")
	}

	clientID, err := parseRegisterClientID(req.ClientID)
	if err != nil {
		return nil, response.ErrBadRequest(40033, "invalid_desktop_client_id", "desktop client_id must be a non-empty uuid generated by the desktop client")
	}

	client, err := s.repo.Register(ctx, repository.RegisterDesktopClientParams{
		ID:            clientID,
		TenantID:      actor.TenantID,
		WorkspaceID:   actor.PrimaryWorkspaceID,
		UserID:        actor.UserID,
		TokenHash:     tokenHash,
		DeviceName:    requiredStringPtr(req.DeviceName),
		OS:            requiredStringPtr(req.OS),
		CPUArch:       trimmedStringPtr(req.CPUArch),
		ClientVersion: requiredStringPtr(req.ClientVersion),
		Channel:       trimmedStringPtr(req.Channel),
	})
	if err != nil {
		return nil, response.ErrInternal(50072, "desktop_client_register_failed", "failed to persist desktop client")
	}

	markDesktopClientPresentForUser(ctx, s.redis, client.ID, client.TenantID, client.WorkspaceID, client.UserID)

	return &RegisterDesktopClientResponse{
		ClientID:    client.ID.String(),
		ClientToken: rawToken,
		ExpiresAt:   time.Now().Add(DesktopClientTokenTTL).Unix(),
		Client:      buildDesktopClientView(client),
	}, nil
}

func (s *DesktopClientService) Rotate(ctx context.Context, client *repository.DesktopClient) (*RotateDesktopClientResponse, error) {
	if client == nil {
		return nil, response.ErrUnauthorized(40108, "desktop_client_missing", "desktop client context is required")
	}

	rawToken, tokenHash, err := newDesktopClientToken()
	if err != nil {
		return nil, response.ErrInternal(50073, "desktop_client_rotate_token_failed", "failed to generate next desktop client token")
	}

	updated, err := s.repo.RotateToken(ctx, repository.RotateDesktopClientTokenParams{
		ID:          client.ID,
		WorkspaceID: client.WorkspaceID,
		TokenHash:   tokenHash,
	})
	if err != nil {
		if err == pgx.ErrNoRows {
			return nil, response.ErrUnauthorized(40109, "desktop_client_revoked", "desktop client is revoked")
		}
		return nil, response.ErrInternal(50074, "desktop_client_rotate_failed", "failed to rotate desktop client token")
	}

	markDesktopClientPresentForUser(ctx, s.redis, updated.ID, updated.TenantID, updated.WorkspaceID, updated.UserID)

	return &RotateDesktopClientResponse{
		ClientToken: rawToken,
		ExpiresAt:   time.Now().Add(DesktopClientTokenTTL).Unix(),
		Client:      buildDesktopClientView(updated),
	}, nil
}

func (s *DesktopClientService) Heartbeat(ctx context.Context, client *repository.DesktopClient, req HeartbeatDesktopClientRequest) (*HeartbeatDesktopClientResponse, error) {
	if client == nil {
		return nil, response.ErrUnauthorized(40108, "desktop_client_missing", "desktop client context is required")
	}

	updated, err := s.repo.Heartbeat(ctx, repository.HeartbeatDesktopClientParams{
		ID:            client.ID,
		WorkspaceID:   client.WorkspaceID,
		DeviceName:    req.DeviceName,
		OS:            req.OS,
		CPUArch:       req.CPUArch,
		ClientVersion: req.ClientVersion,
		Channel:       req.Channel,
	})
	if err != nil {
		if err == pgx.ErrNoRows {
			return nil, response.ErrUnauthorized(40109, "desktop_client_revoked", "desktop client is revoked")
		}
		return nil, response.ErrInternal(50075, "desktop_client_heartbeat_failed", "failed to persist desktop client heartbeat")
	}

	markDesktopClientPresentForUser(ctx, s.redis, updated.ID, updated.TenantID, updated.WorkspaceID, updated.UserID)
	markDesktopAccountsPresent(ctx, s.redis, updated.ID, parseDesktopAccountIDs(req.AccountIDs))
	onlineClientCount, err := s.currentUserOnlineClientCount(ctx, updated)
	if err != nil {
		return nil, err
	}

	if req.Startup && s.taskSvc != nil {
		if err := s.taskSvc.recoverClientTasksOnStartup(ctx, updated); err != nil {
			return nil, err
		}
	}

	return &HeartbeatDesktopClientResponse{
		Client:                       buildDesktopClientView(updated),
		ServerTime:                   time.Now().UTC(),
		OnlineClientCount:            onlineClientCount,
		CurrentUserOnlineClientCount: onlineClientCount,
	}, nil
}

func (s *DesktopClientService) Revoke(ctx context.Context, client *repository.DesktopClient) (*DesktopClientView, error) {
	if client == nil {
		return nil, response.ErrUnauthorized(40108, "desktop_client_missing", "desktop client context is required")
	}

	updated, err := s.repo.Revoke(ctx, client.ID, client.WorkspaceID)
	if err != nil {
		if err == pgx.ErrNoRows {
			return nil, response.ErrUnauthorized(40109, "desktop_client_revoked", "desktop client is revoked")
		}
		return nil, response.ErrInternal(50076, "desktop_client_revoke_failed", "failed to revoke desktop client")
	}

	clearDesktopClientPresenceStateForUser(ctx, s.redis, updated.ID, updated.TenantID, updated.WorkspaceID, updated.UserID)
	if s.taskSvc != nil {
		if err := s.taskSvc.recoverClientTasksOnDisconnect(ctx, updated); err != nil {
			return nil, err
		}
	}

	view := buildDesktopClientView(updated)
	return &view, nil
}

func (s *DesktopClientService) Offline(ctx context.Context, client *repository.DesktopClient) (*OfflineDesktopClientResponse, error) {
	if client == nil {
		return nil, response.ErrUnauthorized(40108, "desktop_client_missing", "desktop client context is required")
	}

	clearDesktopClientPresenceStateForUser(ctx, s.redis, client.ID, client.TenantID, client.WorkspaceID, client.UserID)
	if s.taskSvc != nil {
		if err := s.taskSvc.recoverClientTasksOnDisconnect(ctx, client); err != nil {
			return nil, err
		}
	}

	return &OfflineDesktopClientResponse{
		ServerTime: time.Now().UTC(),
	}, nil
}

func (s *DesktopClientService) currentUserOnlineClientCount(ctx context.Context, client *repository.DesktopClient) (int, error) {
	if count, ok := loadDesktopUserOnlineClientCount(ctx, s.redis, client.TenantID, client.WorkspaceID, client.UserID, time.Now().UTC()); ok {
		return count, nil
	}

	clients, err := s.repo.ListByUser(ctx, client.TenantID, client.WorkspaceID, client.UserID)
	if err != nil {
		return 0, response.ErrInternal(50077, "desktop_client_online_count_failed", "failed to inspect online desktop clients")
	}

	return countOnlineDesktopClients(ctx, s.redis, clients, time.Now().UTC()), nil
}

func countOnlineDesktopClients(ctx context.Context, redis *goredis.Client, clients []repository.DesktopClient, now time.Time) int {
	if len(clients) == 0 {
		return 0
	}

	clientIDs := make([]uuid.UUID, 0, len(clients))
	for _, client := range clients {
		if client.ID == uuid.Nil || client.RevokedAt != nil {
			continue
		}
		clientIDs = append(clientIDs, client.ID)
	}
	presence := loadDesktopClientPresence(ctx, redis, clientIDs)

	onlineCount := 0
	for _, client := range clients {
		if client.ID == uuid.Nil || client.RevokedAt != nil {
			continue
		}
		if resolveDesktopClientOnline(&client, presence[client.ID], presence != nil, now) {
			onlineCount++
		}
	}
	return onlineCount
}

func newDesktopClientToken() (string, []byte, error) {
	raw := make([]byte, 32)
	if _, err := rand.Read(raw); err != nil {
		return "", nil, fmt.Errorf("generate desktop client token: %w", err)
	}
	token := hex.EncodeToString(raw)
	return token, HashDesktopClientToken(token), nil
}

func buildDesktopClientView(client *repository.DesktopClient) DesktopClientView {
	return DesktopClientView{
		ID:            client.ID.String(),
		TenantID:      client.TenantID,
		WorkspaceID:   client.WorkspaceID,
		UserID:        client.UserID,
		DeviceName:    client.DeviceName,
		OS:            client.OS,
		CPUArch:       client.CPUArch,
		ClientVersion: client.ClientVersion,
		Channel:       client.Channel,
		CreatedAt:     client.CreatedAt,
		LastSeenAt:    client.LastSeenAt,
		LastRotatedAt: client.LastRotatedAt,
		RevokedAt:     client.RevokedAt,
	}
}

func requiredStringPtr(value string) *string {
	trimmed := value
	return &trimmed
}

func trimmedStringPtr(value string) *string {
	if value == "" {
		return nil
	}
	trimmed := value
	return &trimmed
}

func parseDesktopAccountIDs(values []string) []uuid.UUID {
	if len(values) == 0 {
		return nil
	}

	accountIDs := make([]uuid.UUID, 0, len(values))
	for _, value := range values {
		parsed, err := uuid.Parse(strings.TrimSpace(value))
		if err != nil {
			continue
		}
		accountIDs = append(accountIDs, parsed)
	}
	return accountIDs
}

func parseRegisterClientID(value string) (uuid.UUID, error) {
	parsed, err := uuid.Parse(strings.TrimSpace(value))
	if err != nil || parsed == uuid.Nil {
		return uuid.Nil, fmt.Errorf("invalid desktop client id")
	}
	return parsed, nil
}