Skip to content

Commit

Permalink
Implemented option to store token in named container (#4607)
Browse files Browse the repository at this point in the history 
* Implemented option to store token in named container

* cleanup

* test fix

* test fix

* test fix
  • Loading branch information
@merlynomsft
merlynomsft authored Jan 25, 2024
1 parent 864c7bb commit 2e829bc
Show file tree
Hide file tree
Showing 10 changed files with 216 additions and 27 deletions.
127 changes: 117 additions & 10 deletions src/Agent.Listener/Configuration.Windows/RSAEncryptedFileKeyManager.cs
View file
Original file line number Diff line number Diff line change
@@ -1,10 +1,11 @@
// Copyright (c) Microsoft Corporation.
// Licensed under the MIT License.

using Microsoft.VisualStudio.Services.Agent.Util;
using System;
using System.IO;
using System.Security.Cryptography;
using System.Text;
using Microsoft.VisualStudio.Services.Agent.Util;

namespace Microsoft.VisualStudio.Services.Agent.Listener.Configuration
{
Expand All @@ -13,7 +14,61 @@ public class RSAEncryptedFileKeyManager : AgentService, IRSAKeyManager
private string _keyFile;
private IHostContext _context;

public RSACryptoServiceProvider CreateKey()
public RSACryptoServiceProvider CreateKey(bool enableAgentKeyStoreInNamedContainer)
{
if (enableAgentKeyStoreInNamedContainer)
{
return CreateKeyStoreKeyInNamedContainer();
}
else
{
return CreateKeyStoreKeyInFile();
}
}

private RSACryptoServiceProvider CreateKeyStoreKeyInNamedContainer()
{
RSACryptoServiceProvider rsa;
if (!File.Exists(_keyFile))
{
Trace.Info("Creating new RSA key using 2048-bit key length");

CspParameters Params = new CspParameters();
Params.KeyContainerName = "AgentKeyContainer" + Guid.NewGuid().ToString();
Params.Flags |= CspProviderFlags.UseNonExportableKey | CspProviderFlags.UseMachineKeyStore;
rsa = new RSACryptoServiceProvider(2048, Params);

// Now write the parameters to disk
SaveParameters(default(RSAParameters), Params.KeyContainerName);
Trace.Info("Successfully saved containerName to file {0} in container {1}", _keyFile, Params.KeyContainerName);
}
else
{
Trace.Info("Found existing RSA key parameters file {0}", _keyFile);

var result = LoadParameters();

if(string.IsNullOrEmpty(result.containerName))
{
Trace.Info("Container name not present; reading RSA key from file");
return CreateKeyStoreKeyInFile();
}

CspParameters Params = new CspParameters();
Params.KeyContainerName = result.containerName;
Params.Flags |= CspProviderFlags.UseNonExportableKey | CspProviderFlags.UseMachineKeyStore;
rsa = new RSACryptoServiceProvider(Params);
}

return rsa;

// References:
// https://stackoverflow.com/questions/2274596/how-to-store-a-public-key-in-a-machine-level-rsa-key-container
// https://social.msdn.microsoft.com/Forums/en-US/e3902420-3a82-42cf-a4a3-de230ebcea56/how-to-store-a-public-key-in-a-machinelevel-rsa-key-container?forum=netfxbcl
// https://security.stackexchange.com/questions/234477/windows-certificates-where-is-private-key-located
}

private RSACryptoServiceProvider CreateKeyStoreKeyInFile()
{
RSACryptoServiceProvider rsa = null;
if (!File.Exists(_keyFile))
Expand All @@ -23,15 +78,23 @@ public RSACryptoServiceProvider CreateKey()
rsa = new RSACryptoServiceProvider(2048);

// Now write the parameters to disk
SaveParameters(rsa.ExportParameters(true));
SaveParameters(rsa.ExportParameters(true), string.Empty);
Trace.Info("Successfully saved RSA key parameters to file {0}", _keyFile);
}
else
{
Trace.Info("Found existing RSA key parameters file {0}", _keyFile);

var result = LoadParameters();

if(!string.IsNullOrEmpty(result.containerName))
{
Trace.Info("Keyfile has ContainerName, so we must read from named container");
return CreateKeyStoreKeyInNamedContainer();
}

rsa = new RSACryptoServiceProvider();
rsa.ImportParameters(LoadParameters());
rsa.ImportParameters(result.rsaParameters);
}

return rsa;
Expand All @@ -46,7 +109,19 @@ public void DeleteKey()
}
}

