HttpPipelineExtensions.cs

// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.

#nullable enable

using System.Threading;
using System.Threading.Tasks;
using Azure.Core.Pipeline;

namespace Azure.Core
{
    internal static class HttpPipelineExtensions
    {
        public static async ValueTask<Response> ProcessMessageAsync(this HttpPipeline pipeline, HttpMessage message, RequestContext? requestContext, CancellationToken cancellationToken = default)
        {
            var (userCt, statusOption) = ApplyRequestContext(requestContext);
            if (!userCt.CanBeCanceled || !cancellationToken.CanBeCanceled)
            {
                await pipeline.SendAsync(message, cancellationToken.CanBeCanceled ? cancellationToken : userCt).ConfigureAwait(false);
            }
            else
            {
                using var cts = CancellationTokenSource.CreateLinkedTokenSource(userCt, cancellationToken);
                await pipeline.SendAsync(message, cts.Token).ConfigureAwait(false);
            }

            if (!message.Response.IsError || statusOption == ErrorOptions.NoThrow)
            {
                return message.Response;
            }

            throw new RequestFailedException(message.Response);
        }

        public static Response ProcessMessage(this HttpPipeline pipeline, HttpMessage message, RequestContext? requestContext, CancellationToken cancellationToken = default)
        {
            var (userCt, statusOption) = ApplyRequestContext(requestContext);
            if (!userCt.CanBeCanceled || !cancellationToken.CanBeCanceled)
            {
                pipeline.Send(message, cancellationToken.CanBeCanceled ? cancellationToken : userCt);
            }
            else
            {
                using var cts = CancellationTokenSource.CreateLinkedTokenSource(userCt, cancellationToken);
                pipeline.Send(message, cts.Token);
            }

            if (!message.Response.IsError || statusOption == ErrorOptions.NoThrow)
            {
                return message.Response;
            }

            throw new RequestFailedException(message.Response);
        }

        public static async ValueTask<Response<bool>> ProcessHeadAsBoolMessageAsync(this HttpPipeline pipeline, HttpMessage message, ClientDiagnostics clientDiagnostics, RequestContext? requestContext)
        {
            var response = await pipeline.ProcessMessageAsync(message, requestContext).ConfigureAwait(false);
            switch (response.Status)
            {
                case >= 200 and < 300:
                    return Response.FromValue(true, response);
                case >= 400 and < 500:
                    return Response.FromValue(false, response);
                default:
                    var exception = await clientDiagnostics.CreateRequestFailedExceptionAsync(response).ConfigureAwait(false);
                    return new ErrorResponse<bool>(response, exception);
            }
        }

        public static Response<bool> ProcessHeadAsBoolMessage(this HttpPipeline pipeline, HttpMessage message, ClientDiagnostics clientDiagnostics, RequestContext? requestContext)
        {
            var response = pipeline.ProcessMessage(message, requestContext);
            switch (response.Status)
            {
                case >= 200 and < 300:
                    return Response.FromValue(true, response);
                case >= 400 and < 500:
                    return Response.FromValue(false, response);
                default:
                    var exception = clientDiagnostics.CreateRequestFailedException(response);
                    return new ErrorResponse<bool>(response, exception);
            }
        }

        private static (CancellationToken CancellationToken, ErrorOptions ErrorOptions) ApplyRequestContext(RequestContext? requestContext)
        {
            if (requestContext == null)
            {
                return (CancellationToken.None, ErrorOptions.Default);
            }

            return (requestContext.CancellationToken, requestContext.ErrorOptions);
        }
    }
}