Change single use installer keys to usage-based with a supplied count.

bitbound · Dec 31, 2024 · 3326817 · 3326817
1 parent 21d6a61
commit 3326817
Show file tree

Hide file tree

Showing 8 changed files with 104 additions and 49 deletions.
diff --git a/ControlR.Web.Client/Components/Pages/Deploy.razor b/ControlR.Web.Client/Components/Pages/Deploy.razor
@@ -13,32 +13,39 @@
     You must first generate an expirable installer token.
   </MudText>
 
-  <MudGrid>
+  <MudGrid Class="mt-2">
     <MudItem xs="12" sm="6">
 
-      <MudInputLabel>
+      <MudText Color="Color.Info">
         Token Type
-      </MudInputLabel>
+      </MudText>
       <MudRadioGroup T="@InstallerKeyType" @bind-Value="@_installerKeyType">
-        <MudRadio Value="@InstallerKeyType.SingleUse">
-          Single Use
+        <MudRadio Value="@InstallerKeyType.UsageBased">
+          Usage-Based
         </MudRadio>
-        <MudRadio Value="@InstallerKeyType.AbsoluteExpiration">
-          Absolute Expiration
+        <MudRadio Value="@InstallerKeyType.TimeBased">
+          Time-Based
         </MudRadio>
       </MudRadioGroup>
 
-      @if (_installerKeyType == InstallerKeyType.AbsoluteExpiration)
+      @if (_installerKeyType == InstallerKeyType.UsageBased)
+      {
+        <div class="mt-2">
+          <MudNumericField T="uint?" @bind-Value="@_totalUsesAllowed" Min="1" Label="Total uses allowed" Required />
+        </div>
+      }
+
+      @if (_installerKeyType == InstallerKeyType.TimeBased)
       {
-        <div class="mt-4">
+        <div class="mt-2">
           <MudDatePicker @bind-Date="@_inputExpirationDate" Label="Expiration Date" MinDate="@DateTime.Now" Required />
         </div>
         <div class="mt-2">
           <MudTimePicker @bind-Time="@_inputExpirationTime" Label="Expiration Time" Required />
         </div>
       }
-      <div>
-        <MudButton Variant="Variant.Filled" Color="Color.Primary" Class="mt-4" OnClick="@GenerateToken">
+      <div class="mt-4">
+        <MudButton Variant="Variant.Filled" Color="Color.Primary" OnClick="@GenerateKey">
           Generate Token
         </MudButton>
       </div>
@@ -51,14 +58,13 @@ else
     Run these scripts on a remote device to install the agent.
   </MudText>
 
-  @if (_keyExpiration is {} keyExpiration)
+  @if (_keyExpiration is { } keyExpiration)
   {
     <MudAlert Severity="Severity.Info">
       Installation scripts will expire at @(keyExpiration).
     </MudAlert>
   }
 
-
   <div class="mt-10">
     <MudText Typo="Typo.h6" Color="Color.Primary" GutterBottom>
       Windows (PowerShell)

diff --git a/ControlR.Web.Client/Components/Pages/Deploy.razor.cs b/ControlR.Web.Client/Components/Pages/Deploy.razor.cs
@@ -15,6 +15,7 @@ public partial class Deploy
   private DateTime? _inputExpirationDate;
   private TimeSpan? _inputExpirationTime;
   private string? _keyExpiration;
+  private uint? _totalUsesAllowed;
 
   [Inject]
   public required AuthenticationStateProvider AuthState { get; init; }
@@ -138,28 +139,26 @@ private async Task CopyWindowsScript()
     Snackbar.Add("Install script copied to clipboard", Severity.Success);
   }
 
-  private async Task GenerateToken()
+  private async Task GenerateKey()
   {
     switch (_installerKeyType)
     {
       case InstallerKeyType.Unknown:
         Snackbar.Add("Token type is required", Severity.Error);
         return;
-      case InstallerKeyType.SingleUse:
-        await GenerateSingleUseToken();
+      case InstallerKeyType.UsageBased:
+        await GenerateUsageBasedKey();
         break;
-      case InstallerKeyType.AbsoluteExpiration:
-        await GenerateExpiringToken();
+      case InstallerKeyType.TimeBased:
+        await GenerateTimeBasedKey();
         break;
       default:
         break;
     }
-
   }
 
