From 40e59ab6a0c9defcd8a2ac9a0e95824ec7935acf Mon Sep 17 00:00:00 2001
From: crschnick <crschnick@xpipe.io>
Date: Fri, 9 Aug 2024 07:00:29 +0000
Subject: [PATCH] Fix shutdown requests

---
 .../main/java/io/xpipe/app/beacon/BeaconRequestHandler.java  | 3 ++-
 beacon/src/main/java/io/xpipe/beacon/BeaconInterface.java    | 4 ++++
 .../src/main/java/io/xpipe/beacon/api/AskpassExchange.java   | 5 +++++
 .../src/main/java/io/xpipe/beacon/api/HandshakeExchange.java | 5 +++++
 4 files changed, 16 insertions(+), 1 deletion(-)

diff --git a/app/src/main/java/io/xpipe/app/beacon/BeaconRequestHandler.java b/app/src/main/java/io/xpipe/app/beacon/BeaconRequestHandler.java
index 2d6f98b11..019b70096 100644
--- a/app/src/main/java/io/xpipe/app/beacon/BeaconRequestHandler.java
+++ b/app/src/main/java/io/xpipe/app/beacon/BeaconRequestHandler.java
@@ -28,7 +28,8 @@ public class BeaconRequestHandler<T> implements HttpHandler {
 
     @Override
     public void handle(HttpExchange exchange) {
-        if (OperationMode.isInShutdown()) {
+        if (OperationMode.isInShutdown() && !beaconInterface.acceptInShutdown()) {
+            writeError(exchange, new BeaconClientErrorResponse("Daemon is currently in shutdown"), 400);
             return;
         }
 
diff --git a/beacon/src/main/java/io/xpipe/beacon/BeaconInterface.java b/beacon/src/main/java/io/xpipe/beacon/BeaconInterface.java
index 0e28a2f03..91742971c 100644
--- a/beacon/src/main/java/io/xpipe/beacon/BeaconInterface.java
+++ b/beacon/src/main/java/io/xpipe/beacon/BeaconInterface.java
@@ -62,6 +62,10 @@ public abstract class BeaconInterface<T> {
         return (Class<T>) Class.forName(name);
     }
 
+    public boolean acceptInShutdown() {
+        return false;
+    }
+
     public boolean requiresCompletedStartup() {
         return true;
     }
diff --git a/beacon/src/main/java/io/xpipe/beacon/api/AskpassExchange.java b/beacon/src/main/java/io/xpipe/beacon/api/AskpassExchange.java
index a2674cc76..eb7587ef1 100644
--- a/beacon/src/main/java/io/xpipe/beacon/api/AskpassExchange.java
+++ b/beacon/src/main/java/io/xpipe/beacon/api/AskpassExchange.java
@@ -12,6 +12,11 @@ import java.util.UUID;
 
 public class AskpassExchange extends BeaconInterface<AskpassExchange.Request> {
 
+    @Override
+    public boolean acceptInShutdown() {
+        return true;
+    }
+
     @Override
     public String getPath() {
         return "/askpass";
diff --git a/beacon/src/main/java/io/xpipe/beacon/api/HandshakeExchange.java b/beacon/src/main/java/io/xpipe/beacon/api/HandshakeExchange.java
index ccd7cd38b..3ccadce8a 100644
--- a/beacon/src/main/java/io/xpipe/beacon/api/HandshakeExchange.java
+++ b/beacon/src/main/java/io/xpipe/beacon/api/HandshakeExchange.java
@@ -11,6 +11,11 @@ import lombok.extern.jackson.Jacksonized;
 
 public class HandshakeExchange extends BeaconInterface<HandshakeExchange.Request> {
 
+    @Override
+    public boolean acceptInShutdown() {
+        return true;
+    }
+
     @Override
     public String getPath() {
         return "/handshake";