From 87400ca5d4fd004331b525ee0d41c0cfa45904cf Mon Sep 17 00:00:00 2001
From: "Artur M. Wolff" <artur@storj.io>
Date: Wed, 27 Apr 2022 17:07:19 +0200
Subject: [PATCH] pkg/auth: limit access grant length @ DRPC

This change introduces an access grant limiting mechanism, similar to
what https://review.dev.storj.io/c/storj/gateway-mt/+/7281 brought.

Closes storj/storj-private#30

Change-Id: I8f1a823cb6ad020667796c20aa4cb588130d041f
---
 pkg/auth/drpcauth/server.go      | 27 +++++++++++++++++------
 pkg/auth/drpcauth/server_test.go | 37 +++++++++++++++++++++++++++++---
 pkg/auth/peer.go                 |  2 +-
 3 files changed, 56 insertions(+), 10 deletions(-)

diff --git a/pkg/auth/drpcauth/server.go b/pkg/auth/drpcauth/server.go
index fde96e77..c7e67f45 100644
--- a/pkg/auth/drpcauth/server.go
+++ b/pkg/auth/drpcauth/server.go
@@ -16,8 +16,10 @@ import (
 	"net/url"
 
 	"github.com/spacemonkeygo/monkit/v3"
+	"github.com/zeebo/errs"
 	"go.uber.org/zap"
 
+	"storj.io/common/memory"
 	"storj.io/common/pb"
 	"storj.io/common/rpc/rpcstatus"
 	"storj.io/drpc/drpcmux"
@@ -36,8 +38,9 @@ type Server struct {
 
 	// This is duplicated with package storj.io/gateway-mt/pkg/auth/httpauth/resources
 	// TODO: factor out common functionality
-	db       *authdb.Database
-	endpoint *url.URL
+	db                   *authdb.Database
+	endpoint             *url.URL
+	accessGrantSizeLimit memory.Size
 }
 
 // NewServer creates a Server that is not running.
@@ -45,11 +48,13 @@ func NewServer(
 	log *zap.Logger,
 	db *authdb.Database,
 	endpoint *url.URL,
+	accessGrantSizeLimit memory.Size,
 ) *Server {
 	return &Server{
-		log:      log,
-		db:       db,
-		endpoint: endpoint,
+		log:                  log,
+		db:                   db,
+		endpoint:             endpoint,
+		accessGrantSizeLimit: accessGrantSizeLimit,
 	}
 }
 
@@ -63,8 +68,18 @@ func (g *Server) RegisterAccess(
 
 	g.log.Debug("DRPC RegisterAccess request")
 
-	response, err := g.registerAccessImpl(ctx, request)
+	// NOTE(artur): DRPC's default message limit is 4 MiB, so we will read such
+	// messages anyway, but Auth Service would blow up memory consumption
+	// because it copies this access grant several times later on. Avoiding
+	// processing the access grant should effectively mitigate this kind of DoS
+	// attack.
+	if len(request.AccessGrant) > g.accessGrantSizeLimit.Int() {
+		err = errs.New("provided access grant is too large")
+		g.log.Error("DRPC RegisterAccess failed", zap.Error(err))
+		return nil, rpcstatus.Wrap(rpcstatus.InvalidArgument, err)
+	}
 
+	response, err := g.registerAccessImpl(ctx, request)
 	if err != nil {
 		g.log.Error("DRPC RegisterAccess failed", zap.Error(err))
 		err = rpcstatus.Wrap(rpcstatus.Internal, err)
diff --git a/pkg/auth/drpcauth/server_test.go b/pkg/auth/drpcauth/server_test.go
index 74758c40..8d3d3787 100644
--- a/pkg/auth/drpcauth/server_test.go
+++ b/pkg/auth/drpcauth/server_test.go
@@ -6,10 +6,13 @@ import (
 	"net/url"
 	"testing"
 
+	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 	"go.uber.org/zap/zaptest"
 
+	"storj.io/common/memory"
 	"storj.io/common/pb"
+	"storj.io/common/rpc/rpcstatus"
 	"storj.io/common/storj"
 	"storj.io/common/testcontext"
 	"storj.io/gateway-mt/pkg/auth/authdb"
@@ -30,21 +33,21 @@ var minimalAccessSatelliteID = func() storj.NodeID {
 	return url.ID
 }()
 
-func createBackend(t *testing.T) (*Server, *authdb.Database) {
+func createBackend(t *testing.T, sizeLimit memory.Size) (*Server, *authdb.Database) {
 	endpoint, err := url.Parse("http://gateway.test")
 	require.NoError(t, err)
 	allowedSatelliteIDs := map[storj.NodeID]struct{}{minimalAccessSatelliteID: {}}
 
 	db := authdb.NewDatabase(memauth.New(), allowedSatelliteIDs)
 
-	return NewServer(zaptest.NewLogger(t), db, endpoint), db
+	return NewServer(zaptest.NewLogger(t), db, endpoint, sizeLimit), db
 }
 
 func TestRegisterAccess(t *testing.T) {
 	ctx := testcontext.New(t)
 	defer ctx.Cleanup()
 
-	server, db := createBackend(t)
+	server, db := createBackend(t, 4*memory.KiB)
 
 	response, err := server.RegisterAccess(
 		ctx,
@@ -73,3 +76,31 @@ func TestRegisterAccess(t *testing.T) {
 	require.Equal(t, minimalAccess, storedAccessGrant)
 	require.Equal(t, response.SecretKey, storedSecretKey.ToBase32())
 }
+
+func TestRegisterAccessTooLarge(t *testing.T) {
+	ctx := testcontext.New(t)
+	defer ctx.Cleanup()
+
+	server, _ := createBackend(t, memory.Size(len(minimalAccess)))
+
+	_, err := server.RegisterAccess(
+		ctx,
+		&pb.EdgeRegisterAccessRequest{
+			AccessGrant: minimalAccess + "a",
+			Public:      false,
+		},
+	)
+	require.Error(t, err)
+
+	assert.Equal(t, rpcstatus.Code(err), rpcstatus.InvalidArgument)
+	assert.EqualError(t, err, "provided access grant is too large")
+
+	_, err = server.RegisterAccess(
+		ctx,
+		&pb.EdgeRegisterAccessRequest{
+			AccessGrant: minimalAccess,
+			Public:      false,
+		},
+	)
+	require.NoError(t, err)
+}
diff --git a/pkg/auth/peer.go b/pkg/auth/peer.go
index 35e15eb5..02dd46b6 100644
--- a/pkg/auth/peer.go
+++ b/pkg/auth/peer.go
@@ -156,7 +156,7 @@ func New(ctx context.Context, log *zap.Logger, config Config, configDir string)
 		Handler: handler,
 	}
 
-	drpcServer := drpcauth.NewServer(log, adb, endpoint)
+	drpcServer := drpcauth.NewServer(log, adb, endpoint, config.POSTSizeLimit)
 
 	return &Peer{
 		log:        log,