DockerService.cs

#nullable disable
namespace Microsoft.ComponentDetection.Common;

using System;
using System.Collections.Concurrent;
using System.Collections.Generic;
using System.Linq;
using System.Text.Json;
using System.Threading;
using System.Threading.Tasks;
using Docker.DotNet;
using Docker.DotNet.Models;
using Microsoft.ComponentDetection.Common.Telemetry;
using Microsoft.ComponentDetection.Common.Telemetry.Records;
using Microsoft.ComponentDetection.Contracts;
using Microsoft.ComponentDetection.Contracts.BcdeModels;
using Microsoft.Extensions.Logging;

internal class DockerService : IDockerService
{
    // Base image annotations from ADO dockerTask
    private const string BaseImageRefAnnotation = "image.base.ref.name";
    private const string BaseImageDigestAnnotation = "image.base.digest";

    private static readonly DockerClient Client = new DockerClientConfiguration().CreateClient();

    /// <summary>
    /// Tracks in-flight image pulls so each image is pulled at most once concurrently.
    /// Concurrent callers for the same image await the same task.
    /// </summary>
    private static readonly ConcurrentDictionary<string, Task<bool>> PullCache = new();

    private static int incrementingContainerId;

    private readonly ILogger logger;

    public DockerService(ILogger<DockerService> logger) => this.logger = logger;

    public async Task<bool> CanPingDockerAsync(CancellationToken cancellationToken = default)
    {
        try
        {
            await Client.System.PingAsync(cancellationToken);
            return true;
        }
        catch (Exception e)
        {
            this.logger.LogError(e, "Failed to ping docker");
            cancellationToken.ThrowIfCancellationRequested();
            return false;
        }
    }

    public async Task<bool> CanRunLinuxContainersAsync(CancellationToken cancellationToken = default)
    {
        using var record = new DockerServiceSystemInfoTelemetryRecord();
        if (!await this.CanPingDockerAsync(cancellationToken))
        {
            return false;
        }

        try
        {
            var systemInfoResponse = await Client.System.GetSystemInfoAsync(cancellationToken);
            record.SystemInfo = JsonSerializer.Serialize(systemInfoResponse);
            return string.Equals(systemInfoResponse.OSType, "linux", StringComparison.OrdinalIgnoreCase);
        }
        catch (Exception e)
        {
            record.ExceptionMessage = e.Message;
            cancellationToken.ThrowIfCancellationRequested();
        }

        return false;
    }

    public async Task<bool> ImageExistsLocallyAsync(string image, CancellationToken cancellationToken = default)
    {
        using var record = new DockerServiceImageExistsLocallyTelemetryRecord
        {
            Image = image,
        };
        try
        {
            var imageInspectResponse = await this.InspectImageAndSanitizeVarsAsync(image, cancellationToken);
            record.ImageInspectResponse = JsonSerializer.Serialize(imageInspectResponse);
            this.logger.LogDebug("Image {Image} found locally", image);
            return true;
        }
        catch (Exception e)
        {
            record.ExceptionMessage = e.Message;
            this.logger.LogDebug("Image {Image} not found locally", image);
            cancellationToken.ThrowIfCancellationRequested();
            return false;
        }
    }

    private async Task<ImageInspectResponse> InspectImageAndSanitizeVarsAsync(string image, CancellationToken cancellationToken = default)
    {
        var imageInspectResponse = await Client.Images.InspectImageAsync(image, cancellationToken);
        this.SanitizeEnvironmentVariables(imageInspectResponse);
        return imageInspectResponse;
    }