public RSACryptoServiceProvider GetKey()
public RSACryptoServiceProvider GetKey(bool enableAgentKeyStoreInNamedContainer)
{
if (enableAgentKeyStoreInNamedContainer)
{
return GetKeyFromNamedContainer();
}
else
{
return GetKeyFromFile();
}
}

private RSACryptoServiceProvider GetKeyFromNamedContainer()
{
if (!File.Exists(_keyFile))
{
Expand All @@ -55,21 +130,53 @@ public RSACryptoServiceProvider GetKey()

Trace.Info("Loading RSA key parameters from file {0}", _keyFile);

var result = LoadParameters();

if (string.IsNullOrEmpty(result.containerName))
{
return GetKeyFromFile();
}

CspParameters Params = new CspParameters();
Params.KeyContainerName = result.containerName;
Params.Flags |= CspProviderFlags.UseNonExportableKey | CspProviderFlags.UseMachineKeyStore;
var rsa = new RSACryptoServiceProvider(Params);
return rsa;
}

private RSACryptoServiceProvider GetKeyFromFile()
{
if (!File.Exists(_keyFile))
{
throw new CryptographicException(StringUtil.Loc("RSAKeyFileNotFound", _keyFile));
}

Trace.Info("Loading RSA key parameters from file {0}", _keyFile);

var result = LoadParameters();

if(!string.IsNullOrEmpty(result.containerName))
{
Trace.Info("Keyfile has ContainerName, reading from NamedContainer");
return GetKeyFromNamedContainer();
}

var rsa = new RSACryptoServiceProvider();
rsa.ImportParameters(LoadParameters());
rsa.ImportParameters(result.rsaParameters);
return rsa;
}

private RSAParameters LoadParameters()
private (string containerName, RSAParameters rsaParameters) LoadParameters()
{
var encryptedBytes = File.ReadAllBytes(_keyFile);
var parametersString = Encoding.UTF8.GetString(ProtectedData.Unprotect(encryptedBytes, null, DataProtectionScope.LocalMachine));
return StringUtil.ConvertFromJson<RSAParametersSerializable>(parametersString).RSAParameters;
var deserialized = StringUtil.ConvertFromJson<RSAParametersSerializable>(parametersString);
return (deserialized.ContainerName, deserialized.RSAParameters);
}

private void SaveParameters(RSAParameters parameters)
private void SaveParameters(RSAParameters parameters, string containerName)
{
var parametersString = StringUtil.ConvertToJson(new RSAParametersSerializable(parameters));
var parametersString = StringUtil.ConvertToJson(new RSAParametersSerializable(containerName, parameters));
var encryptedBytes = ProtectedData.Protect(Encoding.UTF8.GetBytes(parametersString), null, DataProtectionScope.LocalMachine);
File.WriteAllBytes(_keyFile, encryptedBytes);
File.SetAttributes(_keyFile, File.GetAttributes(_keyFile) | FileAttributes.Hidden);
Expand Down
5 changes: 3 additions & 2 deletions src/Agent.Listener/Configuration/ConfigurationManager.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -178,11 +178,12 @@ public async Task ConfigureAsync(CommandSettings command)
_term.WriteError(StringUtil.Loc("FailedToConnect"));
}
}

// We want to use the native CSP of the platform for storage, so we use the RSACSP directly
RSAParameters publicKey;
var keyManager = HostContext.GetService<IRSAKeyManager>();
using (var rsa = keyManager.CreateKey())
var enableAgentKeyStoreInNamedContainer = await keyManager.GetStoreAgentTokenInNamedContainerFF(HostContext, Trace, agentSettings, creds);
using (var rsa = keyManager.CreateKey(enableAgentKeyStoreInNamedContainer))
{
publicKey = rsa.ExportParameters(false);
}
Expand Down
22 changes: 17 additions & 5 deletions src/Agent.Listener/Configuration/FeatureFlagProvider.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,7 @@ public interface IFeatureFlagProvider : IAgentService
/// <exception cref="InvalidOperationException">Thrown if agent is not configured</exception>
public Task<FeatureFlag> GetFeatureFlagAsync(IHostContext context, string featureFlagName, ITraceWriter traceWriter, CancellationToken ctk = default);

