Working code

15 files changed, 499 insertions, 0 deletions

diff --git a/src/main/java/moe/nea/libautoupdate/CurrentVersion.java b/src/main/java/moe/nea/libautoupdate/CurrentVersion.java
new file mode 100644
index 0000000..52b8372
--- /dev/null
+++ b/src/main/java/moe/nea/libautoupdate/CurrentVersion.java
@@ -0,0 +1,9 @@
+package moe.nea.libautoupdate;
+
+public interface CurrentVersion {
+    int getCurrentVersionNumber();
+
+    static CurrentVersion of(int number) {
+        return () -> number;
+    }
+}
diff --git a/src/main/java/moe/nea/libautoupdate/DeleteAndSaveInSameFolderUpdateTarget.java b/src/main/java/moe/nea/libautoupdate/DeleteAndSaveInSameFolderUpdateTarget.java
new file mode 100644
index 0000000..3e07f45
--- /dev/null
+++ b/src/main/java/moe/nea/libautoupdate/DeleteAndSaveInSameFolderUpdateTarget.java
@@ -0,0 +1,22 @@
+package moe.nea.libautoupdate;
+
+import lombok.SneakyThrows;
+import lombok.Value;
+
+import java.io.File;
+import java.util.Arrays;
+import java.util.List;
+
+@Value
+public class DeleteAndSaveInSameFolderUpdateTarget implements UpdateTarget {
+    File file;
+
+    @SneakyThrows
+    @Override
+    public List<UpdateAction> generateUpdateActions(PotentialUpdate update) {
+        return Arrays.asList(
+                new UpdateAction.DeleteFile(file),
+                new UpdateAction.MoveDownloadedFile(update.getUpdateJarStorage(), new File(file.getParentFile(), update.getFileName()))
+        );
+    }
+}
diff --git a/src/main/java/moe/nea/libautoupdate/ExitHookInvoker.java b/src/main/java/moe/nea/libautoupdate/ExitHookInvoker.java
new file mode 100644
index 0000000..55436a2
--- /dev/null
+++ b/src/main/java/moe/nea/libautoupdate/ExitHookInvoker.java
@@ -0,0 +1,62 @@
+package moe.nea.libautoupdate;
+
+import java.io.File;
+import java.util.ArrayList;
+import java.util.List;
+
+public class ExitHookInvoker {
+
+    private static boolean isExitHookRegistered = false;
+    private static List<UpdateAction> actions;
+    private static File updaterJar;
+    private static boolean cancelled = false;
+
+
+    public static synchronized void setExitHook(File updaterJar, List<UpdateAction> actions) {
+        if (!isExitHookRegistered) {
+            Runtime.getRuntime().addShutdownHook(new Thread(ExitHookInvoker::runExitHook));
+
+            isExitHookRegistered = true;
+        }
+        ExitHookInvoker.cancelled = false;
+        ExitHookInvoker.actions = actions;
+        ExitHookInvoker.updaterJar = updaterJar;
+    }
+
+    public static synchronized void cancelUpdate() {
+        cancelled = true;
+    }
+
+    private static synchronized String[] buildInvocation() {
+        boolean isWindows = System.getProperty("os.name", "").startsWith("Windows");
+        File javaBinary = new File(System.getProperty("java.home"), "bin/java" + (isWindows ? ".exe" : ""));
+
+
+        List<String> arguments = new ArrayList<>();
+        arguments.add(javaBinary.getAbsolutePath());
+        arguments.add("-jar");
+        arguments.add(updaterJar.getAbsolutePath());
+
+        for (UpdateAction action : actions) {
+            action.encode(arguments);
+        }
+
+        return arguments.toArray(new String[0]);
+    }
+
+    private static void runExitHook() {
+        try {
+            if (cancelled) {
+                System.out.println("Skipping cancelled update");
+                return;
+            }
+            String[] invocation = buildInvocation();
+            System.out.println("Running post updater using: " + String.join(" ", invocation));
+            Runtime.getRuntime().exec(invocation);
+        } catch (Throwable t) {
+            t.printStackTrace();
+        }
+    }
+
+
+}
diff --git a/src/main/java/moe/nea/libautoupdate/GistSource.java b/src/main/java/moe/nea/libautoupdate/GistSource.java
new file mode 100644
index 0000000..a8ae5df
--- /dev/null
+++ b/src/main/java/moe/nea/libautoupdate/GistSource.java
@@ -0,0 +1,19 @@
+package moe.nea.libautoupdate;
+
+import lombok.EqualsAndHashCode;
+import lombok.Value;
+
+import java.util.concurrent.CompletableFuture;
+
+@Value
+@EqualsAndHashCode(callSuper = false)
+public class GistSource extends JsonUpdateSource {
+    String owner;
+    String gistId;
+    private static final String GIST_RAW_URL = "https://gist.githubusercontent.com/%s/%s/raw/%s.json";
+
+    @Override
+    public CompletableFuture<UpdateData> checkUpdate(String updateStream) {
+        return getJsonFromURL(String.format(GIST_RAW_URL, owner, gistId, updateStream), UpdateData.class);
+    }
+}
diff --git a/src/main/java/moe/nea/libautoupdate/JsonUpdateSource.java b/src/main/java/moe/nea/libautoupdate/JsonUpdateSource.java
new file mode 100644
index 0000000..ffbfa9a
--- /dev/null
+++ b/src/main/java/moe/nea/libautoupdate/JsonUpdateSource.java
@@ -0,0 +1,31 @@
+package moe.nea.libautoupdate;
+
+import com.google.gson.Gson;
+import lombok.val;
+
+import java.io.IOException;
+import java.io.InputStreamReader;
+import java.net.URL;
+import java.nio.charset.StandardCharsets;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CompletionException;
+
+public abstract class JsonUpdateSource implements UpdateSource {
+    private static final Gson gson = new Gson();
+
+    protected Gson getGson() {
+        return gson;
+    }
+
+    protected <T> CompletableFuture<T> getJsonFromURL(String url, Class<T> clazz) {
+        return CompletableFuture.supplyAsync(() -> {
+            try {
+                try (val is = new URL(url).openStream()) {
+                    return getGson().fromJson(new InputStreamReader(is, StandardCharsets.UTF_8), clazz);
+                }
+            } catch (IOException e) {
+                throw new CompletionException(e);
+            }
+        });
+    }
+}
diff --git a/src/main/java/moe/nea/libautoupdate/Main.java b/src/main/java/moe/nea/libautoupdate/Main.java
new file mode 100644
index 0000000..5d98df1
--- /dev/null
+++ b/src/main/java/moe/nea/libautoupdate/Main.java
@@ -0,0 +1,22 @@
+package moe.nea.libautoupdate;
+
+public class Main {
+    public static void main(String[] args) {
+
+        UpdateContext updater = new UpdateContext(
+                UpdateSource.gistSource("romangraef", "9b62fe32bc41c09d2d7e2d3153f14ee8"),
+                UpdateTarget.deleteAndSaveInTheSameFolder(Main.class),
+                CurrentVersion.of(10000),
+                "test"
+        );
+        updater.cleanup();
+        System.out.println("Update cleaned");
+        System.out.println("Created update context: " + updater);
+        updater.checkUpdate("stable").thenCompose(it -> {
+            System.out.println("Checked for update on stable: " + it);
+            System.out.println("Can update: " + it.isUpdateAvailable());
+            System.out.println("Executing update.");
+            return it.launchUpdate();
+        }).join();
+    }
+}
diff --git a/src/main/java/moe/nea/libautoupdate/PotentialUpdate.java b/src/main/java/moe/nea/libautoupdate/PotentialUpdate.java
new file mode 100644
index 0000000..2e6142d
--- /dev/null
+++ b/src/main/java/moe/nea/libautoupdate/PotentialUpdate.java
@@ -0,0 +1,94 @@
+package moe.nea.libautoupdate;
+
+import lombok.Value;
+import lombok.val;
+
+import java.io.File;
+import java.io.FileInputStream;
+import java.io.FileOutputStream;
+import java.io.IOException;
+import java.net.MalformedURLException;
+import java.util.UUID;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CompletionException;
+
+@Value
+public class PotentialUpdate {
+    UpdateData update;
+    UpdateContext context;
+    UUID updateUUID = UUID.randomUUID();
+
+    public File getUpdateDirectory() {
+        return new File(".autoupdates", context.getIdentifier() + "/" + updateUUID);
+    }
+
+
+    public boolean isUpdateAvailable() {
+        if (update == null) return false;
+        return update.getVersionNumber() > context.getCurrentVersion().getCurrentVersionNumber();
+    }
+
+    private File getFile(String name) {
+        getUpdateDirectory().mkdirs();
+        return new File(getUpdateDirectory(), name);
+    }
+
+    public File getUpdateJarStorage() {
+        return getFile("next.jar");
+    }
+
+    public String getFileName() throws MalformedURLException {
+        val split = update.getDownloadAsURL().getPath().split("/");
+        return split[split.length - 1];
+    }
+
+
+    public void extractUpdater() throws IOException {
+        val file = getFile("updater.jar");
+        try (val from = getClass().getResourceAsStream("/updater.jar");
+             val to = new FileOutputStream(file)) {
+            UpdateUtils.connect(from, to);
+        }
+    }
+
+    public void downloadUpdate() throws IOException {
+        try (val from = update.getDownloadAsURL().openStream();
+             val to = new FileOutputStream(getUpdateJarStorage())) {
+            UpdateUtils.connect(from, to);
+        }
+        try (val check = new FileInputStream(getUpdateJarStorage())) {
+            val updateSha = UpdateUtils.sha256sum(check);
+            if (!update.getSha256().equalsIgnoreCase(updateSha)) {
+                throw new UpdateException(
+                        "Hash of downloaded file " + getUpdateJarStorage() +
+                                " (" + updateSha + ") does not match expected hash of " +
+                                update.getSha256());
+            }
+        }
+    }
+
+    public void prepareUpdate() throws IOException {
+        extractUpdater();
+        downloadUpdate();
+    }
+
+
+    public void executeUpdate() throws IOException {
+        prepareUpdate();
+        ExitHookInvoker.setExitHook(getFile("updater.jar"),
+                context.getTarget().generateUpdateActions(this));
+    }
+
+
+    public CompletableFuture<Void> launchUpdate() {
+        return CompletableFuture.supplyAsync(() -> {
+            try {
+                executeUpdate();
+            } catch (IOException e) {
+                throw new CompletionException(e);
+            }
+            return null;
+        });
+    }
+
+}
diff --git a/src/main/java/moe/nea/libautoupdate/ReplaceJarUpdateTarget.java b/src/main/java/moe/nea/libautoupdate/ReplaceJarUpdateTarget.java
new file mode 100644
index 0000000..f31481c
--- /dev/null
+++ b/src/main/java/moe/nea/libautoupdate/ReplaceJarUpdateTarget.java
@@ -0,0 +1,20 @@
+package moe.nea.libautoupdate;
+
+import lombok.Value;
+
+import java.io.File;
+import java.util.Arrays;
+import java.util.List;
+
+@Value
+public class ReplaceJarUpdateTarget implements UpdateTarget {
+    File file;
+
+    @Override
+    public List<UpdateAction> generateUpdateActions(PotentialUpdate update) {
+        return Arrays.asList(
+                new UpdateAction.DeleteFile(file),
+                new UpdateAction.MoveDownloadedFile(update.getUpdateJarStorage(), file)
+        );
+    }
+}
diff --git a/src/main/java/moe/nea/libautoupdate/UpdateAction.java b/src/main/java/moe/nea/libautoupdate/UpdateAction.java
new file mode 100644
index 0000000..ae9852e
--- /dev/null
+++ b/src/main/java/moe/nea/libautoupdate/UpdateAction.java
@@ -0,0 +1,43 @@
+package moe.nea.libautoupdate;
+
+import lombok.EqualsAndHashCode;
+import lombok.NonNull;
+import lombok.Value;
+
+import java.io.File;
+import java.util.List;
+
+public abstract class UpdateAction {
+    private UpdateAction() {
+    }
+
+    public abstract void encode(List<String> arguments);
+
+    @Value
+    @EqualsAndHashCode(callSuper = false)
+    public static class DeleteFile extends UpdateAction {
+        @NonNull File toDelete;
+
+        @Override
+        public void encode(List<String> arguments) {
+            arguments.add("delete");
+            arguments.add(toDelete.getAbsolutePath());
+        }
+    }
+
+    @Value
+    @EqualsAndHashCode(callSuper = false)
+    public static class MoveDownloadedFile extends UpdateAction {
+        @NonNull File whereFrom;
+        @NonNull File whereTo;
+
+        @Override
+        public void encode(List<String> arguments) {
+            arguments.add("move");
+            arguments.add(whereFrom.getAbsolutePath());
+            arguments.add(whereTo.getAbsolutePath());
+        }
+    }
+
+
+}
diff --git a/src/main/java/moe/nea/libautoupdate/UpdateContext.java b/src/main/java/moe/nea/libautoupdate/UpdateContext.java
new file mode 100644
index 0000000..fa8b8d5
--- /dev/null
+++ b/src/main/java/moe/nea/libautoupdate/UpdateContext.java
@@ -0,0 +1,32 @@
+package moe.nea.libautoupdate;
+
+
+import lombok.NonNull;
+import lombok.Value;
+
+import java.io.File;
+import java.io.IOException;
+import java.util.concurrent.CompletableFuture;
+
+@Value
+public class UpdateContext {
+    @NonNull UpdateSource source;
+    @NonNull UpdateTarget target;
+    @NonNull CurrentVersion currentVersion;
+    @NonNull String identifier;
+
+
+    public void cleanup() {
+        File file = new File(".autoupdates", identifier).getAbsoluteFile();
+        try {
+            UpdateUtils.deleteDirectory(file.toPath());
+        } catch (IOException e) {
+            e.printStackTrace();
+        }
+    }
+
+    public CompletableFuture<PotentialUpdate> checkUpdate(String updateStream) {
+        return source.checkUpdate(updateStream)
+                .thenApply(it -> new PotentialUpdate(it, this));
+    }
+}
diff --git a/src/main/java/moe/nea/libautoupdate/UpdateData.java b/src/main/java/moe/nea/libautoupdate/UpdateData.java
new file mode 100644
index 0000000..276ae5f
--- /dev/null
+++ b/src/main/java/moe/nea/libautoupdate/UpdateData.java
@@ -0,0 +1,22 @@
+package moe.nea.libautoupdate;
+
+import lombok.AllArgsConstructor;
+import lombok.Data;
+import lombok.NoArgsConstructor;
+
+import java.net.MalformedURLException;
+import java.net.URL;
+
+@Data
+@NoArgsConstructor
+@AllArgsConstructor
+public class UpdateData {
+    String versionName;
+    int versionNumber;
+    String sha256;
+    String download;
+
+    public URL getDownloadAsURL() throws MalformedURLException {
+        return new URL(download);
+    }
+}
diff --git a/src/main/java/moe/nea/libautoupdate/UpdateException.java b/src/main/java/moe/nea/libautoupdate/UpdateException.java
new file mode 100644
index 0000000..482263b
--- /dev/null
+++ b/src/main/java/moe/nea/libautoupdate/UpdateException.java
@@ -0,0 +1,7 @@
+package moe.nea.libautoupdate;
+
+public class UpdateException extends RuntimeException {
+    public UpdateException(String message) {
+        super(message);
+    }
+}
diff --git a/src/main/java/moe/nea/libautoupdate/UpdateSource.java b/src/main/java/moe/nea/libautoupdate/UpdateSource.java
new file mode 100644
index 0000000..1407063
--- /dev/null
+++ b/src/main/java/moe/nea/libautoupdate/UpdateSource.java
@@ -0,0 +1,11 @@
+package moe.nea.libautoupdate;
+
+import java.util.concurrent.CompletableFuture;
+
+public interface UpdateSource {
+    static UpdateSource gistSource(String owner, String gistId) {
+        return new GistSource(owner, gistId);
+    }
+
+    CompletableFuture<UpdateData> checkUpdate(String updateStream);
+}
diff --git a/src/main/java/moe/nea/libautoupdate/UpdateTarget.java b/src/main/java/moe/nea/libautoupdate/UpdateTarget.java
new file mode 100644
index 0000000..fcaaa3e
--- /dev/null
+++ b/src/main/java/moe/nea/libautoupdate/UpdateTarget.java
@@ -0,0 +1,19 @@
+package moe.nea.libautoupdate;
+
+import java.io.File;
+import java.util.List;
+
+public interface UpdateTarget {
+    List<UpdateAction> generateUpdateActions(PotentialUpdate update);
+
+    static UpdateTarget replaceJar(Class<?> containedClass) {
+        File file = UpdateUtils.getJarFileContainingClass(containedClass);
+        return new ReplaceJarUpdateTarget(file);
+    }
+
+    static UpdateTarget deleteAndSaveInTheSameFolder(Class<?> containedClass) {
+        File file = UpdateUtils.getJarFileContainingClass(containedClass);
+        return new DeleteAndSaveInSameFolderUpdateTarget(file);
+    }
+
+}
diff --git a/src/main/java/moe/nea/libautoupdate/UpdateUtils.java b/src/main/java/moe/nea/libautoupdate/UpdateUtils.java
new file mode 100644
index 0000000..ce09c0b
--- /dev/null
+++ b/src/main/java/moe/nea/libautoupdate/UpdateUtils.java
@@ -0,0 +1,86 @@
+package moe.nea.libautoupdate;
+
+import lombok.val;
+import lombok.var;
+
+import java.io.File;
+import java.io.IOException;
+import java.io.InputStream;
+import java.io.OutputStream;
+import java.math.BigInteger;
+import java.net.URI;
+import java.net.URISyntaxException;
+import java.nio.file.FileVisitResult;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.nio.file.SimpleFileVisitor;
+import java.nio.file.attribute.BasicFileAttributes;
+import java.security.MessageDigest;
+import java.security.NoSuchAlgorithmException;
+import java.util.Locale;
+
+public class UpdateUtils {
+    private UpdateUtils() {
+    }
+
+    public static File getJarFileContainingClass(Class<?> clazz) {
+        val location = clazz.getProtectionDomain().getCodeSource().getLocation();
+        if (location == null)
+            return null;
+        var path = location.toString();
+        path = path.split("!", 2)[0];
+        if (path.startsWith("jar:")) {
+            path = path.substring(4);
+        }
+        try {
+            return new File(new URI(path));
+        } catch (URISyntaxException e) {
+            e.printStackTrace();
+            return null;
+        }
+    }
+
+    public static void connect(InputStream from, OutputStream to) throws IOException {
+        val buf = new byte[4096];
+        int r;
+        while ((r = from.read(buf)) != -1) {
+            to.write(buf, 0, r);
+        }
+    }
+
+    public static String sha256sum(InputStream stream) throws IOException {
+        try {
+            val digest = MessageDigest.getInstance("SHA-256");
+            int r;
+            val buf = new byte[4096];
+            while ((r = stream.read(buf)) != -1) {
+                digest.update(buf, 0, r);
+            }
+            return String.format("%64s",
+                            new BigInteger(1, digest.digest())
+                                    .toString(16))
+                    .replace(' ', '0')
+                    .toLowerCase(Locale.ROOT);
+        } catch (NoSuchAlgorithmException e) {
+            throw new RuntimeException(e);
+        }
+    }
+
+    public static void deleteDirectory(Path path) throws IOException {
+        if (!Files.exists(path)) return;
+        Files.walkFileTree(path, new SimpleFileVisitor<Path>() {
+            @Override
+            public FileVisitResult postVisitDirectory(Path dir, IOException exc) throws IOException {
+                Files.delete(dir);
+                return FileVisitResult.CONTINUE;
+            }
+
+            @Override
+            public FileVisitResult visitFile(Path file, BasicFileAttributes attrs) throws IOException {
+                Files.delete(file);
+                return FileVisitResult.CONTINUE;
+            }
+        });
+    }
+}
+