    public async Task<bool> TryPullImageAsync(string image, CancellationToken cancellationToken = default)
    {
        // Check if already available locally before attempting a pull
        if (await this.ImageExistsLocallyAsync(image, cancellationToken))
        {
            return true;
        }

        var tcs = new TaskCompletionSource<bool>(TaskCreationOptions.RunContinuationsAsynchronously);
        var existingTask = PullCache.GetOrAdd(image, tcs.Task);

        if (existingTask != tcs.Task)
        {
            // Another caller is already pulling this image — await their result
            this.logger.LogDebug("Image {Image} is already being pulled by another caller, waiting", image);
            return await existingTask.WaitAsync(cancellationToken);
        }

        // We own this cache entry — perform the actual pull.
        try
        {
            this.logger.LogDebug("Pulling image {Image}...", image);
            var result = await this.PullImageCoreAsync(image, cancellationToken);
            this.logger.LogDebug("Pull of image {Image} completed (success={Success})", image, result);
            tcs.SetResult(result);
            return result;
        }
        catch (Exception ex)
        {
            this.logger.LogDebug(ex, "Pull of image {Image} failed", image);
            tcs.SetException(ex);
            throw;
        }
        finally
        {
            // Remove the entry once complete. The cache only deduplicates concurrent
            // in-flight pulls — subsequent callers will hit ImageExistsLocallyAsync
            // for images that were already pulled successfully.
            PullCache.TryRemove(image, out _);
        }
    }

    private async Task<bool> PullImageCoreAsync(string image, CancellationToken cancellationToken)
    {
        using var record = new DockerServiceTryPullImageTelemetryRecord
        {
            Image = image,
        };
        var parameters = new ImagesCreateParameters
        {
            FromImage = image,
        };
        try
        {
            var createImageProgress = new List<string>();
            var progress = new Progress<JSONMessage>(message =>
            {
                createImageProgress.Add(JsonSerializer.Serialize(message));
            });
            await Client.Images.CreateImageAsync(parameters, null, progress, cancellationToken);
            record.CreateImageProgress = JsonSerializer.Serialize(createImageProgress);
            return true;
        }
        catch (Exception e)
        {
            record.ExceptionMessage = e.Message;
            cancellationToken.ThrowIfCancellationRequested();
            return false;
        }
    }

    internal void SanitizeEnvironmentVariables(ImageInspectResponse inspectResponse)
    {
        var envVariables = inspectResponse?.Config?.Env;
        if (envVariables == null || !envVariables.Any())
        {
            return;
        }

        var sanitizedVarList = new List<string>();
        foreach (var variable in inspectResponse.Config.Env)
        {
            sanitizedVarList.Add(variable.RemoveSensitiveInformation());
        }

        inspectResponse.Config.Env = sanitizedVarList;
    }

    public async Task<ContainerDetails> InspectImageAsync(string image, CancellationToken cancellationToken = default)
    {
        using var record = new DockerServiceInspectImageTelemetryRecord
        {
            Image = image,
        };
        try
        {
            var imageInspectResponse = await this.InspectImageAndSanitizeVarsAsync(image, cancellationToken);
            record.ImageInspectResponse = JsonSerializer.Serialize(imageInspectResponse);

            var baseImageRef = string.Empty;
            var baseImageDigest = string.Empty;

            imageInspectResponse.Config.Labels?.TryGetValue(BaseImageRefAnnotation, out baseImageRef);
            imageInspectResponse.Config.Labels?.TryGetValue(BaseImageDigestAnnotation, out baseImageDigest);

            record.BaseImageRef = baseImageRef;
            record.BaseImageDigest = baseImageDigest;

            var layers = imageInspectResponse.RootFS?.Layers
                .Select((diffId, index) =>
                    new DockerLayer
                    {
                        DiffId = diffId,
                        LayerIndex = index,
                    });

            return new ContainerDetails
            {
                Id = GetContainerId(),
                ImageId = imageInspectResponse.ID,
                Digests = imageInspectResponse.RepoDigests,
                Tags = imageInspectResponse.RepoTags,
                CreatedAt = imageInspectResponse.Created,
                BaseImageDigest = baseImageDigest,
                BaseImageRef = baseImageRef,
                Layers = layers ?? [],
            };
        }
        catch (Exception e)
        {
            record.ExceptionMessage = e.Message;
            cancellationToken.ThrowIfCancellationRequested();
            return null;
        }
    }

    public async Task<(string Stdout, string Stderr)> CreateAndRunContainerAsync(string image, IList<string> command, CancellationToken cancellationToken = default)
    {
        return await this.CreateAndRunContainerAsync(image, command, additionalBinds: null, cancellationToken);
    }

