diff --git a/app/src/main/java/io/xpipe/app/beacon/AppBeaconCache.java b/app/src/main/java/io/xpipe/app/beacon/AppBeaconCache.java index 9452e6c723..f68e2bfc38 100644 --- a/app/src/main/java/io/xpipe/app/beacon/AppBeaconCache.java +++ b/app/src/main/java/io/xpipe/app/beacon/AppBeaconCache.java @@ -6,46 +6,71 @@ import lombok.Value; -import java.util.HashSet; -import java.util.Set; +import java.util.Collection; +import java.util.Map; import java.util.UUID; +import java.util.concurrent.ConcurrentHashMap; @Value public class AppBeaconCache { - Set shellSessions = new HashSet<>(); + Map shellSessionsById = new ConcurrentHashMap<>(); + + public Collection getShellSessions() { + return shellSessionsById.values(); + } public BeaconShellSession getShellSession(UUID uuid) throws BeaconClientException { - var found = shellSessions.stream() - .filter(beaconShellSession -> - beaconShellSession.getEntry().getUuid().equals(uuid)) - .findFirst(); - if (found.isEmpty()) { + var session = shellSessionsById.get(uuid); + if (session == null) { throw new BeaconClientException("No active shell session known for id " + uuid); } - return found.get(); + return session; } public BeaconShellSession getOrStart(DataStoreEntryRef ref) throws Exception { - var existing = AppBeaconServer.get().getCache().getShellSessions().stream() - .filter(beaconShellSession -> beaconShellSession.getEntry().equals(ref.get())) - .findFirst(); - var control = (existing.isPresent() - ? existing.get().getControl() - : ref.getStore().standaloneControl().start()); - control.setNonInteractive(); - control.start(); - - var d = control.getShellDialect().getDumbMode(); - if (!d.supportsAnyPossibleInteraction()) { - control.close(); - d.throwIfUnsupported(); + var entry = ref.get(); + try { + var session = shellSessionsById.computeIfAbsent(entry.getUuid(), key -> { + try { + var control = ref.getStore().standaloneControl(); + control.setNonInteractive(); + control = control.start(); + + var d = control.getShellDialect().getDumbMode(); + if (!d.supportsAnyPossibleInteraction()) { + control.close(); + d.throwIfUnsupported(); + } + + return new BeaconShellSession(entry, control); + } catch (Exception ex) { + throw new ShellSessionInitException(ex); + } + }); + session.getControl().setNonInteractive(); + return session; + } catch (ShellSessionInitException ex) { + throw ex.getCause(); } + } - if (existing.isEmpty()) { - AppBeaconServer.get().getCache().getShellSessions().add(new BeaconShellSession(ref.get(), control)); + public void removeShellSession(UUID uuid) { + shellSessionsById.remove(uuid); + } + + private static final class ShellSessionInitException extends RuntimeException { + + private final Exception cause; + + private ShellSessionInitException(Exception cause) { + super(cause); + this.cause = cause; } - return new BeaconShellSession(ref.get(), control); + @Override + public synchronized Exception getCause() { + return cause; + } } } diff --git a/app/src/main/java/io/xpipe/app/beacon/impl/ShellStartExchangeImpl.java b/app/src/main/java/io/xpipe/app/beacon/impl/ShellStartExchangeImpl.java index 45410fa958..eb22f5ddfd 100644 --- a/app/src/main/java/io/xpipe/app/beacon/impl/ShellStartExchangeImpl.java +++ b/app/src/main/java/io/xpipe/app/beacon/impl/ShellStartExchangeImpl.java @@ -1,9 +1,9 @@ package io.xpipe.app.beacon.impl; import io.xpipe.app.beacon.AppBeaconServer; -import io.xpipe.app.beacon.BeaconShellSession; import io.xpipe.app.ext.ShellStore; import io.xpipe.app.storage.DataStorage; +import io.xpipe.app.storage.DataStoreEntryRef; import io.xpipe.beacon.BeaconClientException; import io.xpipe.beacon.api.ShellStartExchange; import io.xpipe.core.JacksonMapper; @@ -19,28 +19,13 @@ public Object handle(HttpExchange exchange, Request msg) { var e = DataStorage.get() .getStoreEntryIfPresent(msg.getConnection()) .orElseThrow(() -> new BeaconClientException("Unknown connection")); - if (!(e.getStore() instanceof ShellStore s)) { + if (!(e.getStore() instanceof ShellStore)) { throw new BeaconClientException("Not a shell connection"); } - var existing = AppBeaconServer.get().getCache().getShellSessions().stream() - .filter(beaconShellSession -> beaconShellSession.getEntry().equals(e)) - .findFirst(); - var control = (existing.isPresent() - ? existing.get().getControl() - : s.standaloneControl().start()); - control.setNonInteractive(); - control.start(); - - var d = control.getShellDialect().getDumbMode(); - if (!d.supportsAnyPossibleInteraction()) { - control.close(); - d.throwIfUnsupported(); - } - - if (existing.isEmpty()) { - AppBeaconServer.get().getCache().getShellSessions().add(new BeaconShellSession(e, control)); - } + var shellSession = + AppBeaconServer.get().getCache().getOrStart(new DataStoreEntryRef(e)); + var control = shellSession.getControl(); var ttyState = JacksonMapper.getDefault().valueToTree(control.getTtyState()).asText(); return Response.builder() .shellDialect(control.getShellDialect().getId()) diff --git a/app/src/main/java/io/xpipe/app/beacon/impl/ShellStopExchangeImpl.java b/app/src/main/java/io/xpipe/app/beacon/impl/ShellStopExchangeImpl.java index 5329275e6a..fc2463970e 100644 --- a/app/src/main/java/io/xpipe/app/beacon/impl/ShellStopExchangeImpl.java +++ b/app/src/main/java/io/xpipe/app/beacon/impl/ShellStopExchangeImpl.java @@ -13,7 +13,7 @@ public class ShellStopExchangeImpl extends ShellStopExchange { public Object handle(HttpExchange exchange, Request msg) { var e = AppBeaconServer.get().getCache().getShellSession(msg.getConnection()); e.getControl().close(); - AppBeaconServer.get().getCache().getShellSessions().remove(e); + AppBeaconServer.get().getCache().removeShellSession(msg.getConnection()); return Response.builder().build(); } }