RockteMQ-AI commented on code in PR #4665:
URL:
https://github.com/apache/rocketmq-dashboard/pull/4665#discussion_r4059283508
##########
server/src/main/java/org/apache/rocketmq/studio/ops/ai/tool/service/ToolTokenService.java:
##########
@@ -91,24 +104,46 @@ private static byte[] resolveSecret(String
configuredSecret) {
}
}
+ /**
+ * Issues a confirmation token for a previewed mutation. Each call embeds
a fresh random
+ * token id, so two previews of the same operation never share an
identifier and consuming
+ * one never invalidates the other.
+ */
public String issue(ToolExecutionContext context) {
requireConfiguredSecret();
long expiresAt = clock.instant().plus(TOKEN_TTL).getEpochSecond();
- byte[] signature = sign(signingPayload(context, expiresAt));
- return new ConfirmationToken(expiresAt, signature).format();
+ String tokenId = newTokenId();
+ byte[] signature = sign(signingPayload(context, expiresAt, tokenId));
+ return new ConfirmationToken(expiresAt, tokenId, signature).format();
}
- public void verify(ToolExecutionContext context) {
+ /**
+ * Validates the confirmation token and atomically consumes it before the
caller may enter
+ * the non-idempotent mutation. Rejections stay distinguishable: an
expired, tampered or
+ * mismatched token is {@code CONFIRMATION_TOKEN_INVALID}; a token that
already admitted an
+ * execution is {@code CONFIRMATION_TOKEN_ALREADY_USED}. Consumption
happens only after all
+ * other checks pass, so a rejected replay never consumes a different
pending token.
+ */
+ public void verifyAndConsume(ToolExecutionContext context) {
requireConfiguredSecret();
String toolName = context.definition().name();
ConfirmationToken token =
ConfirmationToken.parse(context.confirmToken(), toolName);
if (clock.instant().getEpochSecond() >= token.expiresAt()) {
Review Comment:
Correct check ordering: expiry then signature then consumed. This ensures
that an expired token is rejected as `INVALID_ARGUMENT` without leaking whether
it was also already consumed, and a tampered token is rejected before the
consumed-state lookup. Nice defensive design.
##########
server/src/main/java/org/apache/rocketmq/studio/ops/ai/tool/service/ConsumedTokenStore.java:
##########
@@ -0,0 +1,56 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.rocketmq.studio.ops.ai.tool.service;
+
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ConcurrentMap;
+
+/**
+ * Records which confirmation token ids have already been consumed so a token
admits at most
+ * one non-idempotent mutation. The store lives in process memory to match the
single-node
+ * deployment model of the confirmation flow; it introduces no external
dependency.
+ *
+ * <p>Entries are kept until the token itself expires: a replay inside the TTL
is rejected as
+ * already used, and a replay afterwards is rejected as expired by the token
verification, so
+ * dropping expired entries never widens the acceptance window. Purging is
opportunistic to
+ * keep the amortized cost of an apply at O(1).
+ */
+final class ConsumedTokenStore {
+
+ /** Purge expired entries once the map grows past this bound; each entry
is one id plus one long. */
+ private static final int PURGE_THRESHOLD = 4096;
+
+ private final ConcurrentMap<String, Long> consumedExpiries = new
ConcurrentHashMap<>();
+
+ /**
+ * Atomically marks the token id as consumed.
+ *
+ * @return {@code true} when this call is the first consumer, {@code
false} when the id was
+ * already consumed and the caller must be rejected
+ */
+ boolean consume(String tokenId, long expiresAtEpochSecond, long
nowEpochSecond) {
Review Comment:
Good use of `putIfAbsent` for atomic single-use enforcement. The
opportunistic purge after the consume is a clean design - expired entries can
never widen the acceptance window because the expiry check in
`verifyAndConsume` runs first.
##########
server/src/main/java/org/apache/rocketmq/studio/ops/ai/tool/service/ToolTokenService.java:
##########
@@ -66,9 +73,15 @@ public ToolTokenService(
}
ToolTokenService(ObjectMapper objectMapper, Clock clock, byte[] secret) {
+ this(objectMapper, clock, secret, new SecureRandom());
Review Comment:
Good testability: the `Random` parameter allows tests to inject a
deterministic RNG while production uses `SecureRandom`. The 16-hex-char token
id (8 bytes / 64 bits) provides sufficient collision resistance for single-use
tokens within the 10-minute TTL window.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