    public async Task<(string Stdout, string Stderr)> CreateAndRunContainerAsync(string image, IList<string> command, IList<string> additionalBinds, CancellationToken cancellationToken = default)
    {
        var commandJson = JsonSerializer.Serialize(command);

        // Summary record captures overall operation including stdout/stderr
        using var record = new DockerServiceTelemetryRecord
        {
            Image = image,
            Command = commandJson,
        };

        await this.TryPullImageAsync(image, cancellationToken);
        var container = await CreateContainerAsync(image, command, additionalBinds, cancellationToken);
        record.Container = JsonSerializer.Serialize(container);

        try
        {
            var stream = await AttachContainerAsync(container.ID, cancellationToken);
            await StartContainerAsync(container.ID, cancellationToken);

            this.logger.LogInformation("Container {ContainerId} started with image {Image} to execute {Command}, reading output...", container.ID, image, commandJson);

            // Flush telemetry before the long-running ReadOutput so we get mid-scan
            // data in App Insights even if the process hangs during the read.
            TelemetryRelay.Instance.FlushCurrentTelemetry();

            var (stdout, stderr) = await ReadContainerOutputAsync(stream, container.ID, image, cancellationToken);

            record.Stdout = stdout;
            record.Stderr = stderr;

            return (stdout, stderr);
        }
        finally
        {
            // Best-effort container cleanup with a bounded timeout.
            // Use Task.WhenAny as belt-and-suspenders: even if Docker.DotNet's HTTP
            // pipeline doesn't honor the CTS (e.g. kernel-level socket blocking),
            // we abandon the removal rather than hanging indefinitely.
            this.logger.LogInformation("Removing container {ContainerId}...", container.ID);
            using var removeCts = new CancellationTokenSource(TimeSpan.FromSeconds(30));
            try
            {
                var removeTask = RemoveContainerAsync(container.ID, removeCts.Token);
                var removeTimeout = Task.Delay(TimeSpan.FromSeconds(30), CancellationToken.None);

                if (await Task.WhenAny(removeTask, removeTimeout) == removeTimeout)
                {
                    this.logger.LogWarning(
                        "RemoveContainerAsync timed out for container {ContainerId}; abandoning cleanup",
                        container.ID);

                    // Observe the abandoned task to prevent unobserved task exceptions
                    _ = removeTask.ContinueWith(
                        static _ => { },
                        CancellationToken.None,
                        TaskContinuationOptions.OnlyOnFaulted,
                        TaskScheduler.Default);
                }
                else
                {
                    await removeTask; // Observe any exception from completed task
                }
            }
            catch (Exception ex)
            {
                this.logger.LogWarning(
                    ex,
                    "Failed to remove container {ContainerId}; abandoning cleanup",
                    container.ID);
            }

            this.logger.LogInformation("Container {ContainerId} cleanup complete", container.ID);
        }
    }

    /// <summary>
    /// Reads container output with proper cancellation support.
    /// ReadOutputToEndAsync doesn't properly honor cancellation when blocked on socket read,
    /// so we race it against a cancellation-aware delay and dispose the stream if cancelled.
    /// </summary>
    private static async Task<(string Stdout, string Stderr)> ReadContainerOutputAsync(
        MultiplexedStream stream,
        string containerId,
        string image,
        CancellationToken cancellationToken)
    {
        using var record = new DockerServiceStepTelemetryRecord
        {
            Step = "ReadOutput",
            ContainerId = containerId,
            Image = image,
        };

        try
        {
            var readTask = stream.ReadOutputToEndAsync(CancellationToken.None);
            var delayTask = Task.Delay(Timeout.Infinite, cancellationToken);

            var completedTask = await Task.WhenAny(readTask, delayTask);

            if (completedTask == delayTask)
            {
                record.WasCancelled = true;

                // Dispose the stream to unblock any pending read operation.
                // Run in fire-and-forget: if the underlying socket close() blocks
                // (e.g. Docker daemon in kernel D-state), we don't want to hang here.
                _ = Task.Run(
                    () =>
                    {
                        try
                        {
                            stream.Dispose();
                        }
                        catch
                        {
                            // best effort
                        }
                    },
                    CancellationToken.None);

                // Observe the readTask to prevent unobserved task exceptions.
                // Running any continuation automatically marks the exception as observed.
                _ = readTask.ContinueWith(
                    static _ => { },
                    CancellationToken.None,
                    TaskContinuationOptions.OnlyOnFaulted,
                    TaskScheduler.Default);

                // Caller is responsible for container cleanup via finally block
                cancellationToken.ThrowIfCancellationRequested();
            }

            return await readTask;
        }
        catch (Exception ex)
        {
            record.ExceptionMessage = ex.Message;
            throw;
        }
    }

