api: add cache to avoid store overload

Co-authored-by: vsaller <vsaller@yahoo.com> Co-authored-by: Luis Gustavo S. Barreto <gustavo@ossystems.com.br>
shellhub-io · Jul 28, 2021 · bddbed3 · bddbed3
1 parent ce9877c
commit bddbed3
Show file tree

Hide file tree

Showing 15 changed files with 78 additions and 50 deletions.
diff --git a/api/store/cache/cache.go → api/cache/cache.go b/api/store/cache/cache.go → api/cache/cache.go
diff --git a/api/store/cache/cache_null.go → api/cache/cache_null.go b/api/store/cache/cache_null.go → api/cache/cache_null.go
diff --git a/api/store/cache/cache_redis.go → api/cache/cache_redis.go b/api/store/cache/cache_redis.go → api/cache/cache_redis.go
diff --git a/api/server.go b/api/server.go
@@ -7,10 +7,10 @@ import (
 	"github.com/labstack/echo/v4"
 	"github.com/labstack/echo/v4/middleware"
 	"github.com/shellhub-io/shellhub/api/apicontext"
+	storecache "github.com/shellhub-io/shellhub/api/cache"
 	"github.com/shellhub-io/shellhub/api/routes"
 	"github.com/shellhub-io/shellhub/api/routes/middlewares"
 	"github.com/shellhub-io/shellhub/api/services"
-	storecache "github.com/shellhub-io/shellhub/api/store/cache"
 	"github.com/shellhub-io/shellhub/api/store/mongo"
 	"github.com/sirupsen/logrus"
 	"github.com/spf13/cobra"
@@ -81,7 +81,7 @@ func startServer() error {
 
 	// apply dependency injection through project layers
 	store := mongo.NewStore(client.Database("main"), cache)
-	service := services.NewService(store, nil, nil)
+	service := services.NewService(store, nil, nil, cache)
 	handler := routes.NewHandler(service)
 
 	e.Use(func(next echo.HandlerFunc) echo.HandlerFunc {

diff --git a/api/services/auth.go b/api/services/auth.go
@@ -35,6 +35,35 @@ type AuthService interface {
 func (s *service) AuthDevice(ctx context.Context, req *models.DeviceAuthRequest) (*models.DeviceAuthResponse, error) {
 	uid := sha256.Sum256(structhash.Dump(req.DeviceAuth, 1))
 
+	key := hex.EncodeToString(uid[:])
+
+	token := jwt.NewWithClaims(jwt.SigningMethodRS256, models.DeviceAuthClaims{
+		UID: hex.EncodeToString(uid[:]),
+		AuthClaims: models.AuthClaims{
+			Claims: "device",
+		},
+	})
+
+	tokenStr, err := token.SignedString(s.privKey)
+	if err != nil {
+		return nil, err
+	}
+
+	type Device struct {
+		Name      string
+		Namespace string
+	}
+
+	var value *Device
+
+	if err := s.cache.Get(ctx, strings.Join([]string{"auth_device", key}, "/"), &value); err == nil && value != nil {
+		return &models.DeviceAuthResponse{
+			UID:       hex.EncodeToString(uid[:]),
+			Token:     tokenStr,
+			Name:      value.Name,
+			Namespace: value.Namespace,
+		}, nil
+	}
 	device := models.Device{
 		UID:       hex.EncodeToString(uid[:]),
 		Identity:  req.Identity,
@@ -60,18 +89,6 @@ func (s *service) AuthDevice(ctx context.Context, req *models.DeviceAuthRequest)
 		return nil, err
 	}
 
-	token := jwt.NewWithClaims(jwt.SigningMethodRS256, models.DeviceAuthClaims{
-		UID: hex.EncodeToString(uid[:]),
-		AuthClaims: models.AuthClaims{
-			Claims: "device",
-		},
-	})
-
-	tokenStr, err := token.SignedString(s.privKey)
-	if err != nil {
-		return nil, err
-	}
-
 	if err := s.store.DeviceSetOnline(ctx, models.UID(device.UID), true); err != nil {
 		return nil, err
 	}
@@ -86,6 +103,9 @@ func (s *service) AuthDevice(ctx context.Context, req *models.DeviceAuthRequest)
 	if err != nil {
 		return nil, err
 	}
+	if err := s.cache.Set(ctx, strings.Join([]string{"auth_device", key}, "/"), &Device{Name: dev.Name, Namespace: namespace.Name}, time.Second*30); err != nil {
+		return nil, err
+	}
 
 	return &models.DeviceAuthResponse{
 		UID:       hex.EncodeToString(uid[:]),

diff --git a/api/services/auth_test.go b/api/services/auth_test.go
@@ -11,6 +11,7 @@ import (
 	"time"
 
 	"github.com/cnf/structhash"
+	storecache "github.com/shellhub-io/shellhub/api/cache"
 	"github.com/shellhub-io/shellhub/api/store"
 	"github.com/shellhub-io/shellhub/api/store/mocks"
 	"github.com/shellhub-io/shellhub/pkg/clock"
@@ -25,7 +26,7 @@ func TestAuthDevice(t *testing.T) {
 	privateKey, err := rsa.GenerateKey(rand.Reader, 2048)
 	assert.NoError(t, err)
 
-	s := NewService(store.Store(mock), privateKey, &privateKey.PublicKey)
+	s := NewService(store.Store(mock), privateKey, &privateKey.PublicKey, storecache.NewNullCache())
 
 	ctx := context.TODO()
 
@@ -83,7 +84,7 @@ func TestAuthUser(t *testing.T) {
 	privateKey, err := rsa.GenerateKey(rand.Reader, 2048)
 	assert.NoError(t, err)
 
-	s := NewService(store.Store(mock), privateKey, &privateKey.PublicKey)
+	s := NewService(store.Store(mock), privateKey, &privateKey.PublicKey, storecache.NewNullCache())
 
 	ctx := context.TODO()
 
@@ -123,7 +124,7 @@ func TestAuthUserInfo(t *testing.T) {
 	privateKey, err := rsa.GenerateKey(rand.Reader, 2048)
 	assert.NoError(t, err)
 
-	s := NewService(store.Store(mock), privateKey, &privateKey.PublicKey)
+	s := NewService(store.Store(mock), privateKey, &privateKey.PublicKey, storecache.NewNullCache())
 
 	ctx := context.TODO()
 

diff --git a/api/services/device_test.go b/api/services/device_test.go
@@ -7,6 +7,7 @@ import (
 	"errors"
 	"testing"
 
+	storecache "github.com/shellhub-io/shellhub/api/cache"
 	"github.com/shellhub-io/shellhub/api/store"
 	"github.com/shellhub-io/shellhub/api/store/mocks"
 	"github.com/shellhub-io/shellhub/pkg/api/paginator"
@@ -17,7 +18,7 @@ import (
 
 func TestListDevices(t *testing.T) {
 	mock := &mocks.Store{}
-	s := NewService(store.Store(mock), privateKey, publicKey)
+	s := NewService(store.Store(mock), privateKey, publicKey, storecache.NewNullCache())
 
 	ctx := context.TODO()
 
@@ -109,7 +110,7 @@ func TestListDevices(t *testing.T) {
 
 func TestGetDevice(t *testing.T) {
 	mock := &mocks.Store{}
-	s := NewService(store.Store(mock), privateKey, publicKey)
+	s := NewService(store.Store(mock), privateKey, publicKey, storecache.NewNullCache())
 
 	Err := errors.New("error")
 
@@ -167,7 +168,7 @@ func TestGetDevice(t *testing.T) {
 
 func TestDeleteDevice(t *testing.T) {
 	mock := &mocks.Store{}
-	s := NewService(store.Store(mock), privateKey, publicKey)
+	s := NewService(store.Store(mock), privateKey, publicKey, storecache.NewNullCache())
 
 	ctx := context.TODO()
 
@@ -262,7 +263,7 @@ func TestDeleteDevice(t *testing.T) {
 
 func TestRenameDevice(t *testing.T) {
 	mock := &mocks.Store{}
-	s := NewService(store.Store(mock), privateKey, publicKey)
+	s := NewService(store.Store(mock), privateKey, publicKey, storecache.NewNullCache())
 
 	ctx := context.TODO()
 
@@ -427,7 +428,7 @@ func TestRenameDevice(t *testing.T) {
 
 func TestLookupDevice(t *testing.T) {
 	mock := &mocks.Store{}
-	s := NewService(store.Store(mock), privateKey, publicKey)
+	s := NewService(store.Store(mock), privateKey, publicKey, storecache.NewNullCache())
 
 	ctx := context.TODO()
 
@@ -487,7 +488,7 @@ func TestLookupDevice(t *testing.T) {
 
 func TestUpdateDeviceStatus(t *testing.T) {
 	mock := &mocks.Store{}
-	s := NewService(store.Store(mock), privateKey, publicKey)
+	s := NewService(store.Store(mock), privateKey, publicKey, storecache.NewNullCache())
 
 	Err := errors.New("error")
 
@@ -535,7 +536,7 @@ func TestUpdateDeviceStatus(t *testing.T) {
 
 func TestUpdatePendingStatus(t *testing.T) {
 	mock := &mocks.Store{}
-	s := NewService(store.Store(mock), privateKey, publicKey)
+	s := NewService(store.Store(mock), privateKey, publicKey, storecache.NewNullCache())
 
 	user := &models.User{Name: "name", Username: "username", ID: "id"}
 	user2 := &models.User{Name: "name2", Username: "username2", ID: "id2"}

diff --git a/api/services/namespace_test.go b/api/services/namespace_test.go
@@ -7,6 +7,7 @@ import (
 	"strings"
 	"testing"
 
+	storecache "github.com/shellhub-io/shellhub/api/cache"
 	utils "github.com/shellhub-io/shellhub/api/pkg/namespace"
 	"github.com/shellhub-io/shellhub/api/store"
 	"github.com/shellhub-io/shellhub/api/store/mocks"
@@ -116,7 +117,7 @@ func TestIsNamespaceOwner(t *testing.T) {
 
 func TestListNamespaces(t *testing.T) {
 	mock := &mocks.Store{}
-	s := NewService(store.Store(mock), privateKey, publicKey)
+	s := NewService(store.Store(mock), privateKey, publicKey, storecache.NewNullCache())
 
 	ctx := context.TODO()
 	Err := errors.New("error")
@@ -185,7 +186,7 @@ func TestListNamespaces(t *testing.T) {
 
 func TestGetNamespace(t *testing.T) {
 	mock := &mocks.Store{}
-	s := NewService(store.Store(mock), privateKey, publicKey)
+	s := NewService(store.Store(mock), privateKey, publicKey, storecache.NewNullCache())
 
 	ctx := context.TODO()
 
@@ -244,7 +245,7 @@ func TestGetNamespace(t *testing.T) {
 
 func TestListMembers(t *testing.T) {
 	mock := &mocks.Store{}
-	s := NewService(store.Store(mock), privateKey, publicKey)
+	s := NewService(store.Store(mock), privateKey, publicKey, storecache.NewNullCache())
 
 	ctx := context.TODO()
 	envMock := &env_mocks.Backend{}
@@ -353,7 +354,7 @@ func TestListMembers(t *testing.T) {
 
 func TestCreateNamespace(t *testing.T) {
 	mock := &mocks.Store{}
-	s := NewService(store.Store(mock), privateKey, publicKey)
+	s := NewService(store.Store(mock), privateKey, publicKey, storecache.NewNullCache())
 
 	ctx := context.TODO()
 	envMock := &env_mocks.Backend{}
@@ -588,7 +589,7 @@ func TestCreateNamespace(t *testing.T) {
 
 func TestEditNamespace(t *testing.T) {
 	mock := &mocks.Store{}
-	s := NewService(store.Store(mock), privateKey, publicKey)
+	s := NewService(store.Store(mock), privateKey, publicKey, storecache.NewNullCache())
 
 	ctx := context.TODO()
 
@@ -711,7 +712,7 @@ func TestEditNamespace(t *testing.T) {
 
 func TestDeleteNamespace(t *testing.T) {
 	mock := &mocks.Store{}
-	s := NewService(store.Store(mock), privateKey, publicKey)
+	s := NewService(store.Store(mock), privateKey, publicKey, storecache.NewNullCache())
 
 	Err := errors.New("error")
 
@@ -775,7 +776,7 @@ func TestDeleteNamespace(t *testing.T) {
 
 func TestAddNamespaceUser(t *testing.T) {
 	mock := &mocks.Store{}
-	s := NewService(store.Store(mock), privateKey, publicKey)
+	s := NewService(store.Store(mock), privateKey, publicKey, storecache.NewNullCache())
 	ctx := context.TODO()
 
 	namespace := &models.Namespace{Name: "group1", Owner: "ID1", TenantID: "a736a52b-5777-4f92-b0b8-e359bf484713"}
@@ -887,7 +888,7 @@ func TestAddNamespaceUser(t *testing.T) {
 
 func TestRemoveNamespaceUser(t *testing.T) {
 	mock := &mocks.Store{}
-	s := NewService(store.Store(mock), privateKey, publicKey)
+	s := NewService(store.Store(mock), privateKey, publicKey, storecache.NewNullCache())
 	ctx := context.TODO()
 	user := &models.User{Name: "user1", Username: "username1", ID: "hash1"}
 	user2 := &models.User{Name: "user2", Username: "username2", ID: "hash2"}
@@ -997,7 +998,7 @@ func TestRemoveNamespaceUser(t *testing.T) {
 
 func TestGetSessionRecord(t *testing.T) {
 	mock := &mocks.Store{}
-	s := NewService(store.Store(mock), privateKey, publicKey)
+	s := NewService(store.Store(mock), privateKey, publicKey, storecache.NewNullCache())
 
 	ctx := context.TODO()
 
@@ -1073,7 +1074,7 @@ func TestGetSessionRecord(t *testing.T) {
 
 func TestEditSessionRecord(t *testing.T) {
 	mock := &mocks.Store{}
-	s := NewService(store.Store(mock), privateKey, publicKey)
+	s := NewService(store.Store(mock), privateKey, publicKey, storecache.NewNullCache())
 
 	ctx := context.TODO()
 

diff --git a/api/services/service.go b/api/services/service.go
@@ -3,13 +3,15 @@ package services
 import (
 	"crypto/rsa"
 
+	"github.com/shellhub-io/shellhub/api/cache"
 	"github.com/shellhub-io/shellhub/api/store"
 )
 
 type service struct {
 	store   store.Store
 	privKey *rsa.PrivateKey
 	pubKey  *rsa.PublicKey
+	cache   cache.Cache
 }
 
 type Service interface {
@@ -22,7 +24,7 @@ type Service interface {
 	StatsService
 }
 
-func NewService(store store.Store, privKey *rsa.PrivateKey, pubKey *rsa.PublicKey) Service {
+func NewService(store store.Store, privKey *rsa.PrivateKey, pubKey *rsa.PublicKey, cache cache.Cache) Service {
 	if privKey == nil || pubKey == nil {
 		var err error
 		privKey, pubKey, err = LoadKeys()
@@ -31,5 +33,5 @@ func NewService(store store.Store, privKey *rsa.PrivateKey, pubKey *rsa.PublicKe
 		}
 	}
 
-	return &service{store, privKey, pubKey}
+	return &service{store, privKey, pubKey, cache}
 }
diff --git a/api/services/session_test.go b/api/services/session_test.go
@@ -5,6 +5,7 @@ import (
 	"errors"
 	"testing"
 
+	storecache "github.com/shellhub-io/shellhub/api/cache"
 	"github.com/shellhub-io/shellhub/api/store"
 	"github.com/shellhub-io/shellhub/api/store/mocks"
 	"github.com/shellhub-io/shellhub/pkg/api/paginator"
@@ -14,7 +15,7 @@ import (
 
 func TestListSessions(t *testing.T) {
 	mock := &mocks.Store{}
-	s := NewService(store.Store(mock), privateKey, publicKey)
+	s := NewService(store.Store(mock), privateKey, publicKey, storecache.NewNullCache())
 
 	ctx := context.TODO()
 
@@ -81,7 +82,7 @@ func TestListSessions(t *testing.T) {
 
 func TestGetSession(t *testing.T) {
 	mock := &mocks.Store{}
-	s := NewService(store.Store(mock), privateKey, publicKey)
+	s := NewService(store.Store(mock), privateKey, publicKey, storecache.NewNullCache())
 
 	ctx := context.TODO()
 
@@ -140,7 +141,7 @@ func TestGetSession(t *testing.T) {
 
 func TestCreateSession(t *testing.T) {
 	mock := &mocks.Store{}
-	s := NewService(store.Store(mock), privateKey, publicKey)
+	s := NewService(store.Store(mock), privateKey, publicKey, storecache.NewNullCache())
 
 	ctx := context.TODO()
 
@@ -198,7 +199,7 @@ func TestCreateSession(t *testing.T) {
 
 func TestDeactivateSession(t *testing.T) {
 	mock := &mocks.Store{}
-	s := NewService(store.Store(mock), privateKey, publicKey)
+	s := NewService(store.Store(mock), privateKey, publicKey, storecache.NewNullCache())
 
 	ctx := context.TODO()
 
@@ -243,7 +244,7 @@ func TestDeactivateSession(t *testing.T) {
 
 func TestSetSessionAuthenticated(t *testing.T) {
 	mock := &mocks.Store{}
-	s := NewService(store.Store(mock), privateKey, publicKey)
+	s := NewService(store.Store(mock), privateKey, publicKey, storecache.NewNullCache())
 
 	ctx := context.TODO()