Update build modification service and cover logic with test cases

closes #2553
This commit is contained in:
Dane Everitt 2020-10-19 21:07:07 -07:00
parent b3598b3b98
commit 1f5e0c0334
No known key found for this signature in database
GPG key ID: EEA66103B3D71F53
2 changed files with 296 additions and 90 deletions

View file

@ -4,22 +4,14 @@ namespace Pterodactyl\Services\Servers;
use Illuminate\Support\Arr;
use Pterodactyl\Models\Server;
use GuzzleHttp\Exception\RequestException;
use Pterodactyl\Models\Allocation;
use Illuminate\Database\ConnectionInterface;
use Pterodactyl\Exceptions\DisplayException;
use Illuminate\Database\Eloquent\ModelNotFoundException;
use Pterodactyl\Repositories\Wings\DaemonServerRepository;
use Pterodactyl\Exceptions\Repository\RecordNotFoundException;
use Pterodactyl\Contracts\Repository\ServerRepositoryInterface;
use Pterodactyl\Contracts\Repository\AllocationRepositoryInterface;
use Pterodactyl\Exceptions\Http\Connection\DaemonConnectionException;
class BuildModificationService
{
/**
* @var \Pterodactyl\Contracts\Repository\AllocationRepositoryInterface
*/
private $allocationRepository;
/**
* @var \Illuminate\Database\ConnectionInterface
*/
@ -30,11 +22,6 @@ class BuildModificationService
*/
private $daemonServerRepository;
/**
* @var \Pterodactyl\Contracts\Repository\ServerRepositoryInterface
*/
private $repository;
/**
* @var \Pterodactyl\Services\Servers\ServerConfigurationStructureService
*/
@ -43,23 +30,17 @@ class BuildModificationService
/**
* BuildModificationService constructor.
*
* @param \Pterodactyl\Contracts\Repository\AllocationRepositoryInterface $allocationRepository
* @param \Pterodactyl\Services\Servers\ServerConfigurationStructureService $structureService
* @param \Illuminate\Database\ConnectionInterface $connection
* @param \Pterodactyl\Repositories\Wings\DaemonServerRepository $daemonServerRepository
* @param \Pterodactyl\Contracts\Repository\ServerRepositoryInterface $repository
*/
public function __construct(
AllocationRepositoryInterface $allocationRepository,
ServerConfigurationStructureService $structureService,
ConnectionInterface $connection,
DaemonServerRepository $daemonServerRepository,
ServerRepositoryInterface $repository
DaemonServerRepository $daemonServerRepository
) {
$this->allocationRepository = $allocationRepository;
$this->daemonServerRepository = $daemonServerRepository;
$this->connection = $connection;
$this->repository = $repository;
$this->structureService = $structureService;
}
@ -70,9 +51,8 @@ class BuildModificationService
* @param array $data
* @return \Pterodactyl\Models\Server
*
* @throws \Throwable
* @throws \Pterodactyl\Exceptions\DisplayException
* @throws \Pterodactyl\Exceptions\Repository\RecordNotFoundException
* @throws \Pterodactyl\Exceptions\Model\DataValidationException
*/
public function handle(Server $server, array $data)
{
@ -82,48 +62,35 @@ class BuildModificationService
if (isset($data['allocation_id']) && $data['allocation_id'] != $server->allocation_id) {
try {
$this->allocationRepository->findFirstWhere([
['id', '=', $data['allocation_id']],
['server_id', '=', $server->id],
]);
} catch (RecordNotFoundException $ex) {
throw new DisplayException(trans('admin/server.exceptions.default_allocation_not_found'));
Allocation::query()->where('id', $data['allocation_id'])->where('server_id', $server->id)->firstOrFail();
} catch (ModelNotFoundException $ex) {
throw new DisplayException('The requested default allocation is not currently assigned to this server.');
}
}
/* @var \Pterodactyl\Models\Server $server */
$server = $this->repository->withFreshModel()->update($server->id, [
'oom_disabled' => array_get($data, 'oom_disabled'),
'memory' => array_get($data, 'memory'),
'swap' => array_get($data, 'swap'),
'io' => array_get($data, 'io'),
'cpu' => array_get($data, 'cpu'),
'threads' => array_get($data, 'threads'),
'disk' => array_get($data, 'disk'),
'allocation_id' => array_get($data, 'allocation_id'),
'database_limit' => array_get($data, 'database_limit', 0) ?? null,
'allocation_limit' => array_get($data, 'allocation_limit', 0) ?? null,
'backup_limit' => array_get($data, 'backup_limit', 0) ?? null,
]);
// If any of these values are passed through in the data array go ahead and set
// them correctly on the server model.
$merge = Arr::only($data, ['oom_disabled', 'memory', 'swap', 'io', 'cpu', 'threads', 'disk', 'allocation_id']);
$server->forceFill(array_merge($merge, [
'database_limit' => Arr::get($data, 'database_limit', 0) ?? null,
'allocation_limit' => Arr::get($data, 'allocation_limit', 0) ?? null,
'backup_limit' => Arr::get($data, 'backup_limit', 0) ?? 0,
]))->saveOrFail();
$server = $server->fresh();
$updateData = $this->structureService->handle($server);
try {
$this->daemonServerRepository
->setServer($server)
->update(Arr::only($updateData, ['build']));
$this->daemonServerRepository->setServer($server)->update($updateData['build'] ?? []);
$this->connection->commit();
} catch (RequestException $exception) {
throw new DaemonConnectionException($exception);
}
$this->connection->commit();
return $server;
}
/**
* Process the allocations being assigned in the data and ensure they
* are available for a server.
* Process the allocations being assigned in the data and ensure they are available for a server.
*
* @param \Pterodactyl\Models\Server $server
* @param array $data
@ -132,55 +99,53 @@ class BuildModificationService
*/
private function processAllocations(Server $server, array &$data)
{
$firstAllocationId = null;
if (! array_key_exists('add_allocations', $data) && ! array_key_exists('remove_allocations', $data)) {
if (empty($data['add_allocations']) && empty($data['remove_allocations'])) {
return;
}
// Handle the addition of allocations to this server.
if (array_key_exists('add_allocations', $data) && ! empty($data['add_allocations'])) {
$unassigned = $this->allocationRepository->getUnassignedAllocationIds($server->node_id);
// Handle the addition of allocations to this server. Only assign allocations that are not currently
// assigned to a different server, and only allocations on the same node as the server.
if (! empty($data['add_allocations'])) {
$query = Allocation::query()
->where('node_id', $server->node_id)
->whereIn('id', $data['add_allocations'])
->whereNull('server_id');
$updateIds = [];
foreach ($data['add_allocations'] as $allocation) {
if (! in_array($allocation, $unassigned)) {
continue;
}
// Keep track of all the allocations we're just now adding so that we can use the first
// one to reset the default allocation to.
$freshlyAllocated = $query->pluck('id')->first();
$firstAllocationId = $firstAllocationId ?? $allocation;
$updateIds[] = $allocation;
}
if (! empty($updateIds)) {
$this->allocationRepository->updateWhereIn('id', $updateIds, ['server_id' => $server->id]);
}
$query->update(['server_id' => $server->id]);
}
// Handle removal of allocations from this server.
if (array_key_exists('remove_allocations', $data) && ! empty($data['remove_allocations'])) {
$assigned = $server->allocations->pluck('id')->toArray();
$updateIds = [];
if (! empty($data['remove_allocations'])) {
foreach ($data['remove_allocations'] as $allocation) {
if (! in_array($allocation, $assigned)) {
continue;
}
if ($allocation == $data['allocation_id']) {
if (is_null($firstAllocationId)) {
throw new DisplayException(trans('admin/server.exceptions.no_new_default_allocation'));
// If we are attempting to remove the default allocation for the server, see if we can reassign
// to the first provided value in add_allocations. If there is no new first allocation then we
// will throw an exception back.
if ($allocation === ($data['allocation_id'] ?? $server->allocation_id)) {
if (empty($freshlyAllocated)) {
throw new DisplayException(
'You are attempting to delete the default allocation for this server but there is no fallback allocation to use.'
);
}
$data['allocation_id'] = $firstAllocationId;
// Update the default allocation to be the first allocation that we are creating.
$data['allocation_id'] = $freshlyAllocated;
}
$updateIds[] = $allocation;
}
if (! empty($updateIds)) {
$this->allocationRepository->updateWhereIn('id', $updateIds, ['server_id' => null]);
}
// Remove any of the allocations we got that are currently assigned to this server on
// this node. Also set the notes to null, otherwise when re-allocated to a new server those
// notes will be carried over.
Allocation::query()->where('node_id', $server->node_id)
->where('server_id', $server->id)
// Only remove the allocations that we didn't also attempt to add to the server...
->whereIn('id', array_diff($data['remove_allocations'], $data['add_allocations'] ?? []))
->update([
'notes' => null,
'server_id' => null,
]);
}
}
}

View file

@ -0,0 +1,241 @@
<?php
namespace Pterodactyl\Tests\Integration\Services\Servers;
use Mockery;
use Exception;
use Pterodactyl\Models\Server;
use Pterodactyl\Models\Allocation;
use Pterodactyl\Exceptions\DisplayException;
use GuzzleHttp\Exception\BadResponseException;
use Pterodactyl\Tests\Integration\IntegrationTestCase;
use Pterodactyl\Repositories\Wings\DaemonServerRepository;
use Pterodactyl\Services\Servers\BuildModificationService;
use Pterodactyl\Exceptions\Http\Connection\DaemonConnectionException;
class BuildModificationServiceTest extends IntegrationTestCase
{
/** @var \Mockery\MockInterface */
private $daemonServerRepository;
/**
* Setup tests.
*/
public function setUp(): void
{
parent::setUp();
$this->daemonServerRepository = Mockery::mock(DaemonServerRepository::class);
$this->swap(DaemonServerRepository::class, $this->daemonServerRepository);
}
/**
* Test that allocations can be added and removed from a server. Only the allocations on the
* current node and belonging to this server should be modified.
*/
public function testAllocationsCanBeModifiedForTheServer()
{
$server = $this->createServerModel();
$server2 = $this->createServerModel();
/** @var \Pterodactyl\Models\Allocation[] $allocations */
$allocations = factory(Allocation::class)->times(4)->create(['node_id' => $server->node_id]);
$initialAllocationId = $server->allocation_id;
$allocations[0]->update(['server_id' => $server->id, 'notes' => 'Test notes']);
// Some additional test allocations for the other server, not the server we are attempting
// to modify.
$allocations[2]->update(['server_id' => $server2->id]);
$allocations[3]->update(['server_id' => $server2->id]);
$this->daemonServerRepository->expects('setServer->update')->andReturnUndefined();
$response = $this->getService()->handle($server, [
// Attempt to add one new allocation, and an allocation assigned to another server. The
// other server allocation should be ignored, and only the allocation for this server should
// be used.
'add_allocations' => [$allocations[2]->id, $allocations[1]->id],
// Remove the default server allocation, ensuring that the new allocation passed through
// in the data becomes the default allocation.
'remove_allocations' => [$server->allocation_id, $allocations[0]->id, $allocations[3]->id],
]);
$this->assertInstanceOf(Server::class, $response);
// Only one allocation should exist for this server now.
$this->assertCount(1, $response->allocations);
$this->assertSame($allocations[1]->id, $response->allocation_id);
// These two allocations should not have been touched.
$this->assertDatabaseHas('allocations', ['id' => $allocations[2]->id, 'server_id' => $server2->id]);
$this->assertDatabaseHas('allocations', ['id' => $allocations[3]->id, 'server_id' => $server2->id]);
// Both of these allocations should have been removed from the server, and have had their
// notes properly reset.
$this->assertDatabaseHas('allocations', ['id' => $initialAllocationId, 'server_id' => null, 'notes' => null]);
$this->assertDatabaseHas('allocations', ['id' => $allocations[0]->id, 'server_id' => null, 'notes' => null]);
}
/**
* Test that an exception is thrown if removing the default allocation without also assigning
* new allocations to the server.
*/
public function testExceptionIsThrownIfRemovingTheDefaultAllocation()
{
$server = $this->createServerModel();
/** @var \Pterodactyl\Models\Allocation[] $allocations */
$allocations = factory(Allocation::class)->times(4)->create(['node_id' => $server->node_id]);
$allocations[0]->update(['server_id' => $server->id]);
$this->expectException(DisplayException::class);
$this->expectExceptionMessage('You are attempting to delete the default allocation for this server but there is no fallback allocation to use.');
$this->getService()->handle($server, [
'add_allocations' => [],
'remove_allocations' => [$server->allocation_id, $allocations[0]->id],
]);
}
/**
* Test that the build data for the server is properly passed along to the Wings instance so that
* the server data is updated in realtime. This test also ensures that only certain fields get updated
* for the server, and not just any arbitrary field.
*/
public function testServerBuildDataIsProperlyUpdatedOnWings()
{
$server = $this->createServerModel();
$this->daemonServerRepository->expects('setServer')->with(Mockery::on(function (Server $s) use ($server) {
return $s->id === $server->id;
}))->andReturnSelf();
$this->daemonServerRepository->expects('update')->with(Mockery::on(function ($data) {
$this->assertEquals([
'memory_limit' => 256,
'swap' => 128,
'io_weight' => 600,
'cpu_limit' => 150,
'threads' => '1,2',
'disk_space' => 1024,
], $data);
return true;
}))->andReturnUndefined();
$response = $this->getService()->handle($server, [
'oom_disabled' => false,
'memory' => 256,
'swap' => 128,
'io' => 600,
'cpu' => 150,
'threads' => '1,2',
'disk' => 1024,
'backup_limit' => null,
'database_limit' => 10,
'allocation_limit' => 20,
]);
$this->assertFalse($response->oom_disabled);
$this->assertSame(256, $response->memory);
$this->assertSame(128, $response->swap);
$this->assertSame(600, $response->io);
$this->assertSame(150, $response->cpu);
$this->assertSame('1,2', $response->threads);
$this->assertSame(1024, $response->disk);
$this->assertSame(0, $response->backup_limit);
$this->assertSame(10, $response->database_limit);
$this->assertSame(20, $response->allocation_limit);
}
/**
* Test that no exception is thrown if we are only removing an allocation.
*/
public function testNoExceptionIsThrownIfOnlyRemovingAllocation()
{
$server = $this->createServerModel();
/** @var \Pterodactyl\Models\Allocation[] $allocations */
$allocation = factory(Allocation::class)->create(['node_id' => $server->node_id, 'server_id' => $server->id]);
$this->daemonServerRepository->expects('setServer->update')->andReturnUndefined();
$this->getService()->handle($server, [
'remove_allocations' => [$allocation->id],
]);
$this->assertDatabaseHas('allocations', ['id' => $allocation->id, 'server_id' => null]);
}
/**
* Test that allocations in both the add and remove arrays are only added, and not removed.
* This scenario wouldn't really happen in the UI, but it is possible to perform via the API
* so we want to make sure that the logic being used doesn't break if the allocation exists
* in both arrays.
*
* We'll default to adding the allocation in this case.
*/
public function testAllocationInBothAddAndRemoveIsAdded()
{
$server = $this->createServerModel();
/** @var \Pterodactyl\Models\Allocation[] $allocations */
$allocation = factory(Allocation::class)->create(['node_id' => $server->node_id]);
$this->daemonServerRepository->expects('setServer->update')->andReturnUndefined();
$this->getService()->handle($server, [
'add_allocations' => [$allocation->id],
'remove_allocations' => [$allocation->id],
]);
$this->assertDatabaseHas('allocations', ['id' => $allocation->id, 'server_id' => $server->id]);
}
/**
* Test that using the same allocation ID multiple times in the array does not cause an error.
*/
public function testUsingSameAllocationIdMultipleTimesDoesNotError()
{
$server = $this->createServerModel();
/** @var \Pterodactyl\Models\Allocation[] $allocations */
$allocation = factory(Allocation::class)->create(['node_id' => $server->node_id, 'server_id' => $server->id]);
$allocation2 = factory(Allocation::class)->create(['node_id' => $server->node_id]);
$this->daemonServerRepository->expects('setServer->update')->andReturnUndefined();
$this->getService()->handle($server, [
'add_allocations' => [$allocation2->id, $allocation2->id],
'remove_allocations' => [$allocation->id, $allocation->id],
]);
$this->assertDatabaseHas('allocations', ['id' => $allocation->id, 'server_id' => null]);
$this->assertDatabaseHas('allocations', ['id' => $allocation2->id, 'server_id' => $server->id]);
}
/**
* Test that any changes we made to the server or allocations are rolled back if there is an
* exception while performing any action.
*/
public function testThatUpdatesAreRolledBackIfExceptionIsEncountered()
{
$server = $this->createServerModel();
/** @var \Pterodactyl\Models\Allocation[] $allocations */
$allocation = factory(Allocation::class)->create(['node_id' => $server->node_id]);
$this->daemonServerRepository->expects('setServer->update')->andThrows(new DisplayException('Test'));
$this->expectException(DisplayException::class);
$this->getService()->handle($server, ['add_allocations' => [$allocation->id]]);
$this->assertDatabaseHas('allocations', ['id' => $allocation->id, 'server_id' => null]);
}
/**
* @return \Pterodactyl\Services\Servers\BuildModificationService
*/
private function getService()
{
return $this->app->make(BuildModificationService::class);
}
}