Spaces:

kvashist
/

AskMe

Sleeping

AskMe / KshitijAIProxyServer.java

Create KshitijAIProxyServer.java

72e8ca8 verified 9 months ago

6.05 kB

import com.sun.net.httpserver.Headers;
import com.sun.net.httpserver.HttpServer;

import java.io.IOException;
import java.io.OutputStream;
import java.net.InetSocketAddress;
import java.net.URI;
import java.net.http.HttpClient;
import java.net.http.HttpRequest;
import java.net.http.HttpResponse;
import java.nio.charset.StandardCharsets;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.time.Duration;

public class KshitijAIProxyServer {

    private static final int PORT = 7860;

    // Secrets (set in HF Space settings)
    private static final String BACKEND_URL = getenvOrThrow("BACKEND_URL"); // e.g. https://user-kshitijai-backend.hf.space
    private static final String HF_TOKEN = getenvOrThrow("HF_TOKEN");       // token that can access private space

    public static void main(String[] args) throws Exception {
        final Path baseDir = Paths.get(System.getProperty("user.dir")).toRealPath();

        HttpServer server = HttpServer.create(new InetSocketAddress(PORT), 0);

        server.createContext("/health", exchange -> {
            addCors(exchange.getResponseHeaders());
            if (!exchange.getRequestMethod().equalsIgnoreCase("GET")) {
                exchange.sendResponseHeaders(405, -1);
                return;
            }
            byte[] ok = "ok".getBytes(StandardCharsets.UTF_8);
            exchange.getResponseHeaders().set("Content-Type", "text/plain; charset=utf-8");
            exchange.sendResponseHeaders(200, ok.length);
            try (OutputStream os = exchange.getResponseBody()) { os.write(ok); }
        });

        // Proxy: browser -> public space -> private backend space
        server.createContext("/api/chat", exchange -> {
            Headers h = exchange.getResponseHeaders();
            addCors(h);

            // Handle preflight
            if (exchange.getRequestMethod().equalsIgnoreCase("OPTIONS")) {
                exchange.sendResponseHeaders(204, -1);
                return;
            }

            if (!exchange.getRequestMethod().equalsIgnoreCase("POST")) {
                exchange.sendResponseHeaders(405, -1);
                return;
            }

            byte[] requestBytes = exchange.getRequestBody().readAllBytes();

            HttpClient client = HttpClient.newBuilder()
                    .connectTimeout(Duration.ofSeconds(10))
                    .build();

            String target = BACKEND_URL.endsWith("/")
                    ? BACKEND_URL + "api/chat"
                    : BACKEND_URL + "/api/chat";

            HttpRequest req = HttpRequest.newBuilder()
                    .uri(URI.create(target))
                    .timeout(Duration.ofSeconds(300))
                    .header("Content-Type", "application/json")
                    .header("Authorization", "Bearer " + HF_TOKEN)
                    .POST(HttpRequest.BodyPublishers.ofByteArray(requestBytes))
                    .build();

            HttpResponse<byte[]> resp;
            try {
                resp = client.send(req, HttpResponse.BodyHandlers.ofByteArray());
            } catch (Exception e) {
                byte[] err = ("Failed to reach backend at " + target + "\n" + e)
                        .getBytes(StandardCharsets.UTF_8);
                h.set("Content-Type", "text/plain; charset=utf-8");
                exchange.sendResponseHeaders(502, err.length);
                try (OutputStream os = exchange.getResponseBody()) { os.write(err); }
                return;
            }

            h.set("Content-Type", "application/json; charset=utf-8");
            exchange.sendResponseHeaders(resp.statusCode(), resp.body().length);
            try (OutputStream os = exchange.getResponseBody()) { os.write(resp.body()); }
        });

        // Static files: / -> index.html, /style.css, /app.js
        server.createContext("/", exchange -> {
            addCors(exchange.getResponseHeaders());

            if (!exchange.getRequestMethod().equalsIgnoreCase("GET")) {
                exchange.sendResponseHeaders(405, -1);
                return;
            }

            String urlPath = exchange.getRequestURI().getPath();
            if (urlPath == null || urlPath.isBlank() || urlPath.equals("/")) urlPath = "/index.html";

            String rel = urlPath.startsWith("/") ? urlPath.substring(1) : urlPath;
            Path requested = baseDir.resolve(rel).normalize();

            if (!requested.startsWith(baseDir)) { exchange.sendResponseHeaders(403, -1); return; }
            if (!Files.exists(requested) || Files.isDirectory(requested)) { exchange.sendResponseHeaders(404, -1); return; }

            String ct = guessContentType(requested);
            byte[] bytes = Files.readAllBytes(requested);

            exchange.getResponseHeaders().set("Content-Type", ct);
            exchange.sendResponseHeaders(200, bytes.length);
            try (OutputStream os = exchange.getResponseBody()) { os.write(bytes); }
        });

        server.start();
        System.out.println("Public UI + Proxy running on port " + PORT);
    }

