Validate shard id in shard controller (#3776)

service/history/shard/controller_impl.go

@@ -33,6 +33,7 @@ import (
  "time"
 
  "go.opentelemetry.io/otel/trace"
+ "go.temporal.io/api/serviceerror"
 
  "go.temporal.io/server/api/historyservice/v1"
  "go.temporal.io/server/client"
@@ -58,6 +59,11 @@ const (
  shardControllerMembershipUpdateListenerName = "ShardController"
 )
 
+var (
+ invalidShardIdLowerBound = serviceerror.NewInvalidArgument("shard Id cannot be equal or lower than zero")
+ invalidShardIdUpperBound = serviceerror.NewInvalidArgument("shard Id cannot be larger than max shard count")
+)
+
 type (
  ControllerImpl struct {
  membershipUpdateCh chan *membership.ChangedEvent
@@ -211,6 +217,17 @@ func (c *ControllerImpl) CloseShardByID(shardID int32) {
  }
 }
 
+func (c *ControllerImpl) ShardIDs() []int32 {
+ c.RLock()
+ defer c.RUnlock()
+
+ ids := make([]int32, 0, len(c.historyShards))
+ for id := range c.historyShards {
+ ids = append(ids, id)
+ }
+ return ids
+}
+
 func (c *ControllerImpl) shardClosedCallback(shard *ContextImpl) {
  startTime := time.Now().UTC()
  defer func() {
@@ -231,6 +248,10 @@ func (c *ControllerImpl) shardClosedCallback(shard *ContextImpl) {
 // if necessary. If a shard context is created, it will initialize in the background.
 // This function won't block on rangeid lease acquisition.
 func (c *ControllerImpl) getOrCreateShardContext(shardID int32) (*ContextImpl, error) {
+ err := c.validateShardId(shardID)
+ if err != nil {
+ return nil, err
+ }
  c.RLock()
  if shard, ok := c.historyShards[shardID]; ok {
  if shard.isValid() {
@@ -443,15 +464,14 @@ func (c *ControllerImpl) doShutdown() {
  c.historyShards = nil
 }
 
-func (c *ControllerImpl) ShardIDs() []int32 {
- c.RLock()
- defer c.RUnlock()
-
- ids := make([]int32, 0, len(c.historyShards))
- for id := range c.historyShards {
- ids = append(ids, id)
+func (c *ControllerImpl) validateShardId(shardID int32) error {
+ if shardID <= 0 {
+ return invalidShardIdLowerBound
  }
- return ids
+ if shardID > c.config.NumberOfShards {
+ return invalidShardIdUpperBound
+ }
+ return nil
 }
 
 func IsShardOwnershipLostError(err error) bool {

service/history/shard/controller_test.go

@@ -594,9 +594,9 @@ func (s *controllerSuite) TestShardExplicitUnload() {
  s.mockClusterMetadata.EXPECT().GetAllClusterInfo().Return(cluster.TestSingleDCClusterInfo).AnyTimes()
  mockEngine := NewMockEngine(s.controller)
  mockEngine.EXPECT().Stop().AnyTimes()
- s.setupMocksForAcquireShard(0, mockEngine, 5, 6, false)
+ s.setupMocksForAcquireShard(1, mockEngine, 5, 6, false)
 
- shard, err := s.shardController.getOrCreateShardContext(0)
+ shard, err := s.shardController.getOrCreateShardContext(1)
  s.NoError(err)
  s.Equal(1, len(s.shardController.ShardIDs()))
 
@@ -618,7 +618,7 @@ func (s *controllerSuite) TestShardExplicitUnloadCancelGetOrCreate() {
  mockEngine := NewMockEngine(s.controller)
  mockEngine.EXPECT().Stop().AnyTimes()
 
- shardID := int32(0)
+ shardID := int32(1)
  s.mockServiceResolver.EXPECT().Lookup(convert.Int32ToString(shardID)).Return(s.hostInfo, nil)
 
  ready := make(chan struct{})
@@ -638,7 +638,7 @@ func (s *controllerSuite) TestShardExplicitUnloadCancelGetOrCreate() {
  })
 
  // get shard, will start initializing in background
- shard, err := s.shardController.getOrCreateShardContext(0)
+ shard, err := s.shardController.getOrCreateShardContext(1)
  s.NoError(err)
 
  <-ready
@@ -659,7 +659,7 @@ func (s *controllerSuite) TestShardExplicitUnloadCancelAcquire() {
  mockEngine := NewMockEngine(s.controller)
  mockEngine.EXPECT().Stop().AnyTimes()
 
- shardID := int32(0)
+ shardID := int32(1)
  s.mockServiceResolver.EXPECT().Lookup(convert.Int32ToString(shardID)).Return(s.hostInfo, nil)
  // return success from GetOrCreateShard
  s.mockShardManager.EXPECT().GetOrCreateShard(gomock.Any(), getOrCreateShardRequestMatcher(shardID)).Return(
@@ -691,7 +691,7 @@ func (s *controllerSuite) TestShardExplicitUnloadCancelAcquire() {
  })
 
  // get shard, will start initializing in background
- shard, err := s.shardController.getOrCreateShardContext(0)
+ shard, err := s.shardController.getOrCreateShardContext(1)
  s.NoError(err)
 
  <-ready
@@ -834,6 +834,17 @@ func (s *controllerSuite) TestShardControllerFuzz() {
  }, 1*time.Second, 50*time.Millisecond, "engine start/stop")
 }
 
+func (s *controllerSuite) Test_GetOrCreateShard_InvalidShardID() {
+ numShards := int32(2)
+ s.config.NumberOfShards = numShards
+
+ _, err := s.shardController.getOrCreateShardContext(0)
+ s.ErrorIs(err, invalidShardIdLowerBound)
+
+ _, err = s.shardController.getOrCreateShardContext(3)
+ s.ErrorIs(err, invalidShardIdUpperBound)
+}
+
 func (s *controllerSuite) setupMocksForAcquireShard(
  shardID int32,
  mockEngine *MockEngine,