public Task<FeatureFlag> GetFeatureFlagWithCred(IHostContext context, string featureFlagName, ITraceWriter traceWriter, AgentSettings settings, VssCredentials creds, CancellationToken ctk = default);
}

public class FeatureFlagProvider : AgentService, IFeatureFlagProvider
Expand All @@ -40,22 +41,33 @@ public async Task<FeatureFlag> GetFeatureFlagAsync(IHostContext context, string
ArgUtil.NotNull(featureFlagName, nameof(featureFlagName));

var credMgr = context.GetService<ICredentialManager>();
VssCredentials creds = credMgr.LoadCredentials();
var configManager = context.GetService<IConfigurationManager>();
AgentSettings settings = configManager.LoadSettings();

return await GetFeatureFlagWithCred(context, featureFlagName, traceWriter, settings, creds, ctk);
}

public async Task<FeatureFlag> GetFeatureFlagWithCred(IHostContext context, string featureFlagName,
ITraceWriter traceWriter, AgentSettings settings, VssCredentials creds, CancellationToken ctk)
{
var agentCertManager = context.GetService<IAgentCertificateManager>();

VssCredentials creds = credMgr.LoadCredentials();
ArgUtil.NotNull(creds, nameof(creds));

AgentSettings settings = configManager.LoadSettings();
using var vssConnection = VssUtil.CreateConnection(new Uri(settings.ServerUrl), creds, traceWriter, agentCertManager.SkipServerCertificateValidation);
var client = vssConnection.GetClient<FeatureAvailabilityHttpClient>();
try
{
return await client.GetFeatureFlagByNameAsync(featureFlagName, checkFeatureExists: false);
} catch (VssServiceException e) {
return await client.GetFeatureFlagByNameAsync(featureFlagName, checkFeatureExists: false, ctk);
}
catch (VssServiceException e)
{
Trace.Warning("Unable to retrieve feature flag status: " + e.ToString());
return new FeatureFlag(featureFlagName, "", "", "Off", "Off");
} catch (VssUnauthorizedException e) {
}
catch (VssUnauthorizedException e)
{
Trace.Warning("Unable to retrieve feature flag with following exception: " + e.ToString());
return new FeatureFlag(featureFlagName, "", "", "Off", "Off");
}
Expand Down
35 changes: 32 additions & 3 deletions src/Agent.Listener/Configuration/IRSAKeyManager.cs
View file
Original file line number Diff line number Diff line change
@@ -1,9 +1,15 @@
// Copyright (c) Microsoft Corporation.
// Licensed under the MIT License.

using Agent.Listener.Configuration;
using Agent.Sdk.Knob;
using Microsoft.VisualStudio.Services.Agent.Util;
using Microsoft.VisualStudio.Services.Common;
using System;
using System.Runtime.Serialization;
using System.Security.Cryptography;
using System.Threading;
using System.Threading.Tasks;

namespace Microsoft.VisualStudio.Services.Agent.Listener.Configuration
{
Expand All @@ -21,7 +27,7 @@ public interface IRSAKeyManager : IAgentService
/// key is returned to the caller.
/// </summary>
/// <returns>An <c>RSACryptoServiceProvider</c> instance representing the key for the agent</returns>
RSACryptoServiceProvider CreateKey();
RSACryptoServiceProvider CreateKey(bool enableAgentKeyStoreInNamedContainer);

/// <summary>
/// Deletes the RSA key managed by the key manager.
Expand All @@ -33,7 +39,23 @@ public interface IRSAKeyManager : IAgentService
/// </summary>
/// <returns>An <c>RSACryptoServiceProvider</c> instance representing the key for the agent</returns>
/// <exception cref="CryptographicException">No key exists in the store</exception>
RSACryptoServiceProvider GetKey();
RSACryptoServiceProvider GetKey(bool enableAgentKeyStoreInNamedContainer);
}

public static class IRSAKeyManagerExtensions
{
public static async Task<bool> GetStoreAgentTokenInNamedContainerFF(this IRSAKeyManager rsaKeyManager, IHostContext hostContext, global::Agent.Sdk.ITraceWriter trace, AgentSettings agentSettings, VssCredentials creds, CancellationToken cancellationToken = default)
{
if(AgentKnobs.StoreAgentKeyInCSPContainer.GetValue(UtilKnobValueContext.Instance()).AsBoolean())
{
return true;
}

var featureFlagProvider = hostContext.GetService<IFeatureFlagProvider>();
var enableAgentKeyStoreInNamedContainer = (await featureFlagProvider.GetFeatureFlagWithCred(hostContext, "DistributedTask.Agent.StoreAgentTokenInNamedContainer", trace, agentSettings, creds, cancellationToken)).EffectiveState == "On";

return enableAgentKeyStoreInNamedContainer;
}
}

// Newtonsoft 10 is not working properly with dotnet RSAParameters class
Expand All @@ -44,6 +66,7 @@ public interface IRSAKeyManager : IAgentService
[Serializable]
internal class RSAParametersSerializable : ISerializable
{
private string _containerName;
private RSAParameters _rsaParameters;

public RSAParameters RSAParameters
Expand All @@ -54,15 +77,18 @@ public RSAParameters RSAParameters
}
}