    private static void addCors(Headers h) {
        h.set("Access-Control-Allow-Origin", "*");
        h.set("Access-Control-Allow-Methods", "GET, POST, OPTIONS");
        h.set("Access-Control-Allow-Headers", "Content-Type, Authorization");
    }

    private static String guessContentType(Path file) throws IOException {
        String name = file.getFileName().toString().toLowerCase();
        if (name.endsWith(".html")) return "text/html; charset=utf-8";
        if (name.endsWith(".css")) return "text/css; charset=utf-8";
        if (name.endsWith(".js")) return "application/javascript; charset=utf-8";
        String probe = Files.probeContentType(file);
        return probe != null ? probe : "application/octet-stream";
    }

    private static String getenvOrThrow(String key) {
        String v = System.getenv(key);
        if (v == null || v.isBlank()) throw new IllegalStateException("Missing env var: " + key);
        return v;
    }
}

	import com.sun.net.httpserver.Headers;
	import com.sun.net.httpserver.HttpServer;

	import java.io.IOException;
	import java.io.OutputStream;
	import java.net.InetSocketAddress;
	import java.net.URI;
	import java.net.http.HttpClient;
	import java.net.http.HttpRequest;
	import java.net.http.HttpResponse;
	import java.nio.charset.StandardCharsets;
	import java.nio.file.Files;
	import java.nio.file.Path;
	import java.nio.file.Paths;
	import java.time.Duration;

	public class KshitijAIProxyServer {

	private static final int PORT = 7860;

	// Secrets (set in HF Space settings)
	private static final String BACKEND_URL = getenvOrThrow("BACKEND_URL"); // e.g. https://user-kshitijai-backend.hf.space
	private static final String HF_TOKEN = getenvOrThrow("HF_TOKEN"); // token that can access private space

	public static void main(String[] args) throws Exception {
	final Path baseDir = Paths.get(System.getProperty("user.dir")).toRealPath();

	HttpServer server = HttpServer.create(new InetSocketAddress(PORT), 0);

	server.createContext("/health", exchange -> {
	addCors(exchange.getResponseHeaders());
	if (!exchange.getRequestMethod().equalsIgnoreCase("GET")) {
	exchange.sendResponseHeaders(405, -1);
	return;
	}
	byte[] ok = "ok".getBytes(StandardCharsets.UTF_8);
	exchange.getResponseHeaders().set("Content-Type", "text/plain; charset=utf-8");
	exchange.sendResponseHeaders(200, ok.length);
	try (OutputStream os = exchange.getResponseBody()) { os.write(ok); }
	});

	// Proxy: browser -> public space -> private backend space
	server.createContext("/api/chat", exchange -> {
	Headers h = exchange.getResponseHeaders();
	addCors(h);

	// Handle preflight
	if (exchange.getRequestMethod().equalsIgnoreCase("OPTIONS")) {
	exchange.sendResponseHeaders(204, -1);
	return;
	}

	if (!exchange.getRequestMethod().equalsIgnoreCase("POST")) {
	exchange.sendResponseHeaders(405, -1);
	return;
	}

	byte[] requestBytes = exchange.getRequestBody().readAllBytes();

	HttpClient client = HttpClient.newBuilder()
	.connectTimeout(Duration.ofSeconds(10))
	.build();

	String target = BACKEND_URL.endsWith("/")
	? BACKEND_URL + "api/chat"
	: BACKEND_URL + "/api/chat";

	HttpRequest req = HttpRequest.newBuilder()
	.uri(URI.create(target))
	.timeout(Duration.ofSeconds(300))
	.header("Content-Type", "application/json")
	.header("Authorization", "Bearer " + HF_TOKEN)
	.POST(HttpRequest.BodyPublishers.ofByteArray(requestBytes))
	.build();

	HttpResponse<byte[]> resp;
	try {
	resp = client.send(req, HttpResponse.BodyHandlers.ofByteArray());
	} catch (Exception e) {
	byte[] err = ("Failed to reach backend at " + target + "\n" + e)
	.getBytes(StandardCharsets.UTF_8);
	h.set("Content-Type", "text/plain; charset=utf-8");
	exchange.sendResponseHeaders(502, err.length);
	try (OutputStream os = exchange.getResponseBody()) { os.write(err); }
	return;
	}

	h.set("Content-Type", "application/json; charset=utf-8");
	exchange.sendResponseHeaders(resp.statusCode(), resp.body().length);
	try (OutputStream os = exchange.getResponseBody()) { os.write(resp.body()); }
	});

	// Static files: / -> index.html, /style.css, /app.js
	server.createContext("/", exchange -> {
	addCors(exchange.getResponseHeaders());

	if (!exchange.getRequestMethod().equalsIgnoreCase("GET")) {
	exchange.sendResponseHeaders(405, -1);
	return;
	}

	String urlPath = exchange.getRequestURI().getPath();
	if (urlPath == null \|\| urlPath.isBlank() \|\| urlPath.equals("/")) urlPath = "/index.html";

	String rel = urlPath.startsWith("/") ? urlPath.substring(1) : urlPath;
	Path requested = baseDir.resolve(rel).normalize();

	if (!requested.startsWith(baseDir)) { exchange.sendResponseHeaders(403, -1); return; }
	if (!Files.exists(requested) \|\| Files.isDirectory(requested)) { exchange.sendResponseHeaders(404, -1); return; }

	String ct = guessContentType(requested);
	byte[] bytes = Files.readAllBytes(requested);

	exchange.getResponseHeaders().set("Content-Type", ct);
	exchange.sendResponseHeaders(200, bytes.length);
	try (OutputStream os = exchange.getResponseBody()) { os.write(bytes); }
	});

	server.start();
	System.out.println("Public UI + Proxy running on port " + PORT);
	}

	private static void addCors(Headers h) {
	h.set("Access-Control-Allow-Origin", "*");
	h.set("Access-Control-Allow-Methods", "GET, POST, OPTIONS");
	h.set("Access-Control-Allow-Headers", "Content-Type, Authorization");
	}

	private static String guessContentType(Path file) throws IOException {
	String name = file.getFileName().toString().toLowerCase();
	if (name.endsWith(".html")) return "text/html; charset=utf-8";
	if (name.endsWith(".css")) return "text/css; charset=utf-8";
	if (name.endsWith(".js")) return "application/javascript; charset=utf-8";
	String probe = Files.probeContentType(file);
	return probe != null ? probe : "application/octet-stream";
	}

	private static String getenvOrThrow(String key) {
	String v = System.getenv(key);
	if (v == null \|\| v.isBlank()) throw new IllegalStateException("Missing env var: " + key);
	return v;
	}
	}