Update build modification service and cover logic with test cases

closes #2553

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,
+        DaemonServerRepository $daemonServerRepository
-        ServerRepositoryInterface $repository
     ) {
-        $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([
+                Allocation::query()->where('id', $data['allocation_id'])->where('server_id', $server->id)->firstOrFail();
-                    ['id', '=', $data['allocation_id']],
+            } catch (ModelNotFoundException $ex) {
-                    ['server_id', '=', $server->id],
+                throw new DisplayException('The requested default allocation is not currently assigned to this server.');
-                ]);
-            } catch (RecordNotFoundException $ex) {
-                throw new DisplayException(trans('admin/server.exceptions.default_allocation_not_found'));
             }
         }
 
-        /* @var \Pterodactyl\Models\Server $server */
+        // If any of these values are passed through in the data array go ahead and set
-        $server = $this->repository->withFreshModel()->update($server->id, [
+        // them correctly on the server model.
-            'oom_disabled' => array_get($data, 'oom_disabled'),
+        $merge = Arr::only($data, ['oom_disabled', 'memory', 'swap', 'io', 'cpu', 'threads', 'disk', 'allocation_id']);
-            'memory' => array_get($data, 'memory'),
+
-            'swap' => array_get($data, 'swap'),
+        $server->forceFill(array_merge($merge, [
-            'io' => array_get($data, 'io'),
+            'database_limit' => Arr::get($data, 'database_limit', 0) ?? null,
-            'cpu' => array_get($data, 'cpu'),
+            'allocation_limit' => Arr::get($data, 'allocation_limit', 0) ?? null,
-            'threads' => array_get($data, 'threads'),
+            'backup_limit' => Arr::get($data, 'backup_limit', 0) ?? 0,
-            'disk' => array_get($data, 'disk'),
+        ]))->saveOrFail();
-            'allocation_id' => array_get($data, 'allocation_id'),
+
-            'database_limit' => array_get($data, 'database_limit', 0) ?? null,
+        $server = $server->fresh();
-            'allocation_limit' => array_get($data, 'allocation_limit', 0) ?? null,
-            'backup_limit' => array_get($data, 'backup_limit', 0) ?? null,
-        ]);
 
         $updateData = $this->structureService->handle($server);
 
-        try {
+        $this->daemonServerRepository->setServer($server)->update($updateData['build'] ?? []);
-            $this->daemonServerRepository
-                ->setServer($server)
-                ->update(Arr::only($updateData, ['build']));
 
-            $this->connection->commit();
+        $this->connection->commit();
-        } catch (RequestException $exception) {
-            throw new DaemonConnectionException($exception);
-        }
 
         return $server;
     }
 
     /**
-     * Process the allocations being assigned in the data and ensure they
+     * Process the allocations being assigned in the data and ensure they are available for a server.
-     * 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 (empty($data['add_allocations']) && empty($data['remove_allocations'])) {
-
-        if (! array_key_exists('add_allocations', $data) && ! array_key_exists('remove_allocations', $data)) {
             return;
         }
 
-        // Handle the addition of allocations to this server.
+        // Handle the addition of allocations to this server. Only assign allocations that are not currently
-        if (array_key_exists('add_allocations', $data) && ! empty($data['add_allocations'])) {
+        // assigned to a different server, and only allocations on the same node as the server.
-            $unassigned = $this->allocationRepository->getUnassignedAllocationIds($server->node_id);
+        if (! empty($data['add_allocations'])) {
+            $query = Allocation::query()
+                ->where('node_id', $server->node_id)
+                ->whereIn('id', $data['add_allocations'])
+                ->whereNull('server_id');
 
-            $updateIds = [];
+            // Keep track of all the allocations we're just now adding so that we can use the first
-            foreach ($data['add_allocations'] as $allocation) {
+            // one to reset the default allocation to.
-                if (! in_array($allocation, $unassigned)) {
+            $freshlyAllocated = $query->pluck('id')->first();
-                    continue;
-                }
 
-                $firstAllocationId = $firstAllocationId ?? $allocation;
+            $query->update(['server_id' => $server->id]);
-                $updateIds[] = $allocation;
-            }
-
-            if (! empty($updateIds)) {
-                $this->allocationRepository->updateWhereIn('id', $updateIds, ['server_id' => $server->id]);
-            }
         }
 
-        // Handle removal of allocations from this server.
+        if (! empty($data['remove_allocations'])) {
-        if (array_key_exists('remove_allocations', $data) && ! empty($data['remove_allocations'])) {
-            $assigned = $server->allocations->pluck('id')->toArray();
-
-            $updateIds = [];
             foreach ($data['remove_allocations'] as $allocation) {
-                if (! in_array($allocation, $assigned)) {
+                // If we are attempting to remove the default allocation for the server, see if we can reassign
-                    continue;
+                // 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 ($allocation == $data['allocation_id']) {
+                    if (empty($freshlyAllocated)) {
-                    if (is_null($firstAllocationId)) {
+                        throw new DisplayException(
-                        throw new DisplayException(trans('admin/server.exceptions.no_new_default_allocation'));
+                            '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)) {
+            // Remove any of the allocations we got that are currently assigned to this server on
-                $this->allocationRepository->updateWhereIn('id', $updateIds, ['server_id' => null]);
+            // 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,
+                ]);
         }
     }
 }

tests/Integration/Services/Servers/BuildModificationServiceTest.php

@@ -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);
+    }
+}