    private static async Task<CreateContainerResponse> CreateContainerAsync(
        string image,
        IList<string> command,
        IList<string> additionalBinds,
        CancellationToken cancellationToken = default)
    {
        using var record = new DockerServiceStepTelemetryRecord
        {
            Step = "CreateContainer",
            Image = image,
            Command = JsonSerializer.Serialize(command),
        };

        try
        {
            var binds = new List<string>
            {
                "/var/run/docker.sock:/var/run/docker.sock",
            };

            if (additionalBinds != null)
            {
                binds.AddRange(additionalBinds);
            }

            var parameters = new CreateContainerParameters
            {
                Image = image,
                Cmd = command,
                NetworkDisabled = true,
                HostConfig = new HostConfig
                {
                    CapDrop =
                    [
                        "all",
                    ],
                    SecurityOpt =
                    [
                        "no-new-privileges",
                    ],
                    Binds = binds,
                },
            };

            var response = await Client.Containers.CreateContainerAsync(parameters, cancellationToken);
            record.ContainerId = response.ID;
            return response;
        }
        catch (Exception ex)
        {
            record.ExceptionMessage = ex.Message;
            throw;
        }
    }

    private static async Task<MultiplexedStream> AttachContainerAsync(string containerId, CancellationToken cancellationToken = default)
    {
        using var record = new DockerServiceStepTelemetryRecord
        {
            Step = "AttachContainer",
            ContainerId = containerId,
        };

        try
        {
            var parameters = new ContainerAttachParameters
            {
                Stdout = true,
                Stderr = true,
                Stream = true,
            };
            return await Client.Containers.AttachContainerAsync(containerId, false, parameters, cancellationToken);
        }
        catch (Exception ex)
        {
            record.ExceptionMessage = ex.Message;
            throw;
        }
    }

    private static async Task StartContainerAsync(string containerId, CancellationToken cancellationToken = default)
    {
        using var record = new DockerServiceStepTelemetryRecord
        {
            Step = "StartContainer",
            ContainerId = containerId,
        };

        try
        {
            var parameters = new ContainerStartParameters();
            await Client.Containers.StartContainerAsync(containerId, parameters, cancellationToken);
        }
        catch (Exception ex)
        {
            record.ExceptionMessage = ex.Message;
            throw;
        }
    }

    private static async Task RemoveContainerAsync(string containerId, CancellationToken cancellationToken = default)
    {
        using var record = new DockerServiceStepTelemetryRecord
        {
            Step = "RemoveContainer",
            ContainerId = containerId,
        };

        try
        {
            var parameters = new ContainerRemoveParameters
            {
                Force = true,
                RemoveVolumes = true,
            };
            await Client.Containers.RemoveContainerAsync(containerId, parameters, cancellationToken);
        }
        catch (DockerContainerNotFoundException)
        {
            // Container already removed - this is expected during cleanup
        }
        catch (Exception ex)
        {
            record.ExceptionMessage = ex.Message;
            throw;
        }
    }

    private static int GetContainerId()
    {
        return Interlocked.Increment(ref incrementingContainerId);
    }

    /// <inheritdoc/>
    public ContainerDetails GetEmptyContainerDetails()
    {
        return new ContainerDetails { Id = GetContainerId() };
    }
}