watchtower/wtmock: validate key index reservation

watchtower/wtdb/client_session.go

@@ -29,6 +29,15 @@ var (
 	// LastApplied value greater than any allocated sequence number.
 	ErrUnallocatedLastApplied = errors.New("tower echoed last appiled " +
 		"greater than allocated seqnum")
+
+	// ErrNoReservedKeyIndex signals that a client session could not be
+	// created because no session key index was reserved.
+	ErrNoReservedKeyIndex = errors.New("key index not reserved")
+
+	// ErrIncorrectKeyIndex signals that the client session could not be
+	// created because session key index differs from the reserved key
+	// index.
+	ErrIncorrectKeyIndex = errors.New("incorrect key index")
 )
 
 // ClientSession encapsulates a SessionInfo returned from a successful

watchtower/wtmock/client_db.go

@@ -106,6 +106,21 @@ func (m *ClientDB) CreateClientSession(session *wtdb.ClientSession) error {
 	m.mu.Lock()
 	defer m.mu.Unlock()
 
+	// Ensure that a session key index has been reserved for this tower.
+	keyIndex, ok := m.indexes[session.TowerID]
+	if !ok {
+		return wtdb.ErrNoReservedKeyIndex
+	}
+
+	// Ensure that the session's index matches the reserved index.
+	if keyIndex != session.KeyIndex {
+		return wtdb.ErrIncorrectKeyIndex
+	}
+
+	// Remove the key index reservation for this tower. Once committed, this
+	// permits us to create another session with this tower.
+	delete(m.indexes, session.TowerID)
+
 	m.activeSessions[session.ID] = &wtdb.ClientSession{
 		TowerID:          session.TowerID,
 		KeyIndex:         session.KeyIndex,