/*
 * Copyright 2002-2022 the original author or authors.
 *
 * Licensed 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
 *
 *      https://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.springframework.security.oauth2.client.web.server;

import java.net.URI;

import reactor.core.publisher.Mono;

import org.springframework.security.oauth2.client.ClientAuthorizationRequiredException;
import org.springframework.security.oauth2.client.registration.ClientRegistration;
import org.springframework.security.oauth2.client.registration.ClientRegistrationRepository;
import org.springframework.security.oauth2.client.registration.ReactiveClientRegistrationRepository;
import org.springframework.security.oauth2.client.web.AuthorizationRequestRepository;
import org.springframework.security.oauth2.core.AuthorizationGrantType;
import org.springframework.security.oauth2.core.endpoint.OAuth2AuthorizationRequest;
import org.springframework.security.web.server.DefaultServerRedirectStrategy;
import org.springframework.security.web.server.ServerRedirectStrategy;
import org.springframework.security.web.server.savedrequest.ServerRequestCache;
import org.springframework.security.web.server.savedrequest.WebSessionServerRequestCache;
import org.springframework.util.Assert;
import org.springframework.web.server.ServerWebExchange;
import org.springframework.web.server.WebFilter;
import org.springframework.web.server.WebFilterChain;
import org.springframework.web.util.UriComponentsBuilder;

/**
 * This {@code WebFilter} initiates the authorization code grant flow by redirecting the
 * End-User's user-agent to the Authorization Server's Authorization Endpoint.
 *
 * <p>
 * It builds the OAuth 2.0 Authorization Request, which is used as the redirect
 * {@code URI} to the Authorization Endpoint. The redirect {@code URI} will include the
 * client identifier, requested scope(s), state, response type, and a redirection URI
 * which the authorization server will send the user-agent back to once access is granted
 * (or denied) by the End-User (Resource Owner).
 *
 * <p>
 * By default, this {@code Filter} responds to authorization requests at the {@code URI}
 * {@code /oauth2/authorization/{registrationId}}. The {@code URI} template variable
 * {@code {registrationId}} represents the {@link ClientRegistration#getRegistrationId()
 * registration identifier} of the client that is used for initiating the OAuth 2.0
 * Authorization Request.
 *
 * @author Rob Winch
 * @since 5.1
 * @see OAuth2AuthorizationRequest
 * @see AuthorizationRequestRepository
 * @see ClientRegistration
 * @see ClientRegistrationRepository
 * @see <a target="_blank" href="https://tools.ietf.org/html/rfc6749#section-4.1">Section
 * 4.1 Authorization Code Grant</a>
 * @see <a target="_blank" href=
 * "https://tools.ietf.org/html/rfc6749#section-4.1.1">Section 4.1.1 Authorization Request
 * (Authorization Code)</a>
 */
public class OAuth2AuthorizationRequestRedirectWebFilter implements WebFilter {

	private ServerRedirectStrategy authorizationRedirectStrategy = new DefaultServerRedirectStrategy();

	private final ServerOAuth2AuthorizationRequestResolver authorizationRequestResolver;

	private ServerAuthorizationRequestRepository<OAuth2AuthorizationRequest> authorizationRequestRepository = new WebSessionOAuth2ServerAuthorizationRequestRepository();

	private ServerRequestCache requestCache = new WebSessionServerRequestCache();

	/**
	 * Constructs an {@code OAuth2AuthorizationRequestRedirectFilter} using the provided
	 * parameters.
	 * @param clientRegistrationRepository the repository of client registrations
	 */
	public OAuth2AuthorizationRequestRedirectWebFilter(
			ReactiveClientRegistrationRepository clientRegistrationRepository) {
		this.authorizationRequestResolver = new DefaultServerOAuth2AuthorizationRequestResolver(
				clientRegistrationRepository);
	}

	/**
	 * Constructs an {@code OAuth2AuthorizationRequestRedirectFilter} using the provided
	 * parameters.
	 * @param authorizationRequestResolver the resolver to use
	 */
	public OAuth2AuthorizationRequestRedirectWebFilter(
			ServerOAuth2AuthorizationRequestResolver authorizationRequestResolver) {
		Assert.notNull(authorizationRequestResolver, "authorizationRequestResolver cannot be null");
		this.authorizationRequestResolver = authorizationRequestResolver;
	}

	/**
	 * Sets the redirect strategy for Authorization Endpoint redirect URI.
	 * @param authorizationRedirectStrategy the redirect strategy
	 */
	public void setAuthorizationRedirectStrategy(ServerRedirectStrategy authorizationRedirectStrategy) {
		Assert.notNull(authorizationRedirectStrategy, "authorizationRedirectStrategy cannot be null");
		this.authorizationRedirectStrategy = authorizationRedirectStrategy;
	}

	/**
	 * Sets the repository used for storing {@link OAuth2AuthorizationRequest}'s.
	 * @param authorizationRequestRepository the repository used for storing
	 * {@link OAuth2AuthorizationRequest}'s
	 */
	public final void setAuthorizationRequestRepository(
			ServerAuthorizationRequestRepository<OAuth2AuthorizationRequest> authorizationRequestRepository) {
		Assert.notNull(authorizationRequestRepository, "authorizationRequestRepository cannot be null");
		this.authorizationRequestRepository = authorizationRequestRepository;
	}

	/**
	 * The request cache to use to save the request before sending a redirect.
	 * @param requestCache the cache to redirect to.
	 */
	public void setRequestCache(ServerRequestCache requestCache) {
		Assert.notNull(requestCache, "requestCache cannot be null");
		this.requestCache = requestCache;
	}

	@Override
	public Mono<Void> filter(ServerWebExchange exchange, WebFilterChain chain) {
		// @formatter:off
		return this.authorizationRequestResolver.resolve(exchange)
				.switchIfEmpty(chain.filter(exchange).then(Mono.empty()))
				.onErrorResume(ClientAuthorizationRequiredException.class,
						(ex) -> this.requestCache.saveRequest(exchange).then(
								this.authorizationRequestResolver.resolve(exchange, ex.getClientRegistrationId()))
				)
				.flatMap((clientRegistration) -> sendRedirectForAuthorization(exchange, clientRegistration));
		// @formatter:on
	}

	private Mono<Void> sendRedirectForAuthorization(ServerWebExchange exchange,
			OAuth2AuthorizationRequest authorizationRequest) {
		return Mono.defer(() -> {
			Mono<Void> saveAuthorizationRequest = Mono.empty();
			if (AuthorizationGrantType.AUTHORIZATION_CODE.equals(authorizationRequest.getGrantType())) {
				saveAuthorizationRequest = this.authorizationRequestRepository
						.saveAuthorizationRequest(authorizationRequest, exchange);
			}
			// @formatter:off
			URI redirectUri = UriComponentsBuilder.fromUriString(authorizationRequest.getAuthorizationRequestUri())
					.build(true)
					.toUri();
			// @formatter:on
			return saveAuthorizationRequest
					.then(this.authorizationRedirectStrategy.sendRedirect(exchange, redirectUri));
		});
	}

}
