feat(api, ssh): implement Bulk Writes in DeviceSetOnline

The `DeviceClaims` now includes a `Tenant` attribute. The `AuthDevice`
will utilize these attribute to populate the `X-Tenant-ID`. Furthermore,
the `/ssh/connection` endpoint has been updated to recognize and utilize
these headers for identifying the device's tenant.

Additionaly `DeviceSetOnline` method will now use bulk writes instead of
single operations. This change should improve performance and reduce
MongoDB I/O needs since the method is exclusively used in the
`api:heartbeat` queue, which processes a list of devices to set as
online. As this method's sole purpose is to set devices as online, it
will no longer accept a boolean parameter to set a device as offline.
Instead, a new method called `DeviceSetOffline` has been added.

Author: heiytor

api/routes/auth.go

@@ -168,6 +168,7 @@ func (h *Handler) AuthRequest(c gateway.Context) error {
 
 		// Extract device UID from JWT and set it into the header.
 		setHeader(c, client.DeviceUIDHeader, claims.UID)
+		setHeader(c, "X-Tenant-ID", claims.Tenant)
 
 		return c.NoContent(http.StatusOK)
 	default:

api/routes/device.go

@@ -175,7 +175,7 @@ func (h *Handler) OfflineDevice(c gateway.Context) error {
 		return err
 	}
 
-	if err := h.service.OfflineDevice(c.Ctx(), models.UID(req.UID), false); err != nil {
+	if err := h.service.OfflineDevice(c.Ctx(), models.UID(req.UID)); err != nil {
 		return err
 	}
 

api/routes/device_test.go

@@ -429,15 +429,15 @@ func TestOfflineDevice(t *testing.T) {
 			title: "fails when try to setting a non-existing device as offline",
 			uid:   "1234",
 			requiredMocks: func() {
-				mock.On("OfflineDevice", gomock.Anything, models.UID("1234"), false).Return(svc.ErrNotFound)
+				mock.On("OfflineDevice", gomock.Anything, models.UID("1234")).Return(svc.ErrNotFound)
 			},
 			expectedStatus: http.StatusNotFound,
 		},
 		{
 			title: "success when try to setting an existing device as offline",
 			uid:   "123",
 			requiredMocks: func() {
-				mock.On("OfflineDevice", gomock.Anything, models.UID("123"), false).Return(nil)
+				mock.On("OfflineDevice", gomock.Anything, models.UID("123")).Return(nil)
 			},
 			expectedStatus: http.StatusOK,
 		},

api/services/auth.go

@@ -58,7 +58,8 @@ func (s *service) AuthDevice(ctx context.Context, req requests.DeviceAuth, remot
 	token, err := jwttoken.New().
 		WithMethod(jwt.SigningMethodRS256).
 		WithClaims(&models.DeviceAuthClaims{
-			UID: key,
+			UID:    key,
+			Tenant: req.TenantID,
 			AuthClaims: models.AuthClaims{
 				Claims: "device",
 			},

api/services/device.go

@@ -2,14 +2,14 @@ package services
 
 import (
 	"context"
+	"errors"
 	"fmt"
 	"strings"
 	"time"
 
 	"github.com/shellhub-io/shellhub/api/store"
 	req "github.com/shellhub-io/shellhub/pkg/api/internalclient"
 	"github.com/shellhub-io/shellhub/pkg/api/query"
-	"github.com/shellhub-io/shellhub/pkg/clock"
 	"github.com/shellhub-io/shellhub/pkg/envs"
 	"github.com/shellhub-io/shellhub/pkg/models"
 	"github.com/shellhub-io/shellhub/pkg/validator"
@@ -24,7 +24,7 @@ type DeviceService interface {
 	DeleteDevice(ctx context.Context, uid models.UID, tenant string) error
 	RenameDevice(ctx context.Context, uid models.UID, name, tenant string) error
 	LookupDevice(ctx context.Context, namespace, name string) (*models.Device, error)
-	OfflineDevice(ctx context.Context, uid models.UID, online bool) error
+	OfflineDevice(ctx context.Context, uid models.UID) error
 	UpdateDeviceStatus(ctx context.Context, tenant string, uid models.UID, status models.DeviceStatus) error
 	UpdateDevice(ctx context.Context, tenant string, uid models.UID, name *string, publicURL *bool) error
 }
@@ -176,13 +176,16 @@ func (s *service) LookupDevice(ctx context.Context, namespace, name string) (*mo
 	return device, nil
 }
 
-func (s *service) OfflineDevice(ctx context.Context, uid models.UID, online bool) error {
-	err := s.store.DeviceSetOnline(ctx, uid, clock.Now(), online)
-	if err == store.ErrNoDocuments {
-		return NewErrDeviceNotFound(uid, err)
+func (s *service) OfflineDevice(ctx context.Context, uid models.UID) error {
+	if err := s.store.DeviceSetOffline(ctx, string(uid)); err != nil {
+		if errors.Is(err, store.ErrNoDocuments) {
+			return NewErrDeviceNotFound(uid, err)
+		}
+
+		return err
 	}
 
-	return err
+	return nil
 }
 
 // UpdateDeviceStatus updates the device status.

api/services/device_test.go

@@ -1613,52 +1613,60 @@ func TestLookupDevice(t *testing.T) {
 }
 
 func TestOfflineDevice(t *testing.T) {
-	mock := new(mocks.Store)
-
-	ctx := context.TODO()
+	storeMock := new(mocks.Store)
 
 	cases := []struct {
-		name          string
-		uid           models.UID
-		online        bool
-		requiredMocks func()
-		expected      error
+		name     string
+		uid      models.UID
+		mocks    func(context.Context)
+		expected error
 	}{
 		{
-			name: "fails when store device online fails",
+			name: "fails when operation does not succeeds",
 			uid:  models.UID("uid"),
-			requiredMocks: func() {
-				clockMock.On("Now").Return(now).Once()
-				mock.On("DeviceSetOnline", ctx, models.UID("uid"), now, false).
-					Return(errors.New("error", "", 0)).Once()
+			mocks: func(ctx context.Context) {
+				storeMock.
+					On("DeviceSetOffline", ctx, "uid").
+					Return(errors.New("error", "", 0)).
+					Once()
 			},
 			expected: errors.New("error", "", 0),
 		},
 		{
-			name:   "succeeds",
-			uid:    models.UID("uid"),
-			online: true,
-			requiredMocks: func() {
-				online := true
-				clockMock.On("Now").Return(now).Once()
-				mock.On("DeviceSetOnline", ctx, models.UID("uid"), now, online).
-					Return(errors.New("error", "", 0)).Once()
+			name: "fails when connected_device does not exist",
+			uid:  models.UID("uid"),
+			mocks: func(ctx context.Context) {
+				storeMock.
+					On("DeviceSetOffline", ctx, "uid").
+					Return(store.ErrNoDocuments).
+					Once()
 			},
-			expected: errors.New("error", "", 0),
+			expected: NewErrDeviceNotFound(models.UID("uid"), store.ErrNoDocuments),
+		},
+		{
+			name: "succeeds",
+			uid:  models.UID("uid"),
+			mocks: func(ctx context.Context) {
+				storeMock.
+					On("DeviceSetOffline", ctx, "uid").
+					Return(nil).
+					Once()
+			},
+			expected: nil,
 		},
 	}
 
+	s := NewService(store.Store(storeMock), privateKey, publicKey, storecache.NewNullCache(), clientMock, nil)
+
 	for _, tc := range cases {
 		t.Run(tc.name, func(t *testing.T) {
-			tc.requiredMocks()
-
-			service := NewService(store.Store(mock), privateKey, publicKey, storecache.NewNullCache(), clientMock, nil)
-			err := service.OfflineDevice(ctx, tc.uid, tc.online)
-			assert.Equal(t, tc.expected, err)
+			ctx := context.Background()
+			tc.mocks(ctx)
+			assert.Equal(t, tc.expected, s.OfflineDevice(ctx, tc.uid))
 		})
 	}
 
-	mock.AssertExpectations(t)
+	storeMock.AssertExpectations(t)
 }
 
 func TestUpdateDeviceStatus_same_mac(t *testing.T) {

api/services/mocks/services.go

@@ -28,7 +28,6 @@ type DeviceStore interface {
 	DeviceCreate(ctx context.Context, d models.Device, hostname string) error
 	DeviceRename(ctx context.Context, uid models.UID, hostname string) error
 	DeviceLookup(ctx context.Context, namespace, hostname string) (*models.Device, error)
-	DeviceSetOnline(ctx context.Context, uid models.UID, timestamp time.Time, online bool) error
 	DeviceUpdateOnline(ctx context.Context, uid models.UID, online bool) error
 	DeviceUpdateLastSeen(ctx context.Context, uid models.UID, ts time.Time) error
 	DeviceUpdateStatus(ctx context.Context, uid models.UID, status models.DeviceStatus) error
@@ -45,4 +44,11 @@ type DeviceStore interface {
 	DeviceRemovedList(ctx context.Context, tenant string, pagination query.Paginator, filters query.Filters, sorter query.Sorter) ([]models.DeviceRemoved, int, error)
 	DeviceCreatePublicURLAddress(ctx context.Context, uid models.UID) error
 	DeviceGetByPublicURLAddress(ctx context.Context, address string) (*models.Device, error)
+
+	// DeviceSetOnline receives a list of devices to mark as online. For each device in the array, it will upsert
+	// a connected device entry; each UID must exists in the "devices" collection.
+	DeviceSetOnline(ctx context.Context, connectedDevices []models.ConnectedDevice) error
+
+	// DeviceSetOffline sets a device's status to offline using its UID.
+	DeviceSetOffline(ctx context.Context, uid string) error
 }