136 lines
6.5 KiB
Java
136 lines
6.5 KiB
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 java.util.concurrent.CountDownLatch;
|
|
|
|
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);
|
|
new Thread(waifu2x::runQueue).start();
|
|
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);
|
|
}
|
|
final var latch = new CountDownLatch(1);
|
|
var resultRef = new Object() {
|
|
Waifu2x.WaifuResult waifuResult;
|
|
};
|
|
var task = new Waifu2x.WaifuTask(filePart.getInputStream().readAllBytes(), inputExtension, outputExtension, (char)(noise+'0'), (char)(scale+'0'), result -> {
|
|
resultRef.waifuResult = result;
|
|
latch.countDown();
|
|
});
|
|
switch (waifu2x.addTask(task)) {
|
|
case 1 -> { return "Invalid image (too big or in non-supported format)"; }
|
|
case 2 -> { return "The processing queue is full. Please try again later."; }
|
|
}
|
|
latch.await();
|
|
if (resultRef.waifuResult.success) {
|
|
resp.raw().setContentType(URLConnection.guessContentTypeFromName(resultRef.waifuResult.filename));
|
|
resp.raw().setHeader("Content-Disposition","attachment; filename="+resultRef.waifuResult.filename);
|
|
resp.raw().setContentLengthLong(resultRef.waifuResult.size);
|
|
resultRef.waifuResult.image.transferTo(resp.raw().getOutputStream());
|
|
resp.raw().getOutputStream().flush();
|
|
resp.raw().getOutputStream().close();
|
|
return null;
|
|
}
|
|
return "Error.\n"+resultRef.waifuResult.stdout;
|
|
});
|
|
}
|
|
}
|