From 1f5e0c033415eaf26aea1eb1b89f773ec8a0ce02 Mon Sep 17 00:00:00 2001 From: Dane Everitt Date: Mon, 19 Oct 2020 21:07:07 -0700 Subject: [PATCH] Update build modification service and cover logic with test cases closes #2553 --- .../Servers/BuildModificationService.php | 145 ++++------- .../Servers/BuildModificationServiceTest.php | 241 ++++++++++++++++++ 2 files changed, 296 insertions(+), 90 deletions(-) create mode 100644 tests/Integration/Services/Servers/BuildModificationServiceTest.php diff --git a/app/Services/Servers/BuildModificationService.php b/app/Services/Servers/BuildModificationService.php index e56d8f9b4..c42264f8d 100644 --- a/app/Services/Servers/BuildModificationService.php +++ b/app/Services/Servers/BuildModificationService.php @@ -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, + ]); } } } diff --git a/tests/Integration/Services/Servers/BuildModificationServiceTest.php b/tests/Integration/Services/Servers/BuildModificationServiceTest.php new file mode 100644 index 000000000..ab3d6160f --- /dev/null +++ b/tests/Integration/Services/Servers/BuildModificationServiceTest.php @@ -0,0 +1,241 @@ +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); + } +}