Various reworks

beacon/src/main/java/io/xpipe/beacon/BeaconClient.java

@@ -1,7 +1,7 @@
 package io.xpipe.beacon;
 
 import com.fasterxml.jackson.core.JsonGenerator;
-import com.fasterxml.jackson.core.JsonParser;
+import com.fasterxml.jackson.core.util.DefaultPrettyPrinter;
 import com.fasterxml.jackson.databind.JsonNode;
 import com.fasterxml.jackson.databind.node.JsonNodeFactory;
 import com.fasterxml.jackson.databind.node.ObjectNode;
@@ -11,12 +11,11 @@ import io.xpipe.beacon.exchange.data.ClientErrorMessage;
 import io.xpipe.beacon.exchange.data.ServerErrorMessage;
 import io.xpipe.core.util.JacksonHelper;
 
-import java.io.IOException;
-import java.io.InputStream;
-import java.io.OutputStream;
+import java.io.*;
 import java.net.InetAddress;
 import java.net.Socket;
 import java.net.SocketException;
+import java.nio.charset.StandardCharsets;
 import java.util.Arrays;
 import java.util.Optional;
 
@@ -50,21 +49,28 @@ public class BeaconClient implements AutoCloseable {
         }
     }
 
-    private final Socket socket;
+    private final Closeable closeable;
     private final InputStream in;
     private final OutputStream out;
 
     public BeaconClient() throws IOException {
-        socket = new Socket(InetAddress.getLoopbackAddress(), BeaconConfig.getUsedPort());
+        var socket = new Socket(InetAddress.getLoopbackAddress(), BeaconConfig.getUsedPort());
+        closeable = socket;
         in = socket.getInputStream();
         out = socket.getOutputStream();
     }
 
+    public BeaconClient(Closeable closeable, InputStream in, OutputStream out) {
+        this.closeable = closeable;
+        this.in = in;
+        this.out = out;
+    }
+
     public void close() throws ConnectorException {
         try {
-            socket.close();
+            closeable.close();
         } catch (IOException ex) {
-            throw new ConnectorException("Couldn't close socket", ex);
+            throw new ConnectorException("Couldn't close client", ex);
         }
     }
 
@@ -74,7 +80,7 @@ public class BeaconClient implements AutoCloseable {
             if (sep.length != 0 && !Arrays.equals(BODY_SEPARATOR, sep)) {
                 throw new ConnectorException("Invalid body separator");
             }
-            return BeaconFormat.readBlocks(socket);
+            return BeaconFormat.readBlocks(in);
         } catch (IOException ex) {
             throw new ConnectorException(ex);
         }