public RSAParametersSerializable(RSAParameters rsaParameters)
public RSAParametersSerializable(string containerName, RSAParameters rsaParameters)
{
_containerName = containerName;
_rsaParameters = rsaParameters;
}

private RSAParametersSerializable()
{
}

public string ContainerName { get { return _containerName; } set { _containerName = value; } }

public byte[] D { get { return _rsaParameters.D; } set { _rsaParameters.D = value; } }

public byte[] DP { get { return _rsaParameters.DP; } set { _rsaParameters.DP = value; } }
Expand All @@ -81,6 +107,8 @@ private RSAParametersSerializable()

public RSAParametersSerializable(SerializationInfo information, StreamingContext context)
{
_containerName = (string)information.GetValue("ContainerName", typeof(string));

_rsaParameters = new RSAParameters()
{
D = (byte[])information.GetValue("d", typeof(byte[])),
Expand All @@ -96,6 +124,7 @@ public RSAParametersSerializable(SerializationInfo information, StreamingContext

public void GetObjectData(SerializationInfo info, StreamingContext context)
{
info.AddValue("ContainerName", _containerName);
info.AddValue("d", _rsaParameters.D);
info.AddValue("dp", _rsaParameters.DP);
info.AddValue("dq", _rsaParameters.DQ);
Expand Down
2 changes: 1 addition & 1 deletion src/Agent.Listener/Configuration/OAuthCredential.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -42,7 +42,7 @@ public override VssCredentials GetVssCredentials(IHostContext context)
// We expect the key to be in the machine store at this point. Configuration should have set all of
// this up correctly so we can use the key to generate access tokens.
var keyManager = context.GetService<IRSAKeyManager>();
var signingCredentials = VssSigningCredentials.Create(() => keyManager.GetKey());
var signingCredentials = VssSigningCredentials.Create(() => keyManager.GetKey(enableAgentKeyStoreInNamedContainer: true));
var clientCredential = new VssOAuthJwtBearerClientCredential(clientId, authorizationUrl, signingCredentials);
var agentCredential = new VssOAuthCredential(new Uri(oathEndpointUrl, UriKind.Absolute), VssOAuthGrant.ClientCredentials, clientCredential);

Expand Down
6 changes: 3 additions & 3 deletions src/Agent.Listener/Configuration/RSAFileKeyManager.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,7 @@ public class RSAFileKeyManager : AgentService, IRSAKeyManager
private string _keyFile;
private IHostContext _context;

public RSACryptoServiceProvider CreateKey()
public RSACryptoServiceProvider CreateKey(bool enableAgentKeyStoreInNamedContainer)
{
RSACryptoServiceProvider rsa = null;
if (!File.Exists(_keyFile))
Expand All @@ -24,7 +24,7 @@ public RSACryptoServiceProvider CreateKey()
rsa = new RSACryptoServiceProvider(2048);

// Now write the parameters to disk
IOUtil.SaveObject(new RSAParametersSerializable(rsa.ExportParameters(true)), _keyFile);
IOUtil.SaveObject(new RSAParametersSerializable("", rsa.ExportParameters(true)), _keyFile);
Trace.Info("Successfully saved RSA key parameters to file {0}", _keyFile);

// Try to lock down the credentials_key file to the owner/group
Expand Down Expand Up @@ -70,7 +70,7 @@ public void DeleteKey()
}
}

public RSACryptoServiceProvider GetKey()
public RSACryptoServiceProvider GetKey(bool enableAgentKeyStoreInNamedContainer)
{
if (!File.Exists(_keyFile))
{
Expand Down
Loading

0 comments on commit 2e829bc

Please sign in to comment.