mirror of
https://github.com/Stirling-Tools/Stirling-PDF.git
synced 2026-09-03 05:10:16 +03:00
Compare commits
1
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
00a36ba0f5 |
@@ -283,6 +283,7 @@ public class ApplicationProperties {
|
||||
|
||||
private Valkey valkey = new Valkey();
|
||||
private Node node = new Node();
|
||||
private Engine engine = new Engine();
|
||||
|
||||
private transient String cachedNodeId;
|
||||
|
||||
@@ -354,6 +355,16 @@ public class ApplicationProperties {
|
||||
/** Heartbeat publish interval for the instance registry, in milliseconds. */
|
||||
private long heartbeatIntervalMs = 5000;
|
||||
}
|
||||
|
||||
@Data
|
||||
public static class Engine {
|
||||
/**
|
||||
* Shared secret for the {@code X-Engine-Auth} header on Java -> Python engine calls.
|
||||
* Must match {@code STIRLING_ENGINE_SHARED_SECRET} on the engine pod. Blank = dev mode
|
||||
* (no auth) - never leave blank in production.
|
||||
*/
|
||||
private String sharedSecret = "";
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
|
||||
@@ -345,6 +345,8 @@ cluster:
|
||||
internalAddress: "" # host:port advertised in the instance registry for peer-to-peer cluster traffic. Blank = derived at startup.
|
||||
scheme: http # 'http' or 'https' - scheme peers use to call this node's /internal/cluster/** endpoints
|
||||
heartbeatIntervalMs: 5000 # Heartbeat publish interval for the instance registry (ms)
|
||||
engine:
|
||||
sharedSecret: "" # Shared secret for Java <-> Python AI engine auth. Generate via 'openssl rand -hex 32' and set BOTH this value and the engine pod's STIRLING_ENGINE_SHARED_SECRET env var to the same string. Blank = dev mode (auth disabled, NEVER in production).
|
||||
|
||||
pdfEditor:
|
||||
fallback-font: classpath:/static/fonts/NotoSans-Regular.ttf # Override to point at a custom fallback font
|
||||
|
||||
+1
-1
@@ -87,7 +87,7 @@ public class AiEngineController {
|
||||
@GetMapping("/health")
|
||||
@Operation(
|
||||
summary = "AI engine health check",
|
||||
description = "Returns the health status of the AI engine including configured models")
|
||||
description = "Returns the health status of the AI engine.")
|
||||
public ResponseEntity<String> health() throws IOException {
|
||||
String response = aiEngineClient.get("/health");
|
||||
return ResponseEntity.ok().contentType(MediaType.APPLICATION_JSON).body(response);
|
||||
|
||||
+134
-33
@@ -1,7 +1,9 @@
|
||||
package stirling.software.proprietary.service;
|
||||
|
||||
import java.io.IOException;
|
||||
import java.net.InetAddress;
|
||||
import java.net.URI;
|
||||
import java.net.UnknownHostException;
|
||||
import java.net.http.HttpClient;
|
||||
import java.net.http.HttpRequest;
|
||||
import java.net.http.HttpResponse;
|
||||
@@ -12,9 +14,15 @@ import java.util.stream.Stream;
|
||||
|
||||
import org.springframework.beans.factory.annotation.Autowired;
|
||||
import org.springframework.http.HttpStatus;
|
||||
import org.springframework.security.authentication.AnonymousAuthenticationToken;
|
||||
import org.springframework.security.core.Authentication;
|
||||
import org.springframework.security.core.context.SecurityContextHolder;
|
||||
import org.springframework.security.core.userdetails.UserDetails;
|
||||
import org.springframework.stereotype.Service;
|
||||
import org.springframework.web.server.ResponseStatusException;
|
||||
|
||||
import jakarta.annotation.PostConstruct;
|
||||
|
||||
import lombok.extern.slf4j.Slf4j;
|
||||
|
||||
import stirling.software.common.model.ApplicationProperties;
|
||||
@@ -23,6 +31,9 @@ import stirling.software.common.model.ApplicationProperties;
|
||||
@Service
|
||||
public class AiEngineClient {
|
||||
|
||||
static final String ENGINE_AUTH_HEADER = "X-Engine-Auth";
|
||||
static final String USER_ID_HEADER = "X-User-Id";
|
||||
|
||||
private final ApplicationProperties applicationProperties;
|
||||
private final HttpClient httpClient;
|
||||
|
||||
@@ -43,15 +54,113 @@ public class AiEngineClient {
|
||||
this.httpClient = httpClient;
|
||||
}
|
||||
|
||||
@PostConstruct
|
||||
void warnIfEngineSecretSentOverPlaintextHttp() {
|
||||
String engineSecret = resolveEngineSecret();
|
||||
if (engineSecret == null || engineSecret.isBlank()) {
|
||||
return;
|
||||
}
|
||||
String engineUrl = applicationProperties.getAiEngine().getUrl();
|
||||
if (engineUrl == null || !engineUrl.startsWith("http://")) {
|
||||
return;
|
||||
}
|
||||
URI uri;
|
||||
try {
|
||||
uri = URI.create(engineUrl);
|
||||
} catch (IllegalArgumentException ex) {
|
||||
log.warn(
|
||||
"Engine URL {} could not be parsed; skipping plaintext-secret check",
|
||||
engineUrl);
|
||||
return;
|
||||
}
|
||||
String host = uri.getHost();
|
||||
if (host == null || isLoopbackOrPrivateHost(host)) {
|
||||
return;
|
||||
}
|
||||
log.error(
|
||||
"SECURITY: engine shared secret will be sent over plaintext HTTP to a non-loopback /"
|
||||
+ " non-private-network host: {}. Use https:// for the engine URL in"
|
||||
+ " production.",
|
||||
engineUrl);
|
||||
}
|
||||
|
||||
private static boolean isLoopbackOrPrivateHost(String host) {
|
||||
// Cover the literal forms first to avoid a DNS lookup on the common cases.
|
||||
if (host.equalsIgnoreCase("localhost")
|
||||
|| host.equals("127.0.0.1")
|
||||
|| host.equals("::1")
|
||||
|| host.equals("0:0:0:0:0:0:0:1")) {
|
||||
return true;
|
||||
}
|
||||
try {
|
||||
InetAddress addr = InetAddress.getByName(host);
|
||||
if (addr.isLoopbackAddress()
|
||||
|| addr.isSiteLocalAddress()
|
||||
|| addr.isLinkLocalAddress()) {
|
||||
return true;
|
||||
}
|
||||
byte[] bytes = addr.getAddress();
|
||||
if (bytes.length == 4) {
|
||||
int b0 = bytes[0] & 0xFF;
|
||||
int b1 = bytes[1] & 0xFF;
|
||||
// 10/8
|
||||
if (b0 == 10) return true;
|
||||
// 172.16/12
|
||||
if (b0 == 172 && b1 >= 16 && b1 <= 31) return true;
|
||||
// 192.168/16
|
||||
if (b0 == 192 && b1 == 168) return true;
|
||||
}
|
||||
return false;
|
||||
} catch (UnknownHostException ex) {
|
||||
// Cannot resolve - treat as remote so we err on the side of warning.
|
||||
return false;
|
||||
}
|
||||
}
|
||||
|
||||
private HttpRequest.Builder decorate(HttpRequest.Builder builder) {
|
||||
String secret = resolveEngineSecret();
|
||||
if (secret != null && !secret.isBlank()) {
|
||||
builder.header(ENGINE_AUTH_HEADER, secret);
|
||||
}
|
||||
String userId = resolveUserId();
|
||||
if (userId != null && !userId.isBlank()) {
|
||||
builder.header(USER_ID_HEADER, userId);
|
||||
}
|
||||
return builder;
|
||||
}
|
||||
|
||||
private String resolveEngineSecret() {
|
||||
ApplicationProperties.Cluster cluster = applicationProperties.getCluster();
|
||||
if (cluster == null || cluster.getEngine() == null) {
|
||||
return "";
|
||||
}
|
||||
return cluster.getEngine().getSharedSecret();
|
||||
}
|
||||
|
||||
private String resolveUserId() {
|
||||
try {
|
||||
Authentication auth = SecurityContextHolder.getContext().getAuthentication();
|
||||
if (auth == null
|
||||
|| !auth.isAuthenticated()
|
||||
|| auth instanceof AnonymousAuthenticationToken) {
|
||||
return null;
|
||||
}
|
||||
Object principal = auth.getPrincipal();
|
||||
if (principal instanceof UserDetails ud) {
|
||||
return ud.getUsername();
|
||||
}
|
||||
return auth.getName();
|
||||
} catch (RuntimeException ex) {
|
||||
return null;
|
||||
}
|
||||
}
|
||||
|
||||
public String post(String path, String jsonBody) throws IOException {
|
||||
ApplicationProperties.AiEngine config = applicationProperties.getAiEngine();
|
||||
return postWithTimeout(path, jsonBody, Duration.ofSeconds(config.getTimeoutSeconds()));
|
||||
}
|
||||
|
||||
/**
|
||||
* POST with an explicit per-call timeout, for heavy operations (e.g. RAG ingestion of a large
|
||||
* document) that legitimately take longer than the default timeout.
|
||||
*/
|
||||
/** POST with an explicit per-call timeout for heavy operations (e.g. RAG ingestion). */
|
||||
public String postLongRunning(String path, String jsonBody) throws IOException {
|
||||
ApplicationProperties.AiEngine config = applicationProperties.getAiEngine();
|
||||
return postWithTimeout(
|
||||
@@ -70,12 +179,13 @@ public class AiEngineClient {
|
||||
log.debug("Proxying AI engine request to {} (timeout {}s)", url, timeout.toSeconds());
|
||||
|
||||
HttpRequest request =
|
||||
HttpRequest.newBuilder()
|
||||
.uri(URI.create(url))
|
||||
.header("Content-Type", "application/json")
|
||||
.header("Accept", "application/json")
|
||||
.timeout(timeout)
|
||||
.POST(HttpRequest.BodyPublishers.ofString(jsonBody))
|
||||
decorate(
|
||||
HttpRequest.newBuilder()
|
||||
.uri(URI.create(url))
|
||||
.header("Content-Type", "application/json")
|
||||
.header("Accept", "application/json")
|
||||
.timeout(timeout)
|
||||
.POST(HttpRequest.BodyPublishers.ofString(jsonBody)))
|
||||
.build();
|
||||
|
||||
HttpResponse<String> response = sendRequest(request);
|
||||
@@ -85,15 +195,7 @@ public class AiEngineClient {
|
||||
return response.body();
|
||||
}
|
||||
|
||||
/**
|
||||
* POST a JSON body and consume the response as a stream of NDJSON lines. Each line is passed to
|
||||
* {@code lineConsumer} in arrival order; the call returns when the engine closes the stream.
|
||||
*
|
||||
* <p>This is the right shape for long-running orchestrator calls that emit incremental
|
||||
* progress. The total HTTP timeout is the long-running timeout (typically 600s+), but in
|
||||
* practice line arrival keeps the connection logically alive: as long as the engine emits
|
||||
* events, the work is progressing. Genuine engine hangs still hit the total timeout.
|
||||
*/
|
||||
/** POST a JSON body and consume the response as a stream of NDJSON lines. */
|
||||
public void streamPost(String path, String jsonBody, Consumer<String> lineConsumer)
|
||||
throws IOException {
|
||||
ApplicationProperties.AiEngine config = applicationProperties.getAiEngine();
|
||||
@@ -110,12 +212,13 @@ public class AiEngineClient {
|
||||
timeout.toSeconds());
|
||||
|
||||
HttpRequest request =
|
||||
HttpRequest.newBuilder()
|
||||
.uri(URI.create(url))
|
||||
.header("Content-Type", "application/json")
|
||||
.header("Accept", "application/x-ndjson")
|
||||
.timeout(timeout)
|
||||
.POST(HttpRequest.BodyPublishers.ofString(jsonBody))
|
||||
decorate(
|
||||
HttpRequest.newBuilder()
|
||||
.uri(URI.create(url))
|
||||
.header("Content-Type", "application/json")
|
||||
.header("Accept", "application/x-ndjson")
|
||||
.timeout(timeout)
|
||||
.POST(HttpRequest.BodyPublishers.ofString(jsonBody)))
|
||||
.build();
|
||||
|
||||
HttpResponse<Stream<String>> response;
|
||||
@@ -160,11 +263,12 @@ public class AiEngineClient {
|
||||
log.debug("Proxying AI engine GET request to {}", url);
|
||||
|
||||
HttpRequest request =
|
||||
HttpRequest.newBuilder()
|
||||
.uri(URI.create(url))
|
||||
.header("Accept", "application/json")
|
||||
.timeout(Duration.ofSeconds(config.getTimeoutSeconds()))
|
||||
.GET()
|
||||
decorate(
|
||||
HttpRequest.newBuilder()
|
||||
.uri(URI.create(url))
|
||||
.header("Accept", "application/json")
|
||||
.timeout(Duration.ofSeconds(config.getTimeoutSeconds()))
|
||||
.GET())
|
||||
.build();
|
||||
|
||||
HttpResponse<String> response = sendRequest(request);
|
||||
@@ -180,9 +284,6 @@ public class AiEngineClient {
|
||||
} catch (HttpTimeoutException e) {
|
||||
throw new ResponseStatusException(HttpStatus.GATEWAY_TIMEOUT, "AI engine timed out", e);
|
||||
} catch (IOException e) {
|
||||
// Connection refused, DNS failure, socket reset, etc. — surface as
|
||||
// SERVICE_UNAVAILABLE so every caller of this client sees a structured
|
||||
// status rather than a raw 500 from an unhandled IOException.
|
||||
throw new ResponseStatusException(
|
||||
HttpStatus.SERVICE_UNAVAILABLE, "AI engine unreachable: " + e.getMessage(), e);
|
||||
} catch (InterruptedException e) {
|
||||
|
||||
+144
-4
@@ -1,31 +1,42 @@
|
||||
package stirling.software.proprietary.service;
|
||||
|
||||
import static org.assertj.core.api.Assertions.assertThat;
|
||||
import static org.junit.jupiter.api.Assertions.assertEquals;
|
||||
import static org.junit.jupiter.api.Assertions.assertSame;
|
||||
import static org.junit.jupiter.api.Assertions.assertThrows;
|
||||
import static org.mockito.ArgumentMatchers.any;
|
||||
import static org.mockito.Mockito.mock;
|
||||
import static org.mockito.Mockito.verify;
|
||||
import static org.mockito.Mockito.when;
|
||||
|
||||
import java.io.IOException;
|
||||
import java.net.ConnectException;
|
||||
import java.net.http.HttpClient;
|
||||
import java.net.http.HttpRequest;
|
||||
import java.net.http.HttpResponse;
|
||||
import java.net.http.HttpTimeoutException;
|
||||
import java.util.List;
|
||||
import java.util.stream.Stream;
|
||||
|
||||
import org.junit.jupiter.api.AfterEach;
|
||||
import org.junit.jupiter.api.BeforeEach;
|
||||
import org.junit.jupiter.api.Test;
|
||||
import org.mockito.ArgumentCaptor;
|
||||
import org.springframework.http.HttpStatus;
|
||||
import org.springframework.security.authentication.AnonymousAuthenticationToken;
|
||||
import org.springframework.security.authentication.UsernamePasswordAuthenticationToken;
|
||||
import org.springframework.security.core.authority.SimpleGrantedAuthority;
|
||||
import org.springframework.security.core.context.SecurityContextHolder;
|
||||
import org.springframework.security.core.userdetails.User;
|
||||
import org.springframework.security.core.userdetails.UserDetails;
|
||||
import org.springframework.web.server.ResponseStatusException;
|
||||
|
||||
import stirling.software.common.model.ApplicationProperties;
|
||||
|
||||
/**
|
||||
* Verifies that AiEngineClient surfaces network-layer failures as structured HTTP statuses so every
|
||||
* AI tool caller sees a consistent, meaningful error rather than a raw 500.
|
||||
*/
|
||||
class AiEngineClientTest {
|
||||
|
||||
private static final String ENGINE_SECRET = "shared-engine-secret";
|
||||
|
||||
private ApplicationProperties applicationProperties;
|
||||
private HttpClient httpClient;
|
||||
private AiEngineClient client;
|
||||
@@ -38,6 +49,12 @@ class AiEngineClientTest {
|
||||
applicationProperties.getAiEngine().setTimeoutSeconds(5);
|
||||
httpClient = mock(HttpClient.class);
|
||||
client = new AiEngineClient(applicationProperties, httpClient);
|
||||
SecurityContextHolder.clearContext();
|
||||
}
|
||||
|
||||
@AfterEach
|
||||
void tearDown() {
|
||||
SecurityContextHolder.clearContext();
|
||||
}
|
||||
|
||||
@Test
|
||||
@@ -84,4 +101,127 @@ class AiEngineClientTest {
|
||||
|
||||
assertEquals(HttpStatus.SERVICE_UNAVAILABLE, ex.getStatusCode());
|
||||
}
|
||||
|
||||
// --- decorate() contract: X-Engine-Auth + X-User-Id stamping ---------------------------
|
||||
|
||||
@Test
|
||||
void postStampsEngineAuthAndUserIdHeaders() throws Exception {
|
||||
applicationProperties.getCluster().getEngine().setSharedSecret(ENGINE_SECRET);
|
||||
authenticateAs("alice");
|
||||
stubOkResponse();
|
||||
|
||||
client.post("/x", "{}");
|
||||
|
||||
HttpRequest sent = captureSentRequest();
|
||||
assertThat(sent.headers().firstValue("X-Engine-Auth")).contains(ENGINE_SECRET);
|
||||
assertThat(sent.headers().firstValue("X-User-Id")).contains("alice");
|
||||
}
|
||||
|
||||
@Test
|
||||
void getStampsEngineAuthAndUserIdHeaders() throws Exception {
|
||||
applicationProperties.getCluster().getEngine().setSharedSecret(ENGINE_SECRET);
|
||||
authenticateAs("bob");
|
||||
stubOkResponse();
|
||||
|
||||
client.get("/x");
|
||||
|
||||
HttpRequest sent = captureSentRequest();
|
||||
assertThat(sent.headers().firstValue("X-Engine-Auth")).contains(ENGINE_SECRET);
|
||||
assertThat(sent.headers().firstValue("X-User-Id")).contains("bob");
|
||||
}
|
||||
|
||||
@Test
|
||||
void streamPostStampsEngineAuthAndUserIdHeaders() throws Exception {
|
||||
applicationProperties.getCluster().getEngine().setSharedSecret(ENGINE_SECRET);
|
||||
applicationProperties.getAiEngine().setLongRunningTimeoutSeconds(60);
|
||||
authenticateAs("carol");
|
||||
@SuppressWarnings("unchecked")
|
||||
HttpResponse<Stream<String>> response =
|
||||
(HttpResponse<Stream<String>>) mock(HttpResponse.class);
|
||||
when(response.statusCode()).thenReturn(200);
|
||||
when(response.body()).thenReturn(Stream.empty());
|
||||
when(httpClient.send(any(), any(HttpResponse.BodyHandler.class))).thenReturn(response);
|
||||
|
||||
client.streamPost("/stream", "{}", line -> {});
|
||||
|
||||
HttpRequest sent = captureSentRequest();
|
||||
assertThat(sent.headers().firstValue("X-Engine-Auth")).contains(ENGINE_SECRET);
|
||||
assertThat(sent.headers().firstValue("X-User-Id")).contains("carol");
|
||||
}
|
||||
|
||||
@Test
|
||||
void postOmitsEngineAuthWhenSharedSecretBlank() throws Exception {
|
||||
applicationProperties.getCluster().getEngine().setSharedSecret("");
|
||||
authenticateAs("alice");
|
||||
stubOkResponse();
|
||||
|
||||
client.post("/x", "{}");
|
||||
|
||||
HttpRequest sent = captureSentRequest();
|
||||
assertThat(sent.headers().firstValue("X-Engine-Auth"))
|
||||
.as("blank engine secret must NOT produce a header")
|
||||
.isEmpty();
|
||||
assertThat(sent.headers().firstValue("X-User-Id")).contains("alice");
|
||||
}
|
||||
|
||||
@Test
|
||||
void postOmitsUserIdHeaderWhenSecurityContextEmpty() throws Exception {
|
||||
// Unauthenticated context: X-User-Id must be omitted; X-Engine-Auth still applies.
|
||||
applicationProperties.getCluster().getEngine().setSharedSecret(ENGINE_SECRET);
|
||||
SecurityContextHolder.clearContext();
|
||||
stubOkResponse();
|
||||
|
||||
client.post("/x", "{}");
|
||||
|
||||
HttpRequest sent = captureSentRequest();
|
||||
assertThat(sent.headers().firstValue("X-Engine-Auth")).contains(ENGINE_SECRET);
|
||||
assertThat(sent.headers().firstValue("X-User-Id"))
|
||||
.as("anonymous context must NOT stamp an X-User-Id header")
|
||||
.isEmpty();
|
||||
}
|
||||
|
||||
@Test
|
||||
void postOmitsUserIdHeaderForAnonymousAuthenticationToken() throws Exception {
|
||||
// AnonymousAuthenticationToken#isAuthenticated() returns true and getName() ==
|
||||
// "anonymousUser". resolveUserId() must special-case it so the engine never sees
|
||||
// anonymous traffic conflated under the literal identity "anonymousUser".
|
||||
applicationProperties.getCluster().getEngine().setSharedSecret(ENGINE_SECRET);
|
||||
SecurityContextHolder.getContext()
|
||||
.setAuthentication(
|
||||
new AnonymousAuthenticationToken(
|
||||
"key",
|
||||
"anonymousUser",
|
||||
List.of(new SimpleGrantedAuthority("ROLE_ANONYMOUS"))));
|
||||
stubOkResponse();
|
||||
|
||||
client.post("/x", "{}");
|
||||
|
||||
HttpRequest sent = captureSentRequest();
|
||||
assertThat(sent.headers().firstValue("X-User-Id"))
|
||||
.as("anonymous principal must NOT stamp an X-User-Id header")
|
||||
.isEmpty();
|
||||
}
|
||||
|
||||
private void authenticateAs(String username) {
|
||||
UserDetails principal =
|
||||
User.withUsername(username).password("n/a").authorities("USER").build();
|
||||
SecurityContextHolder.getContext()
|
||||
.setAuthentication(
|
||||
new UsernamePasswordAuthenticationToken(
|
||||
principal, "n/a", principal.getAuthorities()));
|
||||
}
|
||||
|
||||
@SuppressWarnings("unchecked")
|
||||
private void stubOkResponse() throws Exception {
|
||||
HttpResponse<String> response = (HttpResponse<String>) mock(HttpResponse.class);
|
||||
when(response.statusCode()).thenReturn(200);
|
||||
when(response.body()).thenReturn("{}");
|
||||
when(httpClient.send(any(), any(HttpResponse.BodyHandler.class))).thenReturn(response);
|
||||
}
|
||||
|
||||
private HttpRequest captureSentRequest() throws Exception {
|
||||
ArgumentCaptor<HttpRequest> captor = ArgumentCaptor.forClass(HttpRequest.class);
|
||||
verify(httpClient).send(captor.capture(), any(HttpResponse.BodyHandler.class));
|
||||
return captor.getValue();
|
||||
}
|
||||
}
|
||||
|
||||
@@ -38,8 +38,8 @@ import stirling.software.proprietary.security.database.repository.UserRepository
|
||||
import stirling.software.proprietary.security.model.User;
|
||||
import stirling.software.saas.ai.model.AiCreateSession;
|
||||
import stirling.software.saas.ai.repository.AiCreateSessionRepository;
|
||||
import stirling.software.saas.ai.service.AiCreateProxyService;
|
||||
import stirling.software.saas.ai.service.AiCreateSessionService;
|
||||
import stirling.software.saas.ai.service.AiProxyService;
|
||||
import stirling.software.saas.service.CreditService;
|
||||
import stirling.software.saas.service.TeamCreditService;
|
||||
import stirling.software.saas.util.AuthenticationUtils;
|
||||
@@ -53,7 +53,7 @@ import stirling.software.saas.util.CreditHeaderUtils;
|
||||
public class AiCreateController {
|
||||
|
||||
private final AiCreateSessionService sessionService;
|
||||
private final AiCreateProxyService proxyService;
|
||||
private final AiProxyService proxyService;
|
||||
private final ObjectMapper objectMapper = new ObjectMapper();
|
||||
private final CreditService creditService;
|
||||
private final TeamCreditService teamCreditService;
|
||||
|
||||
@@ -1,145 +0,0 @@
|
||||
package stirling.software.saas.ai.service;
|
||||
|
||||
import java.io.IOException;
|
||||
import java.io.InputStream;
|
||||
import java.io.UncheckedIOException;
|
||||
import java.net.URI;
|
||||
import java.net.http.HttpClient;
|
||||
import java.net.http.HttpRequest;
|
||||
import java.net.http.HttpResponse;
|
||||
|
||||
import org.springframework.beans.factory.annotation.Value;
|
||||
import org.springframework.context.annotation.Profile;
|
||||
import org.springframework.stereotype.Service;
|
||||
|
||||
import jakarta.servlet.http.HttpServletRequest;
|
||||
|
||||
import lombok.extern.slf4j.Slf4j;
|
||||
|
||||
import stirling.software.proprietary.security.database.repository.UserRepository;
|
||||
import stirling.software.proprietary.security.service.UserService;
|
||||
|
||||
@Service
|
||||
@Profile("saas")
|
||||
@Slf4j
|
||||
public class AiCreateProxyService {
|
||||
|
||||
private static final String DEFAULT_AI_BASE_URL = "http://localhost:5001";
|
||||
|
||||
private final String aiServiceBaseUrl;
|
||||
|
||||
private final HttpClient httpClient;
|
||||
private final UserRepository userRepository;
|
||||
private final UserService userService;
|
||||
|
||||
public AiCreateProxyService(
|
||||
@Value("${app.ai.service-base-url:" + DEFAULT_AI_BASE_URL + "}")
|
||||
String aiServiceBaseUrl,
|
||||
UserRepository userRepository,
|
||||
UserService userService) {
|
||||
this.aiServiceBaseUrl = aiServiceBaseUrl;
|
||||
this.httpClient = HttpClient.newBuilder().build();
|
||||
this.userRepository = userRepository;
|
||||
this.userService = userService;
|
||||
}
|
||||
|
||||
public HttpResponse<InputStream> forward(
|
||||
String method, String path, HttpServletRequest request, boolean acceptEventStream)
|
||||
throws IOException, InterruptedException {
|
||||
String targetUrl = buildTargetUrl(path, request.getQueryString());
|
||||
HttpRequest.Builder builder = HttpRequest.newBuilder(URI.create(targetUrl));
|
||||
|
||||
String contentType = request.getContentType();
|
||||
if (contentType != null && !contentType.isBlank()) {
|
||||
builder.header("Content-Type", contentType);
|
||||
}
|
||||
|
||||
String authorization = request.getHeader("Authorization");
|
||||
if (authorization != null && !authorization.isBlank()) {
|
||||
builder.header("Authorization", authorization);
|
||||
}
|
||||
|
||||
// Extract user API key from authenticated user and forward to AI backend
|
||||
String apiKey = request.getHeader("X-API-KEY");
|
||||
if (apiKey == null || apiKey.isBlank()) {
|
||||
apiKey = extractUserApiKey();
|
||||
}
|
||||
if (apiKey != null && !apiKey.isBlank()) {
|
||||
builder.header("X-API-KEY", apiKey);
|
||||
log.debug("Forwarding X-API-KEY header to AI backend");
|
||||
}
|
||||
|
||||
String accept = request.getHeader("Accept");
|
||||
if (acceptEventStream) {
|
||||
builder.header("Accept", "text/event-stream");
|
||||
} else if (accept != null && !accept.isBlank()) {
|
||||
builder.header("Accept", accept);
|
||||
}
|
||||
|
||||
builder.method(method, buildBodyPublisher(method, request));
|
||||
log.debug("Proxying AI create request {} {}", method, targetUrl);
|
||||
return httpClient.send(builder.build(), HttpResponse.BodyHandlers.ofInputStream());
|
||||
}
|
||||
|
||||
private String buildTargetUrl(String path, String queryString) {
|
||||
String baseUrl = aiServiceBaseUrl;
|
||||
if (baseUrl == null || baseUrl.isBlank()) {
|
||||
baseUrl = DEFAULT_AI_BASE_URL;
|
||||
}
|
||||
baseUrl = baseUrl.trim();
|
||||
if (baseUrl.endsWith("/")) {
|
||||
baseUrl = baseUrl.substring(0, baseUrl.length() - 1);
|
||||
}
|
||||
if (!path.startsWith("/")) {
|
||||
path = "/" + path;
|
||||
}
|
||||
String url = baseUrl + path;
|
||||
if (queryString != null && !queryString.isBlank()) {
|
||||
url += "?" + queryString;
|
||||
}
|
||||
return url;
|
||||
}
|
||||
|
||||
private HttpRequest.BodyPublisher buildBodyPublisher(
|
||||
String method, HttpServletRequest request) {
|
||||
if ("GET".equalsIgnoreCase(method) || "DELETE".equalsIgnoreCase(method)) {
|
||||
return HttpRequest.BodyPublishers.noBody();
|
||||
}
|
||||
return HttpRequest.BodyPublishers.ofInputStream(
|
||||
() -> {
|
||||
try {
|
||||
return request.getInputStream();
|
||||
} catch (IOException exc) {
|
||||
throw new UncheckedIOException(exc);
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
/**
|
||||
* Extract the authenticated user's API key from the database. If the user doesn't have an API
|
||||
* key, one will be created automatically.
|
||||
*
|
||||
* @return The user's API key, or null if not authenticated or key creation fails
|
||||
*/
|
||||
private String extractUserApiKey() {
|
||||
try {
|
||||
// Use getCurrentUsername() which handles all auth types including anonymous users
|
||||
String username = userService.getCurrentUsername();
|
||||
if (username == null || username.isBlank()) {
|
||||
log.debug("No authenticated user found for API key extraction");
|
||||
return null;
|
||||
}
|
||||
|
||||
// getApiKeyForUser will create a key if it doesn't exist
|
||||
String apiKey = userService.getApiKeyForUser(username);
|
||||
log.debug("Retrieved API key for user: {}", username);
|
||||
return apiKey;
|
||||
} catch (Exception e) {
|
||||
log.error(
|
||||
"Failed to extract or create user API key for user: {}",
|
||||
userService.getCurrentUsername(),
|
||||
e);
|
||||
return null;
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -20,9 +20,20 @@ import jakarta.servlet.http.Part;
|
||||
|
||||
import lombok.extern.slf4j.Slf4j;
|
||||
|
||||
import stirling.software.common.model.ApplicationProperties;
|
||||
import stirling.software.proprietary.security.database.repository.UserRepository;
|
||||
import stirling.software.proprietary.security.service.UserService;
|
||||
|
||||
/**
|
||||
* Proxies saas-side HTTP requests to the AI engine and enforces the auth-header contract:
|
||||
*
|
||||
* <ul>
|
||||
* <li>Client-supplied {@code Authorization} and {@code X-API-KEY} are stripped; the server-
|
||||
* resolved API key is stamped so callers cannot spoof identity downstream.
|
||||
* <li>{@code X-Engine-Auth} is stamped when the cluster shared secret is configured.
|
||||
* <li>{@code X-User-Id} is stamped for authenticated principals; omitted for anonymous.
|
||||
* </ul>
|
||||
*/
|
||||
@Service
|
||||
@Profile("saas")
|
||||
@Slf4j
|
||||
@@ -30,21 +41,29 @@ public class AiProxyService {
|
||||
|
||||
private static final String DEFAULT_AI_BASE_URL = "http://localhost:5001";
|
||||
|
||||
private static final String ENGINE_AUTH_HEADER = "X-Engine-Auth";
|
||||
private static final String USER_ID_HEADER = "X-User-Id";
|
||||
private static final String API_KEY_HEADER = "X-API-KEY";
|
||||
private static final String AUTHORIZATION_HEADER = "Authorization";
|
||||
|
||||
private final String aiServiceBaseUrl;
|
||||
|
||||
private final HttpClient httpClient;
|
||||
private final UserRepository userRepository;
|
||||
private final UserService userService;
|
||||
private final ApplicationProperties applicationProperties;
|
||||
|
||||
public AiProxyService(
|
||||
@Value("${app.ai.service-base-url:" + DEFAULT_AI_BASE_URL + "}")
|
||||
String aiServiceBaseUrl,
|
||||
UserRepository userRepository,
|
||||
UserService userService) {
|
||||
UserService userService,
|
||||
ApplicationProperties applicationProperties) {
|
||||
this.aiServiceBaseUrl = aiServiceBaseUrl;
|
||||
this.httpClient = HttpClient.newBuilder().build();
|
||||
this.userRepository = userRepository;
|
||||
this.userService = userService;
|
||||
this.applicationProperties = applicationProperties;
|
||||
}
|
||||
|
||||
public HttpResponse<InputStream> forward(
|
||||
@@ -55,21 +74,6 @@ public class AiProxyService {
|
||||
|
||||
String contentType = request.getContentType();
|
||||
|
||||
String authorization = request.getHeader("Authorization");
|
||||
if (authorization != null && !authorization.isBlank()) {
|
||||
builder.header("Authorization", authorization);
|
||||
}
|
||||
|
||||
// Extract user API key from authenticated user and forward to AI backend
|
||||
String apiKey = request.getHeader("X-API-KEY");
|
||||
if (apiKey == null || apiKey.isBlank()) {
|
||||
apiKey = extractUserApiKey();
|
||||
}
|
||||
if (apiKey != null && !apiKey.isBlank()) {
|
||||
builder.header("X-API-KEY", apiKey);
|
||||
log.debug("Forwarding X-API-KEY header to AI backend");
|
||||
}
|
||||
|
||||
String accept = request.getHeader("Accept");
|
||||
if (acceptEventStream) {
|
||||
builder.header("Accept", "text/event-stream");
|
||||
@@ -84,10 +88,68 @@ public class AiProxyService {
|
||||
builder.header("Content-Type", contentType);
|
||||
}
|
||||
builder.method(method, body.publisher);
|
||||
|
||||
// INVARIANT: stamp auth headers LAST. setHeader() overrides any previously-set value,
|
||||
// including any caller-supplied Authorization / X-API-KEY that arrived via this proxy.
|
||||
// Any .header(...) call AFTER this point would defeat the strip-on-output guarantee.
|
||||
stampAuthHeaders(builder);
|
||||
|
||||
log.debug("Proxying AI request {} {}", method, targetUrl);
|
||||
return httpClient.send(builder.build(), HttpResponse.BodyHandlers.ofInputStream());
|
||||
}
|
||||
|
||||
private void stampAuthHeaders(HttpRequest.Builder builder) {
|
||||
// setHeader (not header) with empty string causes the JDK builder to omit the header
|
||||
// entirely in the built request, dropping any caller-supplied value.
|
||||
builder.setHeader(AUTHORIZATION_HEADER, "");
|
||||
|
||||
String apiKey = extractUserApiKey();
|
||||
if (apiKey != null && !apiKey.isBlank()) {
|
||||
builder.setHeader(API_KEY_HEADER, apiKey);
|
||||
log.debug("Attaching server-resolved X-API-KEY for the engine call");
|
||||
} else {
|
||||
builder.setHeader(API_KEY_HEADER, "");
|
||||
}
|
||||
|
||||
String engineSecret = resolveEngineSecret();
|
||||
if (engineSecret != null && !engineSecret.isBlank()) {
|
||||
builder.setHeader(ENGINE_AUTH_HEADER, engineSecret);
|
||||
}
|
||||
|
||||
try {
|
||||
String username = userService.getCurrentUsername();
|
||||
if (username != null && !username.isBlank() && !"anonymousUser".equals(username)) {
|
||||
builder.setHeader(USER_ID_HEADER, username);
|
||||
}
|
||||
} catch (RuntimeException ex) {
|
||||
log.warn("Could not resolve current username for X-User-Id header", ex);
|
||||
}
|
||||
}
|
||||
|
||||
private String resolveEngineSecret() {
|
||||
ApplicationProperties.Cluster cluster = applicationProperties.getCluster();
|
||||
if (cluster == null || cluster.getEngine() == null) {
|
||||
return "";
|
||||
}
|
||||
return cluster.getEngine().getSharedSecret();
|
||||
}
|
||||
|
||||
private String extractUserApiKey() {
|
||||
try {
|
||||
String username = userService.getCurrentUsername();
|
||||
if (username == null || username.isBlank()) {
|
||||
log.debug("No authenticated user found for API key extraction");
|
||||
return null;
|
||||
}
|
||||
String apiKey = userService.getApiKeyForUser(username);
|
||||
log.debug("Retrieved API key for user: {}", username);
|
||||
return apiKey;
|
||||
} catch (Exception e) {
|
||||
log.error("Failed to extract or create user API key", e);
|
||||
return null;
|
||||
}
|
||||
}
|
||||
|
||||
private String buildTargetUrl(String path, String queryString) {
|
||||
String baseUrl = aiServiceBaseUrl;
|
||||
if (baseUrl == null || baseUrl.isBlank()) {
|
||||
@@ -176,34 +238,6 @@ public class AiProxyService {
|
||||
output.write("\r\n".getBytes(StandardCharsets.UTF_8));
|
||||
}
|
||||
|
||||
/**
|
||||
* Extract the authenticated user's API key from the database. If the user doesn't have an API
|
||||
* key, one will be created automatically.
|
||||
*
|
||||
* @return The user's API key, or null if not authenticated or key creation fails
|
||||
*/
|
||||
private String extractUserApiKey() {
|
||||
try {
|
||||
// Use getCurrentUsername() which handles all auth types including anonymous users
|
||||
String username = userService.getCurrentUsername();
|
||||
if (username == null || username.isBlank()) {
|
||||
log.debug("No authenticated user found for API key extraction");
|
||||
return null;
|
||||
}
|
||||
|
||||
// getApiKeyForUser will create a key if it doesn't exist
|
||||
String apiKey = userService.getApiKeyForUser(username);
|
||||
log.debug("Retrieved API key for user: {}", username);
|
||||
return apiKey;
|
||||
} catch (Exception e) {
|
||||
log.error(
|
||||
"Failed to extract or create user API key for user: {}",
|
||||
userService.getCurrentUsername(),
|
||||
e);
|
||||
return null;
|
||||
}
|
||||
}
|
||||
|
||||
private static class BodyPublisherWithContentType {
|
||||
private final HttpRequest.BodyPublisher publisher;
|
||||
private final String contentType;
|
||||
|
||||
@@ -0,0 +1,157 @@
|
||||
package stirling.software.saas.ai.service;
|
||||
|
||||
import static org.assertj.core.api.Assertions.assertThat;
|
||||
import static org.mockito.ArgumentMatchers.any;
|
||||
import static org.mockito.Mockito.lenient;
|
||||
import static org.mockito.Mockito.when;
|
||||
|
||||
import java.io.ByteArrayInputStream;
|
||||
import java.io.InputStream;
|
||||
import java.net.http.HttpClient;
|
||||
import java.net.http.HttpRequest;
|
||||
import java.net.http.HttpResponse;
|
||||
import java.util.Optional;
|
||||
|
||||
import org.junit.jupiter.api.BeforeEach;
|
||||
import org.junit.jupiter.api.Test;
|
||||
import org.junit.jupiter.api.extension.ExtendWith;
|
||||
import org.junit.jupiter.params.ParameterizedTest;
|
||||
import org.junit.jupiter.params.provider.ValueSource;
|
||||
import org.mockito.ArgumentCaptor;
|
||||
import org.mockito.Mock;
|
||||
import org.mockito.junit.jupiter.MockitoExtension;
|
||||
import org.springframework.mock.web.MockHttpServletRequest;
|
||||
import org.springframework.test.util.ReflectionTestUtils;
|
||||
|
||||
import stirling.software.common.model.ApplicationProperties;
|
||||
import stirling.software.proprietary.security.database.repository.UserRepository;
|
||||
import stirling.software.proprietary.security.service.UserService;
|
||||
|
||||
/** Verifies the auth-header strip-and-stamp contract on the saas AI proxy. */
|
||||
@ExtendWith(MockitoExtension.class)
|
||||
class AiProxyServiceTest {
|
||||
|
||||
private static final String ENGINE_URL = "http://engine.local:5001";
|
||||
private static final String ENGINE_SECRET = "shared-engine-secret";
|
||||
private static final String SERVER_API_KEY = "server-resolved-api-key";
|
||||
private static final String CURRENT_USERNAME = "alice";
|
||||
|
||||
@Mock UserRepository userRepository;
|
||||
@Mock UserService userService;
|
||||
@Mock HttpClient httpClient;
|
||||
@Mock HttpResponse<InputStream> mockResponse;
|
||||
|
||||
private ApplicationProperties applicationProperties;
|
||||
private AiProxyService proxy;
|
||||
|
||||
@BeforeEach
|
||||
void setUp() throws Exception {
|
||||
applicationProperties = new ApplicationProperties();
|
||||
applicationProperties.getCluster().getEngine().setSharedSecret(ENGINE_SECRET);
|
||||
|
||||
proxy = new AiProxyService(ENGINE_URL, userRepository, userService, applicationProperties);
|
||||
// Replace the internally-constructed HttpClient with our mock so we can capture requests.
|
||||
ReflectionTestUtils.setField(proxy, "httpClient", httpClient);
|
||||
|
||||
lenient().when(mockResponse.body()).thenReturn(new ByteArrayInputStream(new byte[0]));
|
||||
lenient().when(mockResponse.statusCode()).thenReturn(200);
|
||||
when(httpClient.send(any(HttpRequest.class), any(HttpResponse.BodyHandler.class)))
|
||||
.thenReturn(mockResponse);
|
||||
}
|
||||
|
||||
@ParameterizedTest
|
||||
@ValueSource(strings = {"/api/v1/chat", "/api/create/sessions"})
|
||||
void forward_dropsClientAuthorizationAndApiKeyHeaders(String path) throws Exception {
|
||||
when(userService.getCurrentUsername()).thenReturn(CURRENT_USERNAME);
|
||||
when(userService.getApiKeyForUser(CURRENT_USERNAME)).thenReturn(SERVER_API_KEY);
|
||||
|
||||
MockHttpServletRequest inbound = new MockHttpServletRequest("POST", path);
|
||||
inbound.addHeader("Authorization", "Bearer client-supplied-bearer");
|
||||
inbound.addHeader("X-API-KEY", "client-supplied-api-key");
|
||||
inbound.setContentType("application/json");
|
||||
inbound.setContent("{}".getBytes());
|
||||
|
||||
proxy.forward("POST", path, inbound, false);
|
||||
|
||||
HttpRequest outbound = captureOutboundRequest();
|
||||
|
||||
// setHeader with empty string causes the JDK builder to omit the header or send it empty.
|
||||
assertThat(outbound.headers().firstValue("Authorization").orElse(""))
|
||||
.as("client Authorization must never reach the engine")
|
||||
.isEmpty();
|
||||
assertThat(outbound.headers().allValues("Authorization"))
|
||||
.as("client Authorization value must not be among the outbound values")
|
||||
.doesNotContain("Bearer client-supplied-bearer");
|
||||
|
||||
Optional<String> apiKey = outbound.headers().firstValue("X-API-KEY");
|
||||
assertThat(apiKey)
|
||||
.as("X-API-KEY must be the server-resolved value, not the client-supplied one")
|
||||
.contains(SERVER_API_KEY);
|
||||
assertThat(apiKey.orElse(""))
|
||||
.as("client-supplied X-API-KEY must not leak through")
|
||||
.isNotEqualTo("client-supplied-api-key");
|
||||
assertThat(outbound.headers().allValues("X-API-KEY"))
|
||||
.as("client X-API-KEY must not be among the outbound values")
|
||||
.doesNotContain("client-supplied-api-key");
|
||||
}
|
||||
|
||||
@ParameterizedTest
|
||||
@ValueSource(strings = {"/api/v1/chat", "/api/create/sessions"})
|
||||
void forward_stampsEngineAuthAndUserIdHeaders(String path) throws Exception {
|
||||
when(userService.getCurrentUsername()).thenReturn(CURRENT_USERNAME);
|
||||
when(userService.getApiKeyForUser(CURRENT_USERNAME)).thenReturn(SERVER_API_KEY);
|
||||
|
||||
MockHttpServletRequest inbound = new MockHttpServletRequest("POST", path);
|
||||
inbound.setContentType("application/json");
|
||||
inbound.setContent("{}".getBytes());
|
||||
|
||||
proxy.forward("POST", path, inbound, false);
|
||||
|
||||
HttpRequest outbound = captureOutboundRequest();
|
||||
|
||||
assertThat(outbound.headers().firstValue("X-Engine-Auth"))
|
||||
.as("engine shared secret must be stamped on outbound requests")
|
||||
.contains(ENGINE_SECRET);
|
||||
assertThat(outbound.headers().firstValue("X-User-Id"))
|
||||
.as("authenticated username must be stamped on outbound requests")
|
||||
.contains(CURRENT_USERNAME);
|
||||
}
|
||||
|
||||
@Test
|
||||
void forward_omitsEngineAuthWhenSecretBlank() throws Exception {
|
||||
applicationProperties.getCluster().getEngine().setSharedSecret("");
|
||||
when(userService.getCurrentUsername()).thenReturn(CURRENT_USERNAME);
|
||||
when(userService.getApiKeyForUser(CURRENT_USERNAME)).thenReturn(SERVER_API_KEY);
|
||||
|
||||
MockHttpServletRequest inbound = new MockHttpServletRequest("GET", "/api/v1/health");
|
||||
|
||||
proxy.forward("GET", "/api/v1/health", inbound, false);
|
||||
|
||||
HttpRequest outbound = captureOutboundRequest();
|
||||
assertThat(outbound.headers().firstValue("X-Engine-Auth"))
|
||||
.as("blank engine secret must not produce a header")
|
||||
.isEmpty();
|
||||
}
|
||||
|
||||
@Test
|
||||
void forward_omitsUserIdForAnonymousPrincipal() throws Exception {
|
||||
when(userService.getCurrentUsername()).thenReturn("anonymousUser");
|
||||
|
||||
MockHttpServletRequest inbound =
|
||||
new MockHttpServletRequest("GET", "/api/create/sessions/abc");
|
||||
|
||||
proxy.forward("GET", "/api/create/sessions/abc", inbound, false);
|
||||
|
||||
HttpRequest outbound = captureOutboundRequest();
|
||||
assertThat(outbound.headers().firstValue("X-User-Id"))
|
||||
.as("anonymous principal must not produce an X-User-Id header")
|
||||
.isEmpty();
|
||||
}
|
||||
|
||||
private HttpRequest captureOutboundRequest() throws Exception {
|
||||
ArgumentCaptor<HttpRequest> captor = ArgumentCaptor.forClass(HttpRequest.class);
|
||||
org.mockito.Mockito.verify(httpClient)
|
||||
.send(captor.capture(), any(HttpResponse.BodyHandler.class));
|
||||
return captor.getValue();
|
||||
}
|
||||
}
|
||||
@@ -1,9 +1,9 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from contextlib import asynccontextmanager
|
||||
from typing import Annotated
|
||||
|
||||
from fastapi import Depends, FastAPI
|
||||
from fastapi import FastAPI
|
||||
from pydantic_ai import Agent
|
||||
from pydantic_ai.models.instrumented import InstrumentationSettings
|
||||
|
||||
@@ -16,7 +16,7 @@ from stirling.agents import (
|
||||
)
|
||||
from stirling.agents.ledger import MathAuditorAgent
|
||||
from stirling.agents.pdf_comment import PdfCommentAgent
|
||||
from stirling.api.middleware import UserIdMiddleware
|
||||
from stirling.api.middleware import EngineAuthMiddleware, UserIdMiddleware
|
||||
from stirling.api.routes import (
|
||||
agent_draft_router,
|
||||
document_router,
|
||||
@@ -63,7 +63,21 @@ async def lifespan(fast_api: FastAPI):
|
||||
|
||||
|
||||
app = FastAPI(title="Stirling AI Engine", lifespan=lifespan, version="0.1.0")
|
||||
|
||||
try:
|
||||
_engine_shared_secret = load_settings().engine_shared_secret or ""
|
||||
except (AttributeError, KeyError) as cfg_err:
|
||||
raise RuntimeError(
|
||||
"engine_shared_secret missing from settings; ensure STIRLING_ENGINE_SHARED_SECRET "
|
||||
"is declared in the env (blank value is allowed for dev mode)."
|
||||
) from cfg_err
|
||||
if not _engine_shared_secret:
|
||||
logging.getLogger(__name__).warning(
|
||||
"STIRLING_ENGINE_SHARED_SECRET is blank - running in dev (open) mode."
|
||||
)
|
||||
|
||||
app.add_middleware(UserIdMiddleware)
|
||||
app.add_middleware(EngineAuthMiddleware, expected_secret=_engine_shared_secret)
|
||||
app.include_router(orchestrator_router)
|
||||
app.include_router(pdf_edit_router)
|
||||
app.include_router(pdf_question_router)
|
||||
@@ -75,9 +89,5 @@ app.include_router(pdf_comments_router)
|
||||
|
||||
|
||||
@app.get("/health", response_model=HealthResponse)
|
||||
async def healthcheck(settings: Annotated[AppSettings, Depends(load_settings)]) -> HealthResponse:
|
||||
return HealthResponse(
|
||||
status="ok",
|
||||
smart_model=settings.smart_model_name,
|
||||
fast_model=settings.fast_model_name,
|
||||
)
|
||||
async def healthcheck() -> HealthResponse:
|
||||
return HealthResponse(status="ok")
|
||||
|
||||
@@ -1,16 +1,20 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import hmac
|
||||
|
||||
from starlette.middleware.base import BaseHTTPMiddleware, RequestResponseEndpoint
|
||||
from starlette.requests import Request
|
||||
from starlette.responses import Response
|
||||
from starlette.responses import JSONResponse, Response
|
||||
|
||||
from stirling.services.tracking import current_user_id
|
||||
|
||||
_USER_ID_HEADER = "X-User-Id"
|
||||
_ENGINE_AUTH_HEADER = "X-Engine-Auth"
|
||||
_HEALTH_PATHS = {"/health", "/healthz", "/readyz"}
|
||||
|
||||
|
||||
class UserIdMiddleware(BaseHTTPMiddleware):
|
||||
"""Extract X-User-Id header and set it as the current user for PostHog tracking."""
|
||||
"""Set X-User-Id (stamped by the trusted Java proxy) as request context."""
|
||||
|
||||
async def dispatch(self, request: Request, call_next: RequestResponseEndpoint) -> Response:
|
||||
user_id = request.headers.get(_USER_ID_HEADER)
|
||||
@@ -21,3 +25,21 @@ class UserIdMiddleware(BaseHTTPMiddleware):
|
||||
finally:
|
||||
current_user_id.reset(token)
|
||||
return await call_next(request)
|
||||
|
||||
|
||||
class EngineAuthMiddleware(BaseHTTPMiddleware):
|
||||
"""Validate shared-secret header. Blank secret = dev mode (open); health probes exempt."""
|
||||
|
||||
def __init__(self, app, expected_secret: str) -> None:
|
||||
super().__init__(app)
|
||||
self._expected_secret = expected_secret or ""
|
||||
# Precompute the bytes form once so per-request work is just the constant-time compare.
|
||||
self._expected_secret_bytes = self._expected_secret.encode("utf-8")
|
||||
|
||||
async def dispatch(self, request: Request, call_next: RequestResponseEndpoint) -> Response:
|
||||
if request.url.path in _HEALTH_PATHS or not self._expected_secret:
|
||||
return await call_next(request)
|
||||
presented = request.headers.get(_ENGINE_AUTH_HEADER, "")
|
||||
if not hmac.compare_digest(presented.encode("utf-8"), self._expected_secret_bytes):
|
||||
return JSONResponse({"detail": "engine authentication failed"}, status_code=401)
|
||||
return await call_next(request)
|
||||
|
||||
@@ -98,6 +98,9 @@ class AppSettings(BaseSettings):
|
||||
posthog_api_key: str = Field(validation_alias="STIRLING_POSTHOG_API_KEY")
|
||||
posthog_host: str = Field(validation_alias="STIRLING_POSTHOG_HOST")
|
||||
|
||||
# Engine shared-secret for cluster deployments. Blank = dev mode (open).
|
||||
engine_shared_secret: str = Field(default="", validation_alias="STIRLING_ENGINE_SHARED_SECRET")
|
||||
|
||||
|
||||
def _configure_logging(level_name: str, log_file: str, http_debug: bool) -> None:
|
||||
"""Configure the ``stirling`` logger hierarchy."""
|
||||
|
||||
@@ -5,5 +5,3 @@ from stirling.models import ApiModel
|
||||
|
||||
class HealthResponse(ApiModel):
|
||||
status: str
|
||||
smart_model: str
|
||||
fast_model: str
|
||||
|
||||
@@ -0,0 +1,59 @@
|
||||
"""Tests for EngineAuthMiddleware - shared-secret gating and health probe exemption."""
|
||||
from __future__ import annotations
|
||||
|
||||
from fastapi import FastAPI
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
from stirling.api.middleware import EngineAuthMiddleware
|
||||
|
||||
|
||||
def _app(expected_secret: str) -> FastAPI:
|
||||
app = FastAPI()
|
||||
app.add_middleware(EngineAuthMiddleware, expected_secret=expected_secret)
|
||||
|
||||
@app.get("/health")
|
||||
async def health():
|
||||
return {"status": "ok"}
|
||||
|
||||
@app.get("/v1/agent")
|
||||
async def agent():
|
||||
return {"ok": True}
|
||||
|
||||
return app
|
||||
|
||||
|
||||
def test_correct_secret_allows_request():
|
||||
client = TestClient(_app("s3cret"))
|
||||
res = client.get("/v1/agent", headers={"X-Engine-Auth": "s3cret"})
|
||||
assert res.status_code == 200
|
||||
assert res.json() == {"ok": True}
|
||||
|
||||
|
||||
def test_missing_header_rejected_with_401():
|
||||
client = TestClient(_app("s3cret"))
|
||||
res = client.get("/v1/agent")
|
||||
assert res.status_code == 401
|
||||
|
||||
|
||||
def test_wrong_header_rejected_with_401():
|
||||
client = TestClient(_app("s3cret"))
|
||||
res = client.get("/v1/agent", headers={"X-Engine-Auth": "wrong"})
|
||||
assert res.status_code == 401
|
||||
|
||||
|
||||
def test_blank_secret_dev_mode_allows_unauthenticated():
|
||||
client = TestClient(_app(""))
|
||||
res = client.get("/v1/agent")
|
||||
assert res.status_code == 200
|
||||
|
||||
|
||||
def test_health_endpoint_exempt_from_auth():
|
||||
client = TestClient(_app("s3cret"))
|
||||
res = client.get("/health")
|
||||
assert res.status_code == 200
|
||||
|
||||
|
||||
def test_health_endpoint_exempt_in_dev_mode():
|
||||
client = TestClient(_app(""))
|
||||
res = client.get("/health")
|
||||
assert res.status_code == 200
|
||||
@@ -0,0 +1,50 @@
|
||||
"""Tests for UserIdMiddleware - X-User-Id propagates as context.
|
||||
|
||||
X-Tenant-Id is deliberately NOT accepted from the wire: the Java proxy does not stamp
|
||||
one and a client-supplied tenant id would be spoofable. If/when tenant scoping arrives
|
||||
it must be derived from the server-side security context, not from request headers.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
from fastapi import FastAPI, Request
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
from stirling.api.middleware import UserIdMiddleware
|
||||
from stirling.services.tracking import current_user_id
|
||||
|
||||
|
||||
def _app() -> FastAPI:
|
||||
app = FastAPI()
|
||||
app.add_middleware(UserIdMiddleware)
|
||||
|
||||
@app.get("/me")
|
||||
async def me(request: Request):
|
||||
return {
|
||||
"user_id": current_user_id.get(),
|
||||
"tenant_id": getattr(request.state, "tenant_id", None),
|
||||
}
|
||||
|
||||
return app
|
||||
|
||||
|
||||
def test_user_id_header_is_propagated():
|
||||
client = TestClient(_app())
|
||||
res = client.get("/me", headers={"X-User-Id": "alice"})
|
||||
assert res.status_code == 200
|
||||
assert res.json()["user_id"] == "alice"
|
||||
|
||||
|
||||
def test_tenant_id_header_is_ignored():
|
||||
client = TestClient(_app())
|
||||
res = client.get("/me", headers={"X-User-Id": "alice", "X-Tenant-Id": "acme"})
|
||||
assert res.status_code == 200
|
||||
body = res.json()
|
||||
assert body["user_id"] == "alice"
|
||||
assert body["tenant_id"] is None
|
||||
|
||||
|
||||
def test_missing_user_id_returns_no_user_in_context():
|
||||
client = TestClient(_app())
|
||||
res = client.get("/me")
|
||||
assert res.status_code == 200
|
||||
assert res.json()["user_id"] in ("", None)
|
||||
Reference in New Issue
Block a user