From 66b6fae01798c398b56b84a8ecaeccb7ebe68212 Mon Sep 17 00:00:00 2001
From: crschnick <crschnick@xpipe.io>
Date: Thu, 21 Nov 2024 11:23:36 +0000
Subject: [PATCH] Rework shell license check

---
 app/src/main/java/io/xpipe/app/ext/ShellStore.java     |  3 ---
 .../main/java/io/xpipe/app/storage/DataStorage.java    |  2 +-
 app/src/main/java/io/xpipe/app/util/LocalShell.java    |  1 -
 app/src/main/java/io/xpipe/app/util/ScanDialog.java    |  2 +-
 .../main/java/io/xpipe/core/process/ShellControl.java  |  4 ++--
 .../io/xpipe/core/process/WrapperShellControl.java     | 10 +++++-----
 .../java/io/xpipe/core/store/ConnectionFileSystem.java |  2 ++
 7 files changed, 11 insertions(+), 13 deletions(-)

diff --git a/app/src/main/java/io/xpipe/app/ext/ShellStore.java b/app/src/main/java/io/xpipe/app/ext/ShellStore.java
index 18d00feb7..3dc2961cb 100644
--- a/app/src/main/java/io/xpipe/app/ext/ShellStore.java
+++ b/app/src/main/java/io/xpipe/app/ext/ShellStore.java
@@ -32,9 +32,6 @@ public interface ShellStore extends DataStore, FileSystemStore, ValidatableStore
     default ShellSession newSession() throws Exception {
         var func = shellFunction();
         var c = func.control();
-        if (!isInStorage()) {
-            c.withoutLicenseCheck();
-        }
         return new ShellSession(this, () -> c);
     }
 
diff --git a/app/src/main/java/io/xpipe/app/storage/DataStorage.java b/app/src/main/java/io/xpipe/app/storage/DataStorage.java
index 2fcaa9a7f..b628e15f1 100644
--- a/app/src/main/java/io/xpipe/app/storage/DataStorage.java
+++ b/app/src/main/java/io/xpipe/app/storage/DataStorage.java
@@ -974,6 +974,6 @@ public abstract class DataStorage {
     }
 
     public DataStoreEntry local() {
-        return getStoreEntryIfPresent(LOCAL_ID).orElse(null);
+        return getStoreEntryIfPresent(LOCAL_ID).orElseThrow(() -> new IllegalStateException("Missing local machine connection"));
     }
 }
diff --git a/app/src/main/java/io/xpipe/app/util/LocalShell.java b/app/src/main/java/io/xpipe/app/util/LocalShell.java
index 93604db8f..11a8f3579 100644
--- a/app/src/main/java/io/xpipe/app/util/LocalShell.java
+++ b/app/src/main/java/io/xpipe/app/util/LocalShell.java
@@ -43,7 +43,6 @@ public class LocalShell {
                 localPowershell = ProcessControlProvider.get()
                         .createLocalProcessControl(false)
                         .subShell(ShellDialects.POWERSHELL)
-                        .withoutLicenseCheck()
                         .start();
             } catch (ProcessOutputException ex) {
                 throw ProcessOutputException.withPrefix("Failed to start local powershell process", ex);
diff --git a/app/src/main/java/io/xpipe/app/util/ScanDialog.java b/app/src/main/java/io/xpipe/app/util/ScanDialog.java
index 9c0649527..17321ff06 100644
--- a/app/src/main/java/io/xpipe/app/util/ScanDialog.java
+++ b/app/src/main/java/io/xpipe/app/util/ScanDialog.java
@@ -143,7 +143,7 @@ class ScanDialog extends DialogComp {
 
         ThreadHelper.runFailableAsync(() -> {
             BooleanScope.executeExclusive(busy, () -> {
-                var sc = entry.get().getStore().getOrStartSession().withoutLicenseCheck();
+                var sc = entry.get().getStore().getOrStartSession();
                 var a = applicable.apply(entry.get().get(), sc);
                 Platform.runLater(() -> {
                     if (a == null) {
diff --git a/core/src/main/java/io/xpipe/core/process/ShellControl.java b/core/src/main/java/io/xpipe/core/process/ShellControl.java
index e2bbebf3d..7cf31819c 100644
--- a/core/src/main/java/io/xpipe/core/process/ShellControl.java
+++ b/core/src/main/java/io/xpipe/core/process/ShellControl.java
@@ -54,8 +54,6 @@ public interface ShellControl extends ProcessControl {
 
     ShellControl getMachineRootSession();
 
-    ShellControl withoutLicenseCheck();
-
     String getOsName();
 
     boolean isLicenseCheck();
@@ -106,6 +104,8 @@ public interface ShellControl extends ProcessControl {
 
     ShellControl withErrorFormatter(Function<String, String> formatter);
 
+    void checkLicenseOrThrow();
+
     String prepareIntermediateTerminalOpen(
             TerminalInitFunction content, TerminalInitScriptConfig config, WorkingDirectoryFunction workingDirectory)
             throws Exception;
diff --git a/core/src/main/java/io/xpipe/core/process/WrapperShellControl.java b/core/src/main/java/io/xpipe/core/process/WrapperShellControl.java
index 7b54eaf91..d6baf5e3b 100644
--- a/core/src/main/java/io/xpipe/core/process/WrapperShellControl.java
+++ b/core/src/main/java/io/xpipe/core/process/WrapperShellControl.java
@@ -106,11 +106,6 @@ public class WrapperShellControl implements ShellControl {
         return parent.getMachineRootSession();
     }
 
-    @Override
-    public ShellControl withoutLicenseCheck() {
-        return parent.withoutLicenseCheck();
-    }
-
     @Override
     public String getOsName() {
         return parent.getOsName();
@@ -262,6 +257,11 @@ public class WrapperShellControl implements ShellControl {
         return parent.withErrorFormatter(formatter);
     }
 
+    @Override
+    public void checkLicenseOrThrow() {
+        parent.checkLicenseOrThrow();
+    }
+
     @Override
     public String prepareIntermediateTerminalOpen(
             TerminalInitFunction content, TerminalInitScriptConfig config, WorkingDirectoryFunction workingDirectory)
diff --git a/core/src/main/java/io/xpipe/core/store/ConnectionFileSystem.java b/core/src/main/java/io/xpipe/core/store/ConnectionFileSystem.java
index fb97c273c..00bd1ecc6 100644
--- a/core/src/main/java/io/xpipe/core/store/ConnectionFileSystem.java
+++ b/core/src/main/java/io/xpipe/core/store/ConnectionFileSystem.java
@@ -50,6 +50,8 @@ public class ConnectionFileSystem implements FileSystem {
                     "Shell has a PTY allocated and as a result does not support file system operations");
         }
 
+        shellControl.checkLicenseOrThrow();
+
         return this;
     }