feat: add companion service discovery

CHANGELOG.md

@@ -14,6 +14,7 @@ versioned entry and uses it as the GitHub release notes.
 ### Added
 
 - Add `background` and `transparent` options to Ditaa diagrams, to set the background colour of the image or make it transparent
+- Add companion service discovery: a companion container not built into Kroki can now register itself as a new diagram type via `POST /services` (with a heartbeat to stay registered, since the registry is in-memory and does not survive a restart). Disabled by default, opt in with `KROKI_ENABLE_COMPANION_DISCOVERY`; secure the registration API with `KROKI_COMPANION_REGISTRATION_TOKEN`. Always disabled while `KROKI_SAFE_MODE` is `SECURE` (the default, including on kroki.io) ([#1423](https://github.com/yuzutech/kroki/issues/1423))
 
 ### Security
 

docs/modules/setup/pages/configuration.adoc

@@ -314,6 +314,88 @@ KROKI_EXCALIDRAW_PORT:: Port of the Excalidraw container (default: `8004`).
 
 NOTE: If you are using the default `docker-compose.yaml` file you can rely on the default values.
 
+== Companion Service Discovery
+
+In addition to the companion containers built into Kroki (Mermaid, BPMN, Excalidraw, diagrams.net), you can extend Kroki with your own companion container for a diagram type that isn't built in. A companion registers itself with the gateway over a small REST API, and Kroki then delegates conversion requests for that diagram type to it — see https://github.com/yuzutech/kroki/issues/1423[issue #1423] for the design discussion.
+
+This feature is disabled by default. Enable it with `KROKI_ENABLE_COMPANION_DISCOVERY=true`.
+
+[IMPORTANT]
+====
+Companion service discovery is always disabled while Kroki runs in `SECURE` safe mode (the default, see <<Safe Mode>>), regardless of `KROKI_ENABLE_COMPANION_DISCOVERY`. `SECURE` means the operator does not trust a service's own security checks, which is exactly what registering a new, unvetted companion would ask the gateway to do. This also means the feature is never available on the public https://kroki.io instance.
+
+To use this feature, explicitly set `KROKI_SAFE_MODE` to `SAFE` or `UNSAFE`.
+====
+
+=== Registering a companion
+
+A companion registers itself by sending its name, version, supported output formats, and where it can be reached:
+
+[source,bash]
+----
+curl -X POST https://my.kroki.example/services \
+  -H 'Content-Type: application/json' \
+  -d '{
+    "name": "mscgen",
+    "version": "1.2.3",
+    "formats": ["png", "svg"],
+    "host": "mscgen-companion",
+    "port": 8080
+  }'
+----
+
+`name`:: The diagram type name, used in the conversion URL (`/mscgen/svg/...`). Must be lowercase letters and digits only, and not already used by a built-in diagram type or another registered companion (`422` otherwise).
+`version`:: A free-form version string, reported on the `/health` endpoint and the homepage.
+`formats`:: A non-empty array of the output formats the companion supports (`png`, `svg`, `jpeg`, `pdf`, `base64`, `txt` or `utxt`).
+`host` and `port`:: Where Kroki should reach the companion. Loopback, link-local (which covers every major cloud provider's instance-metadata address) and private addresses are rejected, as well as a small built-in denylist of cloud metadata hostnames — see `KROKI_COMPANION_BLOCKED_HOSTS` below.
+
+On success, the server responds `201 Created` and the diagram type is immediately available, e.g. `POST /mscgen/svg`. The companion is expected to implement the same wire protocol as Kroki's own companion containers: `POST /<name>/<format>` with the decoded diagram source as the request body.
+
+Since Kroki is stateless, the registry does not survive a restart, and a companion is expected to keep renewing its registration with a heartbeat:
+
+[source,bash]
+----
+curl -X PUT https://my.kroki.example/services/mscgen/heartbeat
+----
+
+A companion that misses its heartbeat deadline (`KROKI_COMPANION_HEARTBEAT_TTL_MS`, default `90000` — 90 seconds) is automatically evicted, freeing up its name. A companion can also unregister explicitly on graceful shutdown:
+
+[source,bash]
+----
+curl -X DELETE https://my.kroki.example/services/mscgen
+----
+
+`GET /services/:name` and `GET /services` return the current registration(s) as JSON, for inspection.
+
+=== Securing the registration API
+
+Since registering a companion lets it serve conversion requests under a new diagram type name, the `/services` API should only be reachable from a trusted network (e.g. a private Docker or Kubernetes network with no route from the outside).
+
+As an additional layer, you can require a bearer token on every `/services` request with `KROKI_COMPANION_REGISTRATION_TOKEN`:
+
+[source,bash]
+----
+KROKI_COMPANION_REGISTRATION_TOKEN=s3cr3t
+----
+
+[source,bash]
+----
+curl -X POST https://my.kroki.example/services \
+  -H 'Authorization: Bearer s3cr3t' \
+  -H 'Content-Type: application/json' \
+  -d '{ "name": "mscgen", "version": "1.2.3", "formats": ["svg"], "host": "mscgen-companion", "port": 8080 }'
+----
+
+If `KROKI_ENABLE_COMPANION_DISCOVERY` is set without a token, Kroki logs a warning on startup: the registration API is then reachable by anyone who can reach the gateway.
+
+=== Companion Service Discovery Environment Variables
+
+`KROKI_ENABLE_COMPANION_DISCOVERY`:: Enables the `/services` registration API. Defaults to `false`. Has no effect while `KROKI_SAFE_MODE` is `SECURE`.
+`KROKI_COMPANION_REGISTRATION_TOKEN`:: Bearer token required on every `/services` request. Optional, but strongly recommended.
+`KROKI_COMPANION_HEARTBEAT_TTL_MS`:: How long a companion may go without a heartbeat before being evicted. Defaults to `90000` (90 seconds).
+`KROKI_COMPANION_HEARTBEAT_SWEEP_INTERVAL_MS`:: How often expired registrations are checked for. Defaults to a third of `KROKI_COMPANION_HEARTBEAT_TTL_MS`, with a floor of `5000`.
+`KROKI_COMPANION_BLOCKED_HOSTS`:: Comma-separated list of additional hostnames to reject as a companion's `host`, on top of the built-in denylist (loopback, link-local, private addresses, and known cloud metadata hostnames).
+
 == Max URI length
 
 Some diagrams, like Excalidraw, have verbose textual descriptions that will produce long URI.

server/src/main/java/io/kroki/server/Server.java

@@ -5,6 +5,10 @@
 import io.kroki.server.error.ErrorHandler;
 import io.kroki.server.error.InvalidRequestHandler;
 import io.kroki.server.log.Logging;
+import io.kroki.server.registry.CompanionAuthHandler;
+import io.kroki.server.registry.CompanionRegistry;
+import io.kroki.server.registry.CompanionServiceHandler;
+import io.kroki.server.security.SafeMode;
 import io.kroki.server.service.*;
 import io.vertx.config.ConfigRetriever;
 import io.vertx.core.*;
@@ -20,6 +24,7 @@
 import io.vertx.ext.web.handler.BodyHandler;
 import io.vertx.ext.web.handler.CorsHandler;
 
+import java.time.Duration;
 import java.util.*;
 import java.util.stream.Collectors;
 
@@ -119,6 +124,8 @@ static Future<HttpServer> start(Vertx vertx, VertxOptions vertxOptions, JsonObje
     registry.register(new Wireviz(vertx, config, commander), "wireviz");
     registry.register(new Goat(vertx, config, commander), "goat");
 
+    mountCompanionDiscovery(vertx, config, router, bodyHandler, delegator, registry);
+
     router.post("/")
       .handler(bodyHandler)
       .handler(new DiagramRest(registry).create());
@@ -130,7 +137,9 @@ static Future<HttpServer> start(Vertx vertx, VertxOptions vertxOptions, JsonObje
     router.get("/metrics")
       .handler(metricHandlerService);
     // health
-    HealthHandler healthHandler = new HealthHandler(registry.getVersions(), blockedThreadChecker);
+    // queries registry.getVersions() live (rather than a startup snapshot) so that companion
+    // services registered or evicted at runtime (see issue #1423) are reflected immediately
+    HealthHandler healthHandler = new HealthHandler(registry::getVersions, blockedThreadChecker);
     Handler<RoutingContext> healthHandlerService = healthHandler.create();
     router.get("/health")
       .handler(healthHandlerService);
@@ -140,14 +149,16 @@ static Future<HttpServer> start(Vertx vertx, VertxOptions vertxOptions, JsonObje
       .handler(healthHandlerService);
 
     // hello
-    List<ServiceVersion> serviceVersions = healthHandler.getServiceVersions();
     String krokiBuildHash = healthHandler.getKrokiBuildHash();
     String krokiVersionNumber = healthHandler.getKrokiVersionNumber();
     router.get("/")
-      .handler(new HelloHandler(vertx, serviceVersions, krokiVersionNumber, krokiBuildHash).create());
+      .handler(new HelloHandler(vertx, healthHandler::getServiceVersions, krokiVersionNumber, krokiBuildHash).create());
 
     // Default route
-    Route route = router.route("/*");
+    // Ordered last so that companion services registered dynamically at runtime (see issue
+    // #1423), whose routes are necessarily added to the router after this one, still get a
+    // chance to match instead of always falling through to this catch-all 404.
+    Route route = router.route("/*").order(Integer.MAX_VALUE);
     route.handler(routingContext -> routingContext.fail(404));
     ErrorHandler errorHandler = new ErrorHandler(vertx, config.getBoolean("KROKI_DISPLAY_EXCEPTION_DETAILS", false));
     route.failureHandler(errorHandler);
@@ -158,6 +169,71 @@ static Future<HttpServer> start(Vertx vertx, VertxOptions vertxOptions, JsonObje
       .listen(getListenAddress(config));
   }
 
+  /**
+   * Mounts the companion service discovery REST API (see issue #1423) under {@code /services},
+   * allowing companion containers not built into Kroki to register themselves as a new diagram
+   * type. Disabled by default: this opens a new registration surface that should only be exposed
+   * on a trusted network, ideally with {@code KROKI_COMPANION_REGISTRATION_TOKEN} configured.
+   *
+   * <p>Always disabled under {@code KROKI_SAFE_MODE=SECURE} (the default, including on kroki.io),
+   * regardless of {@code KROKI_ENABLE_COMPANION_DISCOVERY}: SECURE means the operator does not
+   * trust arbitrary services' own security checks, which is precisely what registering a new,
+   * unvetted companion would ask the gateway to do.
+   */
+  private static void mountCompanionDiscovery(Vertx vertx, JsonObject config, Router router, BodyHandler bodyHandler, Delegator delegator, DiagramRegistry registry) {
+    if (!config.getBoolean("KROKI_ENABLE_COMPANION_DISCOVERY", false)) {
+      return;
+    }
+    SafeMode safeMode = SafeMode.get(config.getString("KROKI_SAFE_MODE", "secure"), SafeMode.SECURE);
+    if (safeMode == SafeMode.SECURE) {
+      logger.warn("KROKI_ENABLE_COMPANION_DISCOVERY is enabled but KROKI_SAFE_MODE is SECURE (the default): " +
+        "companion service discovery stays disabled. Set KROKI_SAFE_MODE to SAFE or UNSAFE to allow it.");
+      return;
+    }
+    String token = config.getString("KROKI_COMPANION_REGISTRATION_TOKEN");
+    if (token == null || token.isEmpty()) {
+      logger.warn("KROKI_ENABLE_COMPANION_DISCOVERY is enabled without KROKI_COMPANION_REGISTRATION_TOKEN: " +
+        "the /services registration API is unauthenticated, only expose it on a trusted network.");
+    }
+    long heartbeatTtlMs = config.getLong("KROKI_COMPANION_HEARTBEAT_TTL_MS", 90_000L);
+
+    // additional hostnames to reject as a companion's host, on top of the built-in cloud metadata denylist
+    Set<String> extraBlockedHosts = new LinkedHashSet<>();
+    String blockedHostsVar = config.getString("KROKI_COMPANION_BLOCKED_HOSTS");
+    if (blockedHostsVar != null) {
+      Arrays.stream(blockedHostsVar.split(","))
+        .map(String::trim)
+        .filter(s -> !s.isEmpty())
+        .forEach(extraBlockedHosts::add);
+    }
+
+    CompanionRegistry companionRegistry = new CompanionRegistry(registry, delegator, extraBlockedHosts);
+    CompanionServiceHandler serviceHandler = new CompanionServiceHandler(companionRegistry);
+
+    if (token != null && !token.isEmpty()) {
+      router.route("/services*").handler(new CompanionAuthHandler(token));
+    }
+    router.post("/services")
+      .handler(bodyHandler)
+      .handler(serviceHandler.createRegister());
+    router.put("/services/:name/heartbeat")
+      .handler(serviceHandler.createHeartbeat());
+    router.get("/services/:name")
+      .handler(serviceHandler.createGet());
+    router.get("/services")
+      .handler(serviceHandler.createList());
+    router.delete("/services/:name")
+      .handler(serviceHandler.createUnregister());
+
+    long sweepIntervalMs = config.getLong("KROKI_COMPANION_HEARTBEAT_SWEEP_INTERVAL_MS", Math.max(heartbeatTtlMs / 3, 5_000L));
+    vertx.setPeriodic(sweepIntervalMs, timerId -> {
+      List<String> expired = companionRegistry.sweepExpired(Duration.ofMillis(heartbeatTtlMs));
+      if (!expired.isEmpty()) {
+        logger.info("Evicted companion service(s) that missed their heartbeat deadline: {}", expired);
+      }
+    });
+  }
+
   private static void setPemKeyCertOptions(JsonObject config, HttpServerOptions serverOptions, boolean enableSSL) {
     if (enableSSL) {
       PemKeyCertOptions certOptions = new PemKeyCertOptions();

server/src/main/java/io/kroki/server/registry/CompanionAuthHandler.java

@@ -0,0 +1,42 @@
+package io.kroki.server.registry;
+
+import io.vertx.core.Handler;
+import io.vertx.ext.web.RoutingContext;
+
+import java.nio.charset.StandardCharsets;
+import java.security.MessageDigest;
+
+/**
+ * Guards the {@code /services} management API with a shared-secret bearer token, when configured
+ * via {@code KROKI_COMPANION_REGISTRATION_TOKEN}. Without a token configured, any client able to
+ * reach the gateway can register a companion service, which is only appropriate on a trusted
+ * network (e.g. a private Docker/Kubernetes network with no external route to the gateway).
+ */
+public class CompanionAuthHandler implements Handler<RoutingContext> {
+
+  private static final String BEARER_PREFIX = "Bearer ";
+
+  private final String token;
+
+  public CompanionAuthHandler(String token) {
+    this.token = token;
+  }
+
+  @Override
+  public void handle(RoutingContext routingContext) {
+    String authorization = routingContext.request().getHeader("Authorization");
+    if (authorization != null && authorization.startsWith(BEARER_PREFIX) && constantTimeEquals(authorization.substring(BEARER_PREFIX.length()), token)) {
+      routingContext.next();
+      return;
+    }
+    routingContext.response()
+      .setStatusCode(401)
+      .putHeader("WWW-Authenticate", "Bearer")
+      .putHeader("Content-Type", "application/json")
+      .end("{\"error\":\"Missing or invalid bearer token.\"}");
+  }
+
+  private static boolean constantTimeEquals(String a, String b) {
+    return MessageDigest.isEqual(a.getBytes(StandardCharsets.UTF_8), b.getBytes(StandardCharsets.UTF_8));
+  }
+}

server/src/main/java/io/kroki/server/registry/CompanionDiagramService.java

@@ -0,0 +1,66 @@
+package io.kroki.server.registry;
+
+import io.kroki.server.action.Delegator;
+import io.kroki.server.decode.DiagramSource;
+import io.kroki.server.decode.SourceDecoder;
+import io.kroki.server.error.DecodeException;
+import io.kroki.server.format.FileFormat;
+import io.kroki.server.service.DiagramService;
+import io.vertx.core.Future;
+import io.vertx.core.buffer.Buffer;
+import io.vertx.core.json.JsonObject;
+import io.vertx.ext.web.client.HttpResponse;
+
+import java.util.List;
+
+/**
+ * Delegates diagram conversion to a companion service registered dynamically at runtime
+ * (see issue #1423), using the same wire protocol as the built-in companion containers
+ * (e.g. {@link io.kroki.server.service.Bpmn}): {@code POST /<name>/<format>} with the
+ * decoded source as the request body.
+ */
+public class CompanionDiagramService implements DiagramService {
+
+  private final Delegator delegator;
+  private final String host;
+  private final int port;
+  private final String version;
+  private final List<FileFormat> formats;
+  private final SourceDecoder sourceDecoder;
+
+  public CompanionDiagramService(Delegator delegator, String host, int port, String version, List<FileFormat> formats) {
+    this.delegator = delegator;
+    this.host = host;
+    this.port = port;
+    this.version = version;
+    this.formats = formats;
+    this.sourceDecoder = new SourceDecoder() {
+      @Override
+      public String decode(String encoded) throws DecodeException {
+        return DiagramSource.decode(encoded);
+      }
+    };
+  }
+
+  @Override
+  public List<FileFormat> getSupportedFormats() {
+    return formats;
+  }
+
+  @Override
+  public SourceDecoder getSourceDecoder() {
+    return sourceDecoder;
+  }
+
+  @Override
+  public String getVersion() {
+    return version;
+  }
+
+  @Override
+  public Future<Buffer> convert(String sourceDecoded, String serviceName, FileFormat fileFormat, JsonObject options) {
+    String requestURI = "/" + serviceName + "/" + fileFormat.getName();
+    Future<HttpResponse<Buffer>> httpResponseFuture = this.delegator.delegate(host, port, requestURI, sourceDecoded, options);
+    return Delegator.handle(host, port, requestURI, httpResponseFuture);
+  }
+}