001/*
002 * Licensed to the Apache Software Foundation (ASF) under one or more
003 * contributor license agreements.  See the NOTICE file distributed with
004 * this work for additional information regarding copyright ownership.
005 * The ASF licenses this file to You under the Apache License, Version 2.0
006 * (the "License"); you may not use this file except in compliance with
007 * the License.  You may obtain a copy of the License at
008 *
009 *      http://www.apache.org/licenses/LICENSE-2.0
010 *
011 * Unless required by applicable law or agreed to in writing, software
012 * distributed under the License is distributed on an "AS IS" BASIS,
013 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
014 * See the License for the specific language governing permissions and
015 * limitations under the License.
016 */
017package org.apache.juneau.rest.client.mcp.auth;
018
019import static org.apache.juneau.commons.utils.Shorts.*;
020
021import java.net.*;
022import java.util.*;
023import java.util.concurrent.*;
024
025/**
026 * Client-side SEP-2350 step-up state machine: wraps an MCP call so that, on a {@code 401}/{@code 403 insufficient_scope}
027 * challenge, it computes the required-scope union ({@link McpScopeAccumulator}), drives (re-)authorization via a
028 * caller-supplied {@link Reauthorizer}, and retries the original call — bounded by a hard attempt cap with
029 * per-<code>(resource, operation)</code> attempt tracking so a misconfigured server can't drive an unbounded
030 * re-authorization loop (spec "no more than a few times").
031 *
032 * <p>
033 * Step-up for user-delegation clients is inherently interactive (a fresh auth-code + PKCE round-trip), so this is an
034 * <b>explicit orchestration helper</b> the caller wraps around MCP calls &mdash; not a silent interceptor retry.  The
035 * caller's {@link ScopedCall} performs the actual MCP request and throws {@link McpInsufficientScopeException} (carrying
036 * the parsed {@link WwwAuthenticateChallenge}) when the server responds with a step-up challenge; the
037 * {@link Reauthorizer} performs the (re-)authorization with the widened scope union (e.g. via
038 * {@link McpAuthorizationCodeAcquirer} for user clients or a client-credentials re-acquire) and installs the new token.
039 *
040 * <p>
041 * {@code client_credentials} (own-behalf) clients MAY abort immediately rather than attempt step-up
042 * ({@link Builder#attemptStepUp(boolean) attemptStepUp(false)}).
043 *
044 * <h5 class='topic'>Attempt-counter lifetime</h5>
045 * <p>
046 * Per-<code>(resource, operation)</code> attempt counts exist only to bound a single step-up burst.  A counter is
047 * <b>cleared automatically</b> when {@link #execute(String, ScopedCall) execute} returns successfully for that operation,
048 * so a long-lived authorizer does not permanently refuse a later legitimate step-up for the same operation.  A counter
049 * that reaches the cap without a success (a permanent failure) is <b>not</b> auto-cleared &mdash; the caller can reset it
050 * explicitly with {@link #reset(String)} (one operation) or {@link #resetAll()} (all operations) once the underlying
051 * cause is resolved.
052 *
053 * @since 10.0.0
054 */
055@SuppressWarnings({
056   "java:S1192" // Duplicated literals (argument/property names) read more clearly inline than as constants
057})
058public class McpStepUpAuthorizer {
059
060   /** A caller-supplied MCP call that throws {@link McpInsufficientScopeException} on a step-up challenge. */
061   @FunctionalInterface
062   public interface ScopedCall<T> {
063      /**
064       * Performs the MCP call.
065       *
066       * @return The call result.
067       * @throws McpInsufficientScopeException If the server responded with a step-up {@code insufficient_scope} challenge.
068       */
069      T call() throws McpInsufficientScopeException;
070   }
071
072   /** Drives (re-)authorization with the computed scope union and installs the resulting token. */
073   @FunctionalInterface
074   public interface Reauthorizer {
075      /**
076       * Performs (re-)authorization with the widened scope union.
077       *
078       * @param unionScopes The order-preserving union of previously-requested + challenge scopes.
079       */
080      void reauthorize(Set<String> unionScopes);
081   }
082
083   /**
084    * Static creator.
085    *
086    * @return A new builder.
087    */
088   public static Builder create() {
089      return new Builder();
090   }
091
092   /**
093    * Builder.
094    */
095   public static class Builder {
096      URI resource;
097      URI issuer;
098      McpScopeAccumulator accumulator;
099      Reauthorizer reauthorizer;
100      int maxAttempts = 3;
101      boolean attemptStepUp = true;
102      final Set<String> previouslyRequestedScopes = st();
103
104      /** Constructor. */
105      protected Builder() {}
106
107      /**
108       * Sets the RFC 8707 resource indicator this authorizer tracks attempts for.  Required.
109       *
110       * @param value The resource URI.
111       * @return This object.
112       */
113      public Builder resource(URI value) {
114         resource = reqnn("value", value);
115         return this;
116      }
117
118      /**
119       * Sets the authorization server issuer.  Required.
120       *
121       * @param value The issuer URI.
122       * @return This object.
123       */
124      public Builder issuer(URI value) {
125         issuer = reqnn("value", value);
126         return this;
127      }
128
129      /**
130       * Sets the scope accumulator.  Defaults to a fresh {@link McpScopeAccumulator}.
131       *
132       * @param value The accumulator.  Must not be <jk>null</jk>.
133       * @return This object.
134       */
135      public Builder accumulator(McpScopeAccumulator value) {
136         accumulator = reqnn("value", value);
137         return this;
138      }
139
140      /**
141       * Sets the re-authorization callback invoked with the widened scope union.  Required.
142       *
143       * @param value The reauthorizer.  Must not be <jk>null</jk>.
144       * @return This object.
145       */
146      public Builder reauthorizer(Reauthorizer value) {
147         reauthorizer = reqnn("value", value);
148         return this;
149      }
150
151      /**
152       * Sets the hard cap on re-authorization attempts per {@code (resource, operation)} before a permanent failure.
153       * Default 3.
154       *
155       * @param value The maximum attempts.  Must be &ge; 1.
156       * @return This object.
157       */
158      public Builder maxAttempts(int value) {
159         req(value >= 1, "maxAttempts must be >= 1 (was %s)", value);
160         maxAttempts = value;
161         return this;
162      }
163
164      /**
165       * Whether to attempt step-up at all.  Default <jk>true</jk>; a {@code client_credentials} client MAY set
166       * <jk>false</jk> to abort immediately on the first challenge.
167       *
168       * @param value <jk>false</jk> to abort immediately rather than attempt step-up.
169       * @return This object.
170       */
171      public Builder attemptStepUp(boolean value) {
172         attemptStepUp = value;
173         return this;
174      }
175
176      /**
177       * Seeds the scope accumulator with the client's previously-requested scopes so the first union preserves them.
178       *
179       * @param values The previously-requested scopes.
180       * @return This object.
181       */
182      public Builder previouslyRequestedScopes(String... values) {
183         reqnn("values", values);
184         for (var v : values) {
185            reqnb("scope", v);
186            previouslyRequestedScopes.add(v);
187         }
188         return this;
189      }
190
191      /**
192       * Builds the authorizer.
193       *
194       * @return A new {@link McpStepUpAuthorizer}.
195       */
196      public McpStepUpAuthorizer build() {
197         if (resource == null)
198            throw isex("McpStepUpAuthorizer requires resource(...)");
199         if (issuer == null)
200            throw isex("McpStepUpAuthorizer requires issuer(...)");
201         if (reauthorizer == null)
202            throw isex("McpStepUpAuthorizer requires reauthorizer(...)");
203         if (accumulator == null)
204            accumulator = new McpScopeAccumulator();
205         if (! previouslyRequestedScopes.isEmpty())
206            accumulator.seed(resource, issuer, previouslyRequestedScopes);
207         return new McpStepUpAuthorizer(this);
208      }
209   }
210
211   private final URI resource;
212   private final URI issuer;
213   private final McpScopeAccumulator accumulator;
214   private final Reauthorizer reauthorizer;
215   private final int maxAttempts;
216   private final boolean attemptStepUp;
217   private final ConcurrentMap<String,Integer> attemptsByOperation = new ConcurrentHashMap<>();
218
219   /**
220    * Constructor.
221    *
222    * @param b The builder.
223    */
224   protected McpStepUpAuthorizer(Builder b) {
225      this.resource = b.resource;
226      this.issuer = b.issuer;
227      this.accumulator = b.accumulator;
228      this.reauthorizer = b.reauthorizer;
229      this.maxAttempts = b.maxAttempts;
230      this.attemptStepUp = b.attemptStepUp;
231   }
232
233   /**
234    * Executes an MCP call, transparently handling SEP-2350 step-up: on an {@code insufficient_scope} challenge it
235    * unions the challenge scopes onto the accumulated set, re-authorizes, and retries &mdash; up to the configured
236    * attempt cap, after which it throws a permanent {@link McpAuthException}.
237    *
238    * @param <T> The call result type.
239    * @param operation The operation identifier (e.g. {@code tools/call:deploy}) used for per-operation attempt
240    *    tracking.  Must not be <jk>null</jk> or blank.
241    * @param call The MCP call.  Must not be <jk>null</jk>.
242    * @return The successful call result.
243    * @throws McpAuthException If step-up is disabled, or the attempt cap is exceeded (permanent authorization failure).
244    */
245   public <T> T execute(String operation, ScopedCall<T> call) {
246      reqnb("operation", operation);
247      reqnn("call", call);
248      var opKey = opKey(operation);
249      while (true) {
250         try {
251            var result = call.call();
252            // Success clears the operation's attempt burst so a later legitimate step-up starts fresh.
253            attemptsByOperation.remove(opKey);
254            return result;
255         } catch (McpInsufficientScopeException e) {
256            if (! attemptStepUp)
257               throw new McpAuthException("Step-up authorization aborted for operation '" + operation
258                  + "' (attemptStepUp=false)", e);
259            var used = attemptsByOperation.merge(opKey, 1, Integer::sum);
260            if (used > maxAttempts)
261               throw new McpAuthException("Step-up authorization exhausted after " + maxAttempts
262                  + " attempt(s) for resource '" + resource + "' operation '" + operation
263                  + "'; treating as a permanent authorization failure", e);
264            var union = accumulator.accumulate(resource, issuer, e.challenge().scopes());
265            reauthorizer.reauthorize(union);
266         }
267      }
268   }
269
270   /**
271    * Returns the number of step-up attempts recorded for the given operation.
272    *
273    * @param operation The operation identifier.  Must not be <jk>null</jk> or blank.
274    * @return The recorded attempt count.
275    */
276   public int attempts(String operation) {
277      reqnb("operation", operation);
278      return attemptsByOperation.getOrDefault(opKey(operation), 0);
279   }
280
281   /**
282    * Clears the recorded step-up attempt count for a single operation (e.g. after the caller has resolved the cause of
283    * a prior permanent failure), so a subsequent {@link #execute(String, ScopedCall)} may step up again.
284    *
285    * @param operation The operation identifier.  Must not be <jk>null</jk> or blank.
286    */
287   public void reset(String operation) {
288      reqnb("operation", operation);
289      attemptsByOperation.remove(opKey(operation));
290   }
291
292   /**
293    * Clears all recorded step-up attempt counts for this authorizer's resource across every operation.
294    */
295   public void resetAll() {
296      attemptsByOperation.clear();
297   }
298
299   /** The per-operation attempt-tracking key: {@code resource + '\n' + operation}. */
300   private String opKey(String operation) {
301      return resource.toString() + '\n' + operation;
302   }
303}