NeitzelFXDesigner/designer-contracts/src/main/java/de/neitzel/fxdesigner/protocol/ProtocolSession.java

package de.neitzel.fxdesigner.protocol;

import static de.neitzel.fxdesigner.protocol.WireProtocol.*;

import java.net.*;
import java.nio.charset.StandardCharsets;
import java.security.MessageDigest;
import java.security.SecureRandom;
import java.time.*;
import java.util.*;

/**
 * Host-side authentication/negotiation and bounded request correlation, independent of sockets/UI.
 */
public final class ProtocolSession implements AutoCloseable {
  private enum State {
    NEW,
    READY,
    CLOSED
  }

  private State state = State.NEW;
  private final String session, token;
  private final long workerPid;
  private final Set<String> supported, required;
  private Set<Operation> operations = Set.of();

  private record Pending(Packet request, Instant deadline) {}

  private final Map<String, Pending> pending = new HashMap<>();
  private long sequence;

  public ProtocolSession(
      String session, long workerPid, String token, Set<String> supported, Set<String> required) {
    if (session == null
        || session.isBlank()
        || workerPid < 1
        || token == null
        || !token.matches("[0-9a-f]{64}"))
      throw new IllegalArgumentException("Invalid launch identity");
    this.session = session;
    this.workerPid = workerPid;
    this.token = token;
    this.supported = Set.copyOf(supported);
    this.required = Set.copyOf(required);
  }

  public static String newToken() {
    byte[] bytes = new byte[32];
    new SecureRandom().nextBytes(bytes);
    return HexFormat.of().formatHex(bytes);
  }

  public synchronized Set<Operation> authenticate(InetAddress peer, Hello hello) {
    if (state != State.NEW) throw new IllegalStateException("Handshake already consumed");
    state = State.CLOSED;
    if (!peer.isLoopbackAddress()
        || !session.equals(hello.session())
        || workerPid != hello.workerPid()
        || !MessageDigest.isEqual(
            token.getBytes(StandardCharsets.US_ASCII),
            hello.token().getBytes(StandardCharsets.UTF_8)))
      throw new IllegalArgumentException("Worker authentication failed");
    if (hello.wire().major() != 1)
      throw new IllegalArgumentException("Unsupported wire major version");
    if (!supported.containsAll(hello.required()) || !hello.capabilities().containsAll(required))
      throw new IllegalArgumentException("Required capability unavailable");
    var usable = EnumSet.noneOf(Operation.class);
    if (hello.metadata().major() == 1
        && supported.contains("discover-v1")
        && hello.capabilities().contains("discover-v1")) usable.add(Operation.DISCOVER);
    if (hello.preview().major() == 1
        && supported.contains("render-rgba-v1")
        && hello.capabilities().contains("render-rgba-v1")) usable.add(Operation.RENDER);
    if (hello.preview().major() == 1
        && supported.contains("layout-v1")
        && hello.capabilities().contains("layout-v1")) {
      usable.add(Operation.MEASURE);
      usable.add(Operation.OCCURRENCES);
    }
    if (hello.preview().major() == 1
        && supported.contains("interactive-v1")
        && hello.capabilities().contains("interactive-v1")) usable.add(Operation.INTERACTIVE);
    if (hello.metadata().major() == 1
        && supported.contains("bindings-v2")
        && hello.capabilities().contains("bindings-v2")) usable.add(Operation.BINDINGS);
    var requiredByEither = new HashSet<>(required);
    requiredByEither.addAll(hello.required());
    if ((requiredByEither.contains("discover-v1") && !usable.contains(Operation.DISCOVER))
        || (requiredByEither.contains("render-rgba-v1") && !usable.contains(Operation.RENDER)))
      throw new IllegalArgumentException("Required capability version incompatible");
    if ((requiredByEither.contains("layout-v1") && !usable.contains(Operation.MEASURE))
        || (requiredByEither.contains("interactive-v1") && !usable.contains(Operation.INTERACTIVE)))
      throw new IllegalArgumentException("Required capability version incompatible");
    if (requiredByEither.contains("bindings-v2") && !usable.contains(Operation.BINDINGS))
      throw new IllegalArgumentException("Required capability version incompatible");
    operations = Set.copyOf(usable);
    state = State.READY;
    return operations;
  }

  public synchronized Welcome welcome() {
    ready();
    return new Welcome(session, new Version(1, 0), operations);
  }

  public synchronized Packet request(
      Operation operation,
      String context,
      String snapshot,
      Map<String, Long> revisions,
      int width,
      int height,
      byte[] payload,
      Instant now,
      Duration timeout) {
    ready();
    if (!operations.contains(operation))
      throw new IllegalStateException("Operation not negotiated");
    if (timeout.isNegative() || timeout.isZero() || timeout.compareTo(Duration.ofMinutes(5)) > 0)
      throw new IllegalArgumentException("Invalid request timeout");
    expire(now);
    if (pending.size() >= MAX_PENDING) throw new IllegalStateException("Too many pending requests");
    var packet =
        new Packet(
            Kind.REQUEST,
            operation,
            new Identity(
                session, Long.toString(++sequence), context, snapshot, revisions, width, height),
            payload);
    pending.put(packet.identity().request(), new Pending(packet, now.plus(timeout)));
    return packet;
  }

  /** False means stale/unrelated/mismatched; such results must not mutate the UI. */
  public synchronized boolean accept(Packet result, Instant now) {
    ready();
    expire(now);
    if (result.kind() != Kind.RESULT && result.kind() != Kind.ERROR) return false;
    var p = pending.get(result.identity().request());
    if (p == null
        || p.request.operation() != result.operation()
        || !p.request.identity().equals(result.identity())) return false;
    pending.remove(result.identity().request());
    return true;
  }

  public synchronized Optional<Packet> cancel(String request) {
    ready();
    var p = pending.remove(request);
    return p == null
        ? Optional.empty()
        : Optional.of(
            new Packet(Kind.CANCEL, p.request.operation(), p.request.identity(), new byte[0]));
  }

  /**
   * On any render input/size/snapshot change the caller cancels the superseded request before
   * issuing another.
   */
  public synchronized List<Packet> expire(Instant now) {
    var cancelled = new ArrayList<Packet>();
    var iterator = pending.values().iterator();
    while (iterator.hasNext()) {
      var p = iterator.next();
      if (!now.isBefore(p.deadline)) {
        cancelled.add(
            new Packet(Kind.CANCEL, p.request.operation(), p.request.identity(), new byte[0]));
        iterator.remove();
      }
    }
    return List.copyOf(cancelled);
  }

  private void ready() {
    if (state != State.READY)
      throw new IllegalStateException("Session not authenticated or closed");
  }

  public synchronized void close() {
    state = State.CLOSED;
    pending.clear();
  }
}