-  private async Task GenerateExpiringToken()
+  private async Task GenerateTimeBasedKey()
   {
-
     if (_inputExpirationDate is null || _inputExpirationTime is null)
     {
       Snackbar.Add("Expiration date and time is required", Severity.Error);
@@ -177,7 +176,7 @@ private async Task GenerateExpiringToken()
       return;
     }
 
-    var dto = new CreateInstallerKeyRequestDto(InstallerKeyType.AbsoluteExpiration, expirationDate);
+    var dto = new CreateInstallerKeyRequestDto(InstallerKeyType.TimeBased, expirationDate);
     var createResult = await ControlrApi.CreateInstallerKey(dto);
     if (!createResult.IsSuccess)
     {
@@ -191,9 +190,15 @@ private async Task GenerateExpiringToken()
     }
   }
 
-  private async Task GenerateSingleUseToken()
+  private async Task GenerateUsageBasedKey()
   {
-    var dto = new CreateInstallerKeyRequestDto(InstallerKeyType.SingleUse);
+    if (_totalUsesAllowed is null or < 1)
+    {
+      Snackbar.Add("Total uses must be greater than 0");
+      return;
+    }
+
+    var dto = new CreateInstallerKeyRequestDto(InstallerKeyType.UsageBased, AllowedUses: _totalUsesAllowed);
     var createResult = await ControlrApi.CreateInstallerKey(dto);
     if (!createResult.IsSuccess)
     {
@@ -212,6 +217,7 @@ private async Task GenerateSingleUseToken()
   //  await Clipboard.SetText(MacArm64DeployScript);
   //  Snackbar.Add("Install script copied to clipboard", Severity.Success);
   //}
+
   private string GetCommonArgs()
   {
     var serverUri = GetServerUri();

diff --git a/ControlR.Web.Server/Api/InstallerKeysController.cs b/ControlR.Web.Server/Api/InstallerKeysController.cs
@@ -29,13 +29,18 @@ public async Task<ActionResult<CreateInstallerKeyResponseDto>> Create(
       return BadRequest("Invalid key type.");
     }
 
-    if (requestDto.KeyType == InstallerKeyType.AbsoluteExpiration &&
+    if (requestDto.KeyType == InstallerKeyType.TimeBased &&
        (!requestDto.Expiration.HasValue || requestDto.Expiration.Value < timeProvider.GetLocalNow()))
     {
       return BadRequest("Expiration date must be in the future.");
     }
 
-    var key = await keyManager.CreateKey(tenantId, userId, requestDto.KeyType, requestDto.Expiration);
-    return new CreateInstallerKeyResponseDto(requestDto.KeyType, key.AccessToken, requestDto.Expiration);
+    if (requestDto.KeyType == InstallerKeyType.UsageBased && requestDto.AllowedUses < 1)
+    {
+      return BadRequest("Allowed uses must be more than 0.");
+    }
+
+    var key = await keyManager.CreateKey(tenantId, userId, requestDto.KeyType, requestDto.AllowedUses, requestDto.Expiration);
+    return new CreateInstallerKeyResponseDto(requestDto.KeyType, key.AccessToken, requestDto.AllowedUses, requestDto.Expiration);
   }
 }
diff --git a/ControlR.Web.Server/Models/AgentInstallerKey.cs b/ControlR.Web.Server/Models/AgentInstallerKey.cs
@@ -5,4 +5,6 @@ public record AgentInstallerKey(
   Guid CreatorId,
   string AccessToken,
   InstallerKeyType KeyType,
-  DateTimeOffset? Expiration);
+  uint? AllowedUses,
+  DateTimeOffset? Expiration,
+  uint CurrentUses = 0);
diff --git a/ControlR.Web.Server/Services/AgentInstallerKeyManager.cs b/ControlR.Web.Server/Services/AgentInstallerKeyManager.cs
@@ -5,8 +5,8 @@ namespace ControlR.Web.Server.Services;
 
 public interface IAgentInstallerKeyManager
 {
-  Task<AgentInstallerKey> CreateKey(Guid tenantId, Guid creatorId, InstallerKeyType keyType, DateTimeOffset? expiration);
-  Task<bool> ValidateKey(string token);
+  Task<AgentInstallerKey> CreateKey(Guid tenantId, Guid creatorId, InstallerKeyType keyType, uint? allowedUses, DateTimeOffset? expiration);
+  Task<bool> ValidateKey(string key);
 }
 
 public class AgentInstallerKeyManager(
@@ -27,24 +27,39 @@ public Task<AgentInstallerKey> CreateKey(
     Guid tenantId,
     Guid creatorId,
     InstallerKeyType keyType,
+    uint? allowedUses,
     DateTimeOffset? expiration)
   {
-    var token = RandomGenerator.CreateAccessToken();
-    var installerKey = new AgentInstallerKey(tenantId, creatorId, token, keyType, expiration);
-    if (expiration.HasValue)
-    {
-      _keyCache.Set(token, installerKey, expiration.Value);
-    }
-    else
+    var key = RandomGenerator.CreateAccessToken();
+    var installerKey = new AgentInstallerKey(tenantId, creatorId, key, keyType, allowedUses, expiration);
+
+    switch (keyType)
     {
-      _keyCache.Set(token, installerKey);
+      case InstallerKeyType.UsageBased:
+        {
+          _keyCache.Set(key, installerKey);
+          break;
+        }
+      case InstallerKeyType.TimeBased:
+        {
+          if (!expiration.HasValue)
+          {
+            throw new ArgumentNullException(nameof(expiration));
+          }
+          _keyCache.Set(key, installerKey, expiration.Value);
+          break;
+        }
+      case InstallerKeyType.Unknown:
+      default:
+        throw new ArgumentOutOfRangeException(nameof(keyType), "Unknown installer key type.");
     }
+
     return installerKey.AsTaskResult();
   }
 
-  public Task<bool> ValidateKey(string token)
+  public Task<bool> ValidateKey(string key)
   {
-    if (!_keyCache.TryGetValue(token, out var cachedObject))
+    if (!_keyCache.TryGetValue(key, out var cachedObject))
     {
       return false.AsTaskResult();
     }
@@ -58,12 +73,31 @@ public Task<bool> ValidateKey(string token)
     {
       case InstallerKeyType.Unknown:
         break;
-      case InstallerKeyType.SingleUse:
-        _keyCache.Remove(installerKey);
-        return true.AsTaskResult();
-      case InstallerKeyType.AbsoluteExpiration:
-        var isValid = installerKey.Expiration.HasValue && installerKey.Expiration.Value >= _timeProvider.GetUtcNow();
-        return isValid.AsTaskResult();
+      case InstallerKeyType.UsageBased:
+        {
+          // This operation has a race condition if multiple validations are being
+          // performed on the same key concurrently.  But the risk/impact is so
+          // small that it's not worth locking the resource.
+
+          var isValid = installerKey.CurrentUses < installerKey.AllowedUses;
+          installerKey = installerKey with { CurrentUses = installerKey.CurrentUses + 1 };
+
+          if (installerKey.CurrentUses >= installerKey.AllowedUses)
+          {
+            _keyCache.Remove(key);
+          }
+
+          return isValid.AsTaskResult();
+        }
+      case InstallerKeyType.TimeBased:
+        {
+          var isValid = installerKey.Expiration.HasValue && installerKey.Expiration.Value >= _timeProvider.GetUtcNow();
+          if (!isValid)
+          {
+            _keyCache.Remove(key);
+          }
+          return isValid.AsTaskResult();
+        }
       default:
         _logger.LogError("Unknown installer key type: {KeyType}", installerKey.KeyType);
         break;

diff --git a/Libraries/ControlR.Libraries.Shared/Dtos/ServerApi/CreateInstallerKeyRequestDto.cs b/Libraries/ControlR.Libraries.Shared/Dtos/ServerApi/CreateInstallerKeyRequestDto.cs
@@ -5,4 +5,5 @@ namespace ControlR.Libraries.Shared.Dtos.ServerApi;
 [DataContract]
 public record CreateInstallerKeyRequestDto(
   [property: DataMember] InstallerKeyType KeyType,
-  [property: DataMember] DateTimeOffset? Expiration = null);
+  [property: DataMember] DateTimeOffset? Expiration = null,
+  [property: DataMember] uint? AllowedUses = null);
diff --git a/Libraries/ControlR.Libraries.Shared/Dtos/ServerApi/CreateInstallerKeyResponseDto.cs b/Libraries/ControlR.Libraries.Shared/Dtos/ServerApi/CreateInstallerKeyResponseDto.cs
@@ -6,4 +6,5 @@ namespace ControlR.Libraries.Shared.Dtos.ServerApi;
 public record CreateInstallerKeyResponseDto(
   [property: DataMember] InstallerKeyType TokenType,
   [property: DataMember] string AccessKey,
+  [property: DataMember] uint? AllowedUses = null,
   [property: DataMember] DateTimeOffset? Expiration = null);
diff --git a/Libraries/ControlR.Libraries.Shared/Dtos/ServerApi/InstallerKeyType.cs b/Libraries/ControlR.Libraries.Shared/Dtos/ServerApi/InstallerKeyType.cs
@@ -8,7 +8,7 @@ public enum InstallerKeyType
   [EnumMember]
   Unknown = 0,
   [EnumMember]
-  SingleUse = 1,
+  UsageBased = 1,
   [EnumMember]
-  AbsoluteExpiration = 2
+  TimeBased = 2
 }