waifu2xlemon/src/main/java/moe/lemonsh/waifu2xlemon/App.java

package moe.lemonsh.waifu2xlemon;

import javax.servlet.MultipartConfigElement;
import javax.servlet.http.Part;
import java.io.*;
import java.net.URLConnection;
import java.nio.charset.StandardCharsets;
import java.util.Objects;

import static spark.Spark.*;

public class App {
    private static final int maxFileSize = 8380416;
    public static void main(String[] argv) throws IOException {
        String envport = System.getenv("WAIFU2XLEMON_PORT");
        String envexec = System.getenv("WAIFU2XLEMON_EXECUTABLE");
        if (envport == null || envexec == null) {
            System.out.println("Error: Invalid config. Make sure that the WAIFU2XLEMON_PORT and WAIFU2XLEMON_EXECUTABLE envvars are set.");
            System.exit(1);
        }
        String envpass = System.getenv("WAIFU2XLEMON_PASSWORD");
        if ("RANDOM".equals(envpass)) {
            var passwordFile = new File("password-dont-edit.txt");
            if (passwordFile.exists()) {
                try (var fis = new FileInputStream(passwordFile)) {
                    envpass = new String(fis.readAllBytes(), StandardCharsets.UTF_8);
                }
            } else {
                envpass = AccessTokenGenerator.randomToken(16);
                try (var fos = new FileOutputStream(passwordFile)) {
                    fos.write(envpass.getBytes(StandardCharsets.UTF_8));
                }
            }
        }
        Waifu2x waifu2x = new Waifu2x(envexec);
        threadPool(4, 1, -1);
        try {
            port(Integer.parseInt(envport));
        } catch (NumberFormatException e) {
            System.out.println("Error: Invalid config. WAIFU2XLEMON_PORT contains an invalid value.");
            System.exit(2);
        }
        MultipartConfigElement multipartConfigElement = new MultipartConfigElement("multipartTemp", maxFileSize, maxFileSize, maxFileSize+1);
        String indexPage, loginPage;
        try (InputStream indexResource = App.class.getResourceAsStream("/app.html");
             InputStream loginResource = App.class.getResourceAsStream("/login.html")) {
            indexPage = new String(Objects.requireNonNull(indexResource).readAllBytes(), StandardCharsets.UTF_8);
            loginPage = new String(Objects.requireNonNull(loginResource).readAllBytes(), StandardCharsets.UTF_8);
        }
        String password = envpass;
        get("/", (req, resp) -> {
            if (password == null) {
                return indexPage;
            } else {
                var loginValue = req.session().attribute("login");
                return loginValue != null && (boolean) loginValue ? indexPage : loginPage;
            }
        });
        post("/login", (req, resp) -> {
            String submittedToken = req.queryParams("pass");
            if (submittedToken == null) halt(400);
            if (!submittedToken.equals(password)) return "Invalid password";
            resp.redirect("/");
            req.session().attribute("login", true);
            return null;
        });
        post("/logout", (req, resp) -> {
            resp.redirect("/");
            req.session().attribute("login", false);
            return null;
        });
        post("/", (req, resp) -> {
            req.raw().setAttribute("org.eclipse.jetty.multipartConfig", multipartConfigElement);
            Part filePart, noisePart, typePart, scalePart;
            try {
                filePart = req.raw().getPart("file");
                scalePart = req.raw().getPart("scale");
                noisePart = req.raw().getPart("noise");
                typePart = req.raw().getPart("ftype");
            } catch (IllegalStateException e) {
                return null;
            }
            if (filePart == null || scalePart == null || noisePart == null || typePart == null) halt(400);
            int scale = 0, noise = 0, type = 0;
            try {
                scale = Integer.parseInt(new String(scalePart.getInputStream().readAllBytes(), StandardCharsets.UTF_8));
                noise = Integer.parseInt(new String(noisePart.getInputStream().readAllBytes(), StandardCharsets.UTF_8));
                type = Integer.parseInt(new String(typePart.getInputStream().readAllBytes(), StandardCharsets.UTF_8));
            } catch (NumberFormatException e) {
                halt(400);
            }
            if (scale < 1 || scale > 2 || noise < -1 || noise > 2 || type < 0 || type > 3) halt(400); // boundary check
            String mimeType = URLConnection.guessContentTypeFromName(filePart.getSubmittedFileName());
            String inputExtension, outputExtension = null;
            switch (mimeType) {
                case "image/png" -> inputExtension = ".png";
                case "image/jpeg" -> inputExtension = ".jpg";
                case "image/webp" -> inputExtension = ".webp";
                default -> {
                    return "Only PNG/JPG/WEBP files are allowed.";
                }
            }
            switch (type) {
                case 0 -> outputExtension = ".png";
                case 1 -> outputExtension = ".jpg";
                case 2 -> outputExtension = ".webp";
                default -> halt(400);
            }
            Waifu2x.WaifuResult waifuResult;
            waifuResult = waifu2x.run(filePart.getInputStream(), inputExtension, outputExtension, (char)(noise+'0'), (char)(scale+'0'));
            if (waifuResult.success) {
                resp.raw().setContentType(URLConnection.guessContentTypeFromName(waifuResult.filename));
                resp.raw().setHeader("Content-Disposition","attachment; filename="+waifuResult.filename);
                resp.raw().setContentLengthLong(waifuResult.size);
                waifuResult.image.transferTo(resp.raw().getOutputStream());
                resp.raw().getOutputStream().flush();
                resp.raw().getOutputStream().close();
                return null;
            }
            return "Error.\n"+waifuResult.stdout;
        });
    }
}