@@ -83,13 +89,13 @@ public class BeaconClient implements AutoCloseable {
     public OutputStream sendBody() throws ConnectorException {
         try {
             out.write(BODY_SEPARATOR);
-            return BeaconFormat.writeBlocks(socket);
+            return BeaconFormat.writeBlocks(out);
         } catch (IOException ex) {
             throw new ConnectorException(ex);
         }
     }
 
-    public  <T extends RequestMessage> void sendRequest(T req) throws ClientException, ConnectorException {
+    public <T extends RequestMessage> void sendRequest(T req) throws ClientException, ConnectorException {
         ObjectNode json = JacksonHelper.newMapper().valueToTree(req);
         var prov = MessageExchanges.byRequest(req);
         if (prov.isEmpty()) {
@@ -106,25 +112,31 @@ public class BeaconClient implements AutoCloseable {
             System.out.println("Sending request to server of type " + req.getClass().getName());
         }
 
-        if (BeaconConfig.printMessages()) {
-            System.out.println("Sending raw request:");
-            System.out.println(msg.toPrettyString());
+        var writer = new StringWriter();
+        var mapper = JacksonHelper.newMapper();
+        try (JsonGenerator g = mapper.createGenerator(writer).setPrettyPrinter(new DefaultPrettyPrinter())) {
+            g.writeTree(msg);
+        } catch (IOException ex) {
+            throw new ConnectorException("Couldn't serialize request", ex);
         }
 
-        try {
-            var mapper = JacksonHelper.newMapper().disable(JsonGenerator.Feature.AUTO_CLOSE_TARGET);
-            var gen = mapper.createGenerator(socket.getOutputStream());
-            gen.writeTree(msg);
+        var content = writer.toString();
+        if (BeaconConfig.printMessages()) {
+            System.out.println("Sending raw request:");
+            System.out.println(content);
+        }
+
+        try (OutputStream blockOut = BeaconFormat.writeBlocks(out)) {
+            blockOut.write(content.getBytes(StandardCharsets.UTF_8));
         } catch (IOException ex) {
             throw new ConnectorException("Couldn't write to socket", ex);
         }
     }
 
     public <T extends ResponseMessage> T receiveResponse() throws ConnectorException, ClientException, ServerException {
-        JsonNode read;
-        try {
-            var in = socket.getInputStream();
-            read = JacksonHelper.newMapper().disable(JsonParser.Feature.AUTO_CLOSE_SOURCE).readTree(in);
+        JsonNode node;
+        try (InputStream blockIn = BeaconFormat.readBlocks(in)) {
+            node = JacksonHelper.newMapper().readTree(blockIn);
         } catch (SocketException ex) {
             throw new ConnectorException("Connection to xpipe daemon closed unexpectedly", ex);
         } catch (IOException ex) {
@@ -133,24 +145,24 @@ public class BeaconClient implements AutoCloseable {
 
         if (BeaconConfig.printMessages()) {
             System.out.println("Received response:");
-            System.out.println(read.toPrettyString());
+            System.out.println(node.toPrettyString());
         }
 
-        if (read.isMissingNode()) {
+        if (node.isMissingNode()) {
             throw new ConnectorException("Received unexpected EOF");
         }
 
-        var se = parseServerError(read);
+        var se = parseServerError(node);
         if (se.isPresent()) {
             se.get().throwError();
         }
 
-        var ce = parseClientError(read);
+        var ce = parseClientError(node);
         if (ce.isPresent()) {
             throw ce.get().throwException();
         }
 
-        return parseResponse(read);
+        return parseResponse(node);
     }
 
     private Optional<ClientErrorMessage> parseClientError(JsonNode node) throws ConnectorException {
@@ -206,4 +218,12 @@ public class BeaconClient implements AutoCloseable {
             throw new ConnectorException("Couldn't parse response", ex);
         }
     }
+
+    public InputStream getRawInputStream() {
+        return in;
+    }
+
+    public OutputStream getRawOutputStream() {
+        return out;
+    }
 }

beacon/src/main/java/io/xpipe/beacon/BeaconConnection.java

@@ -6,22 +6,27 @@ import java.io.OutputStream;
 
 public abstract class BeaconConnection implements AutoCloseable {
 
-    protected BeaconClient socket;
+
+    protected BeaconClient beaconClient;
 
     private InputStream bodyInput;
     private OutputStream bodyOutput;
 
     protected abstract void constructSocket();
 
+    public BeaconClient getBeaconClient() {
+        return beaconClient;
+    }
+
     @Override
     public void close() {
         try {
-            if (socket != null) {
-                socket.close();
+            if (beaconClient != null) {
+                beaconClient.close();
             }
-            socket = null;
+            beaconClient = null;
         } catch (Exception e) {
-            socket = null;
+            beaconClient = null;
             throw new BeaconException("Could not close beacon connection", e);
         }
     }
@@ -43,7 +48,7 @@ public abstract class BeaconConnection implements AutoCloseable {
     }
 
     public void checkClosed() {
-        if (socket == null) {
+        if (beaconClient == null) {
             throw new BeaconException("Socket is closed");
         }
     }
@@ -70,7 +75,8 @@ public abstract class BeaconConnection implements AutoCloseable {
 
     public <REQ extends RequestMessage, RES extends ResponseMessage> void performInputExchange(
             REQ req,
-            BeaconClient.FailableBiConsumer<RES, InputStream, Exception> responseConsumer) {
+            BeaconClient.FailableBiConsumer<RES, InputStream, Exception> responseConsumer
+    ) {
         checkClosed();
 
         performInputOutputExchange(req, null, responseConsumer);
@@ -79,33 +85,35 @@ public abstract class BeaconConnection implements AutoCloseable {
     public <REQ extends RequestMessage, RES extends ResponseMessage> void performInputOutputExchange(
             REQ req,
             BeaconClient.FailableConsumer<OutputStream, IOException> reqWriter,
-            BeaconClient.FailableBiConsumer<RES, InputStream, Exception> responseConsumer) {
+            BeaconClient.FailableBiConsumer<RES, InputStream, Exception> responseConsumer
+    ) {
         checkClosed();
 
         try {
-            socket.sendRequest(req);
+            beaconClient.sendRequest(req);
             if (reqWriter != null) {
-                try (var out = socket.sendBody()) {
+                try (var out = sendBody()) {
                     reqWriter.accept(out);
                 }
             }
-            RES res = socket.receiveResponse();
-            try (var in = socket.receiveBody()) {
+            RES res = beaconClient.receiveResponse();
+            try (var in = receiveBody()) {
                 responseConsumer.accept(res, in);
             }
         } catch (Exception e) {
-            throw new BeaconException("Could not communicate with beacon", e);
+            throw unwrapException(e);
         }
     }
 
     public <REQ extends RequestMessage> void sendRequest(
-            REQ req) {
+            REQ req
+    ) {
         checkClosed();
 
         try {
-            socket.sendRequest(req);
+            beaconClient.sendRequest(req);
         } catch (Exception e) {
-            throw new BeaconException("Could not communicate with beacon", e);
+            throw unwrapException(e);
         }
     }
 
@@ -113,9 +121,9 @@ public abstract class BeaconConnection implements AutoCloseable {
         checkClosed();
 
         try {
-            return socket.receiveResponse();
+            return beaconClient.receiveResponse();
         } catch (Exception e) {
-            throw new BeaconException("Could not communicate with beacon", e);
+            throw unwrapException(e);
         }
     }
 
@@ -123,10 +131,10 @@ public abstract class BeaconConnection implements AutoCloseable {
         checkClosed();
 
         try {
-            bodyOutput = socket.sendBody();
+            bodyOutput = beaconClient.sendBody();
             return bodyOutput;
         } catch (Exception e) {
-            throw new BeaconException("Could not communicate with beacon", e);
+            throw unwrapException(e);
         }
     }
 
@@ -134,38 +142,57 @@ public abstract class BeaconConnection implements AutoCloseable {
         checkClosed();
 
         try {
-            bodyInput = socket.receiveBody();
+            bodyInput = beaconClient.receiveBody();
             return bodyInput;
         } catch (Exception e) {
-            throw new BeaconException("Could not communicate with beacon", e);
+            throw unwrapException(e);
         }
     }
 
     public <REQ extends RequestMessage, RES extends ResponseMessage> RES performOutputExchange(
             REQ req,
-            BeaconClient.FailableConsumer<OutputStream, Exception> reqWriter) {
+            BeaconClient.FailableConsumer<OutputStream, Exception> reqWriter
+    ) {
         checkClosed();
 
         try {
-            socket.sendRequest(req);
-            try (var out = socket.sendBody()) {
+            beaconClient.sendRequest(req);
+            try (var out = sendBody()) {
                 reqWriter.accept(out);
             }
-            return socket.receiveResponse();
+            return beaconClient.receiveResponse();
         } catch (Exception e) {
-            throw new BeaconException("Could not communicate with beacon", e);
+            throw unwrapException(e);
         }
     }
 
     public <REQ extends RequestMessage, RES extends ResponseMessage> RES performSimpleExchange(
-            REQ req) {
+            REQ req
+    ) {
         checkClosed();
 
         try {
-            socket.sendRequest(req);
-            return socket.receiveResponse();
+            beaconClient.sendRequest(req);
+            return beaconClient.receiveResponse();
         } catch (Exception e) {
-            throw new BeaconException("Could not communicate with beacon", e);
+            throw unwrapException(e);
         }
     }
+
+    private BeaconException unwrapException(Exception exception) {
+        if (exception instanceof ServerException s) {
+            return new BeaconException("And internal server error occurred", s.getCause());
+        }
+
+        if (exception instanceof ClientException s) {
+            return new BeaconException("A client error occurred", s.getCause());
+        }
+
+        if (exception instanceof ConnectorException s) {
+            return new BeaconException("A beacon connection error occurred", s.getCause());
+        }
+
+
+        return new BeaconException("An unexpected error occurred", exception);
+    }
 }

beacon/src/main/java/io/xpipe/beacon/BeaconFormat.java

@@ -5,7 +5,6 @@ import lombok.experimental.UtilityClass;
 import java.io.IOException;
 import java.io.InputStream;
 import java.io.OutputStream;
-import java.net.Socket;
 import java.nio.ByteBuffer;
 
 @UtilityClass
@@ -13,8 +12,7 @@ public class BeaconFormat {
 
     private static final int SEGMENT_SIZE = 65536;
 
-    public static OutputStream writeBlocks(Socket socket) throws IOException {
-        var out = socket.getOutputStream();
+    public static OutputStream writeBlocks(OutputStream out) throws IOException {
         return new OutputStream() {
             private final byte[] currentBytes = new byte[SEGMENT_SIZE];
             private int index;
@@ -23,6 +21,7 @@ public class BeaconFormat {
             public void close() throws IOException {
                 finishBlock();
                 out.flush();
+                index = -1;
             }
 
             @Override
@@ -49,8 +48,7 @@ public class BeaconFormat {
         };
     }
 
-    public static InputStream readBlocks(Socket socket) throws IOException {
-        var in = socket.getInputStream();
+    public static InputStream readBlocks(InputStream in) throws IOException {
         return new InputStream() {
 
             private byte[] currentBytes;
@@ -60,7 +58,9 @@ public class BeaconFormat {
             @Override
             public int read() throws IOException {
                 if ((currentBytes == null || index == currentBytes.length) && !lastBlock) {
-                    readBlock();
+                    if (!readBlock()) {
+                        return -1;
+                    }
                 }
 
                 if (currentBytes != null && index == currentBytes.length && lastBlock) {
@@ -72,8 +72,12 @@ public class BeaconFormat {
                 return out;
             }
 
-            private void readBlock() throws IOException {
+            private boolean readBlock() throws IOException {
                 var length = in.readNBytes(4);
+                if (length.length < 4) {
+                    return false;
+                }
+
                 var lengthInt = ByteBuffer.wrap(length).getInt();
 
                 if (BeaconConfig.printMessages()) {
@@ -85,6 +89,7 @@ public class BeaconFormat {
                 if (lengthInt < SEGMENT_SIZE) {
                     lastBlock = true;
                 }
+                return true;
             }
         };
     }

beacon/src/main/java/io/xpipe/beacon/XPipeInstance.java

@@ -0,0 +1,17 @@
+package io.xpipe.beacon;
+
+import io.xpipe.core.store.ShellStore;
+import lombok.Value;
+
+import java.util.List;
+import java.util.Map;
+import java.util.UUID;
+
+@Value
+public class XPipeInstance {
+
+    UUID uuid;
+    String name;
+    Map<ShellStore, XPipeInstance> adjacent;
+    List<XPipeInstance> reachable;
+}

beacon/src/main/java/io/xpipe/beacon/exchange/cli/InstanceExchange.java

@@ -0,0 +1,33 @@
+package io.xpipe.beacon.exchange.cli;
+
+import io.xpipe.beacon.RequestMessage;
+import io.xpipe.beacon.ResponseMessage;
+import io.xpipe.beacon.XPipeInstance;
+import io.xpipe.beacon.exchange.MessageExchange;
+import lombok.Builder;
+import lombok.NonNull;
+import lombok.Value;
+import lombok.extern.jackson.Jacksonized;
+
+public class InstanceExchange implements MessageExchange {
+
+    @Override
+    public String getId() {
+        return "instance";
+    }
+
+    @Jacksonized
+    @Builder
+    @Value
+    public static class Request implements RequestMessage {
+
+    }
+
+    @Jacksonized
+    @Builder
+    @Value
+    public static class Response implements ResponseMessage {
+        @NonNull
+        XPipeInstance instance;
+    }
+}

beacon/src/main/java/module-info.java

@@ -23,6 +23,7 @@ module io.xpipe.beacon {
     uses MessageExchange;
     provides io.xpipe.beacon.exchange.MessageExchange with
             ForwardExchange,
+            InstanceExchange,
             EditStoreExchange,
             AddSourceExchange,
             StoreProviderListExchange,