burmatovan/openuds
1
0
Fork 0
forked from shaba/openuds
Code Issues Pull Requests Projects Releases Wiki Activity

Fixed service_multi, so, in case a phisical machines got "stuck on db" forever, could be removed by editing the service (remove, save, readd and done)

Browse Source
This commit is contained in:
Adolfo Gómez García 2020-07-22 07:10:30 +02:00
parent 8a49e7c437
commit 735a093b0b
2 changed files with 17 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

10
server/src/uds/core/util/storage.py
View File

@ -151,7 +151,7 @@ class Storage:
self._owner = owner.decode('utf-8') if isinstance(owner, bytes) else owner
self._bowner = self._owner.encode('utf8')
def __getKey(self, key: typing.Union[str, bytes]) -> str:
def getKey(self, key: typing.Union[str, bytes]) -> str:
return _calcKey(self._bowner, key.encode('utf8') if isinstance(key, str) else key)
def saveData(self, skey: typing.Union[str, bytes], data: typing.Any, attr1: typing.Optional[str] = None) -> None:
@ -160,7 +160,7 @@ class Storage:
self.remove(skey)
return
key = self.__getKey(skey)
key = self.getKey(skey)
if isinstance(data, str):
data = data.encode('utf-8')
data = encoders.encodeAsStr(data, 'base64')
@ -183,7 +183,7 @@ class Storage:
def readData(self, skey: typing.Union[str, bytes], fromPickle: bool = False) -> typing.Optional[typing.Union[str, bytes]]:
try:
key = self.__getKey(skey)
key = self.getKey(skey)
logger.debug('Accesing to %s %s', skey, key)
c: DBStorage = DBStorage.objects.get(pk=key) # @UndefinedVariable
val: bytes = typing.cast(bytes, encoders.decode(c.data, 'base64'))
@ -221,7 +221,7 @@ class Storage:
keys: typing.Iterable[typing.Union[str, bytes]] = [skey] if isinstance(skey, (str, bytes)) else skey
try:
# Process several keys at once
DBStorage.objects.filter(key__in=[self.__getKey(k) for k in keys]).delete()
DBStorage.objects.filter(key__in=[self.getKey(k) for k in keys]).delete()
except Exception:
pass
@ -249,7 +249,7 @@ class Storage:
for v in query:
yield typing.cast(bytes, encoders.decode(v.data, 'base64'))
def filter(self, attr1: typing.Optional[str], forUpdate: bool = False) -> typing.Iterable[typing.Tuple[str, bytes, str]]:
def filter(self, attr1: typing.Optional[str] = None, forUpdate: bool = False) -> typing.Iterable[typing.Tuple[str, bytes, str]]:
if attr1 is None:
query = DBStorage.objects.filter(owner=self._owner) # @UndefinedVariable
else:

12
server/src/uds/services/PhysicalMachines/service_multi.py
View File

@ -35,6 +35,7 @@ import logging
import typing
from django.utils.translation import ugettext_lazy as _
from django.db import transaction
from uds.core.ui import gui
from uds.core.util import log
@ -100,7 +101,18 @@ class IPMachinesService(IPServiceBase):
self._ips = []
else:
self._ips = ['{}~{}'.format(str(ip).strip(), i) for i, ip in enumerate(values['ipList']) if str(ip).strip()] # Allow duplicates right now
active = {v for v in values['ipList']}
# self._ips.sort()
# Remove non existing "locked" ips from storage now
skipKey = self.storage.getKey('ips')
with transaction.atomic():
for key, data, _ in self.storage.filter(forUpdate=True):
if key == skipKey: # Avoid "ips" key
continue
# If not in current active list of ips, remove it
if data.decode() not in active:
logger.info('IP %s locked but not in active list. Removed', data.decode())
self.storage.remove(data.decode())
self._token = self.token.value.strip()
self._port = self.port.value