custom images ms3

lib/travis/api/enqueue/services/restart_model.rb

@@ -27,7 +27,7 @@ def push(event, payload)
         end
 
         def accept?
-          current_user && permission? && resetable? && billing?
+          current_user && permission? && resetable? && billing? && custom_image_create_allowed? && custom_image_use_allowed?
         end
 
         def billing?
@@ -61,12 +61,64 @@ def billing?
           end
         end
 
+        def custom_image_create_allowed?
+          return true if !!Travis.config.enterprise
+
+          @_custom_image_create_allowed ||= begin
+            jobs = target.is_a?(Job) ? [target] : target.matrix
+            jobs.map do |job|
+              next unless job.config
+
+              create_name = job.config.dig(:vm, :create, :name)
+              if create_name
+                return false unless !!artifact_manager.create(owner: repository.owner, image_name: create_name, job_restart: true)
+              end
+            end
+            true
+          rescue Travis::API::V3::Error
+            false
+          end
+        end
+
+
+        def custom_image_use_allowed?
+          return true if !!Travis.config.enterprise
+
+          @_custom_image_use_allowed ||= begin
+            jobs = target.is_a?(Job) ? [target] : target.matrix
+
+            jobs.map do |job|
+              next unless job.config
+
+              use_name = job.config.dig(:vm, :use)
+              use_name = use_name[:name] if use_name.is_a?(Hash)
+
+              if use_name
+                return false unless can_use_custom_image?(owner: repository.owner, image_name: use_name)
+              end
+            end
+            true
+          rescue Travis::API::V3::Error
+            false
+          end
+        end
+
+        def can_use_custom_image?(owner:, image_name:)
+          !!artifact_manager.use(owner: , image_name:)
+        rescue Travis::API::V3::NotFound
+          true
+        rescue Travis::API::V3::Error
+          false
+        end
+
         def messages
           messages = []
           messages << { notice: "The #{type} was successfully restarted." } if accept?
           messages << { error:  'You do not seem to have sufficient permissions.' } unless permission?
           messages << { error:  'You do not have enough credits.' } unless billing?
           messages << { error:  "This #{type} currently can not be restarted." } unless resetable?
+          messages << { error:  "Image creation build restart not allowed." } unless custom_image_create_allowed?
+          messages << { error:  "Can't use the custom image. Make sure it's already created" } unless custom_image_use_allowed?
           messages
         end
 
@@ -88,6 +140,10 @@ def target
 
         private
 
+        def artifact_manager
+          @_artifact_manger = Travis::API::V3::ArtifactManagerClient.new(current_user&.id)
+        end
+
         def subscription
           Subscription.where(owner: repository.owner)&.first
         end

lib/travis/api/v3/artifact_manager_client.rb

@@ -0,0 +1,99 @@
+module Travis::API::V3
+  class ArtifactManagerClient
+    class ConfigurationError < StandardError; end
+
+    def initialize(user_id)
+      @user_id = user_id
+    end
+
+    def create(owner:, image_name:, job_restart: false)
+      params = {
+        owner_type: owner.class.name.downcase,
+        id: owner.id,
+        name: image_name,
+        job_restart:
+      }
+      handle_errors_and_respond(connection.post("/create", params)) do |body|
+        body.include?('image_id') ? body['image_id'] : false
+      end
+    rescue Faraday::Error
+      raise ArtifactManagerConnectionError
+    end
+
+    def use(owner:, image_name:)
+      handle_errors_and_respond(connection.get("/image/#{owner.class.name.downcase}/#{owner.id}/#{image_name}")) do |body|
+        return body['image_id'] if body.include?('image_id')
+
+        return body['image']['id'] if body.include?('image')
+
+        false
+      end
+    rescue Faraday::Error
+      raise ArtifactManagerConnectionError
+    end
+
+    def images(owner_type, owner_id)
+      response = connection.get("/images?owner_type=#{owner_type}&id=#{owner_id}")
+      handle_images_response(response)
+    end
+
+    def delete_images(owner_type, owner_id, image_ids)
+      image_ids.each do |image_id|
+        response = connection.delete("/image/#{owner_type.downcase}/#{owner_id}/#{image_id}")
+        handle_errors_and_respond(response)
+      end
+    end
+
+    private
+
+    def handle_images_response(response)
+      handle_errors_and_respond(response) do |r|
+        r['images'].map { |image_data| Travis::API::V3::Models::CustomImage.new(image_data) }
+      end
+    end
+
+    def handle_errors_and_respond(response)
+      body = response.body.is_a?(String) && response.body.length.positive? ? JSON.parse(response.body) : response.body
+
+      case response.status
+      when 200, 201
+        yield(body) if block_given?
+      when 202
+        true
+      when 204
+        true
+      when 400
+        raise Travis::API::V3::ClientError, body['error']
+      when 403
+        raise Travis::API::V3::InsufficientAccess, body['rejection_code']
+      when 404
+        raise Travis::API::V3::NotFound, body['error']
+      when 422
+        raise Travis::API::V3::UnprocessableEntity, body['error']
+      else
+        raise Travis::API::V3::ServerError, 'Artifact manager failed'
+      end
+    end
+
+    def connection(timeout: 10)
+      @connection ||= Faraday.new(url: artifact_manager_url, ssl: { ca_path: '/usr/lib/ssl/certs' }) do |conn|
+        conn.request(:authorization, :basic, '_', artifact_manager_auth_key)
+        conn.headers['X-Travis-User-Id'] = @user_id.to_s
+        conn.headers['Content-Type'] = 'application/json'
+        conn.request :json
+        conn.response :json
+        conn.options[:open_timeout] = timeout
+        conn.options[:timeout] = timeout
+        conn.adapter :net_http
+      end
+    end
+
+    def artifact_manager_url
+      Travis.config.artifact_manager&.url || raise(ConfigurationError, 'No artifact manager url configured')
+    end
+
+    def artifact_manager_auth_key
+      Travis.config.artifact_manager&.auth_key || raise(ConfigurationError, 'No artifact manager auth key configured')
+    end
+  end
+end

lib/travis/api/v3/billing_client.rb

@@ -47,6 +47,20 @@ def executions(owner_type, owner_id, page, per_page, from, to)
       executions
     end
 
+    def storage_usage(owner_type, owner_id, from, to)
+      response = connection.get("/usage/#{owner_type.downcase}s/#{owner_id}/storage?from=#{from}&to=#{to}")
+      body(response).map do |usage_data|
+        usage_data
+      end
+    end
+
+    def storage_executions_usage(owner_type, owner_id)
+      response = connection.get("/usage/#{owner_type.downcase}s/#{owner_id}/storage_executions")
+      body(response).map do |usage_data|
+        usage_data
+      end
+    end
+
     def calculate_credits(users, executions)
       response = connection.post("/usage/credits_calculator", users: users, executions: executions)
       body(response).map do |calculator_data|

lib/travis/api/v3/models/custom_image.rb

@@ -0,0 +1,20 @@
+module Travis::API::V3
+  class Models::CustomImage < Model
+    after_initialize :readonly!
+    belongs_to :owner, polymorphic: true
+    has_many :custom_image_logs
+
+    scope :available, -> { where(state: 'available') }
+
+    def created_by
+      user_id = custom_image_logs.created.first&.sender_id
+      return unless user_id
+
+      User.find(user_id)
+    end
+
+    def private
+      true
+    end
+  end
+end

lib/travis/api/v3/models/custom_image_log.rb

@@ -0,0 +1,7 @@
+module Travis::API::V3
+  class Models::CustomImageLog < Model
+    after_initialize :readonly!
+    belongs_to :custom_image
+    enum action: { created: 'created', used: 'used', deleted: 'deleted', other: 'other' }
+  end
+end

lib/travis/api/v3/models/custom_image_storage.rb

@@ -0,0 +1,4 @@
+module Travis::API::V3
+  class Models::CustomImageStorage < Model
+  end
+end

lib/travis/api/v3/models/v2_addon_usage.rb

@@ -1,6 +1,7 @@
 module Travis::API::V3
   class Models::V2AddonUsage
-    attr_reader :id, :addon_id, :addon_quantity, :addon_usage, :remaining, :purchase_date, :valid_to, :active, :status
+    attr_reader :id, :addon_id, :addon_quantity, :addon_usage, :remaining, :purchase_date, :valid_to, :active, :status,
+                :quantity_limit_free, :quantity_limit_type, :quantity_limit_charge, :total_usage
 
     def initialize(attrs)
       @id = attrs.fetch('id')
@@ -12,6 +13,10 @@ def initialize(attrs)
       @valid_to = attrs.fetch('valid_to')
       @active = attrs.fetch('active')
       @status = attrs.fetch('status')
+      @quantity_limit_free = attrs.fetch('quantity_limit_free', 0)
+      @quantity_limit_type = attrs.fetch('quantity_limit_type', nil)
+      @quantity_limit_charge = attrs.fetch('quantity_limit_charge', nil)
+      @total_usage = attrs.fetch('total_usage', nil)
     end
   end
 end

lib/travis/api/v3/queries/custom_images.rb

@@ -0,0 +1,31 @@
+module Travis::API::V3
+  class Queries::CustomImages < Query
+    def for_owner(owner)
+      Models::CustomImage.available.where(owner_id: owner.id, owner_type: owner_type(owner)).order('created_at DESC')
+    end
+
+    def delete(image_ids, owner, sender)
+      client = ArtifactManagerClient.new(sender.id)
+      client.delete_images(owner_type(owner), owner.id, image_ids)
+    end
+
+    def usage(owner, user_id, from, to)
+      client = BillingClient.new(user_id)
+      client.storage_usage(owner_type(owner), owner.id, from, to)
+    end
+
+    def current_storage(owner, user_id)
+      Models::CustomImageStorage.where(owner_type: owner_type(owner), owner_id: owner.id).order('id desc').limit(1).first
+    end
+
+    def storage_executions_usage(owner, user_id)
+      BillingClient.new(user_id).storage_executions_usage(owner_type(owner), owner.id)
+    end
+
+    private
+
+    def owner_type(owner)
+      owner.vcs_type =~ /User/ ? 'User' : 'Organization'
+    end
+  end
+end

lib/travis/api/v3/renderer/custom_image.rb

@@ -0,0 +1,22 @@
+module Travis::API::V3
+  class Renderer::CustomImage < ModelRenderer
+    representation :minimal, :id, :owner_id, :owner_type, :name, :usage, :created_at, :updated_at, :os_version,
+                   :created_by, :private, :size_bytes
+    representation :standard, *representations[:minimal]
+
+    def created_by
+      return nil unless user = model.created_by
+
+      {
+        '@type' => 'user',
+        '@href' => "/v3/user/#{user.id}",
+        '@representation' => 'minimal'.freeze,
+        'id' => user.id,
+        'login' => user.login,
+        'name' => user.name,
+      }.tap do |data|
+        data['avatar_url'] = user.avatar_url if user.email.present?
+      end
+    end
+  end
+end

lib/travis/api/v3/renderer/custom_image_storage.rb

@@ -0,0 +1,6 @@
+module Travis::API::V3
+  class Renderer::CustomImageStorage < ModelRenderer
+    representation :minimal, :id, :owner_id, :owner_type, :current_aggregated_storage, :created_at, :updated_at, :end_date
+    representation :standard, *representations[:minimal]
+  end
+end

lib/travis/api/v3/renderer/custom_image_usage.rb

@@ -0,0 +1,7 @@
+module Travis::API::V3
+  class Renderer::CustomImagesUsage < ModelRenderer
+    representation :minimal, :total_usage, :excess_usage, :free_usage, :quantity_limit_free, :quantity_limit_type,
+                   :quantity_limit_charge
+    representation :standard, *representations[:minimal]
+  end
+end

lib/travis/api/v3/renderer/custom_image_usages.rb

@@ -0,0 +1,6 @@
+module Travis::API::V3
+  class Renderer::CustomImagesUsages < CollectionRenderer
+    type           :custom_images_usages
+    collection_key :custom_images_usages
+  end
+end

lib/travis/api/v3/renderer/custom_images.rb

@@ -0,0 +1,6 @@
+module Travis::API::V3
+  class Renderer::CustomImages < CollectionRenderer
+    type           :custom_images
+    collection_key :custom_images
+  end
+end

lib/travis/api/v3/renderer/membership.rb

@@ -0,0 +1,6 @@
+module Travis::API::V3
+  class Renderer::Membership < ModelRenderer
+    representation(:minimal, :organization_id, :build_permission)
+    representation(:standard, :user_id, :organization_id, :role, :build_permission)
+  end
+end

lib/travis/api/v3/renderer/storage_executions_usage.rb

@@ -0,0 +1,6 @@
+module Travis::API::V3
+  class Renderer::StorageExecutionsUsage < ModelRenderer
+    representation :minimal, :estimated_usage
+    representation :standard, *representations[:minimal]
+  end
+end

lib/travis/api/v3/renderer/storage_executions_usages.rb

@@ -0,0 +1,6 @@
+module Travis::API::V3
+  class Renderer::StorageExecutionsUsages < CollectionRenderer
+    type           :storage_executions_usages
+    collection_key :storage_executions_usages
+  end
+end

lib/travis/api/v3/renderer/user.rb

@@ -2,7 +2,7 @@
 
 module Travis::API::V3
   class Renderer::User < Renderer::Owner
-    representation(:standard, :email, :is_syncing, :synced_at, :recently_signed_up, :secure_user_hash, :ro_mode, :confirmed_at, :custom_keys, :internal, :last_activity_at)
+    representation(:standard, :email, :is_syncing, :synced_at, :recently_signed_up, :secure_user_hash, :ro_mode, :confirmed_at, :custom_keys, :internal, :last_activity_at, :memberships)
     representation(:additional, :emails, :collaborator)
 
     def email

lib/travis/api/v3/renderer/v2_addon_usage.rb

@@ -1,6 +1,8 @@
 module Travis::API::V3
   class Renderer::V2AddonUsage < ModelRenderer
-    representation(:standard, :id, :addon_id, :addon_quantity, :addon_usage, :remaining, :purchase_date, :valid_to, :active, :status)
-    representation(:minimal, :id, :addon_id, :addon_quantity, :addon_usage, :remaining, :purchase_date, :valid_to, :active, :status)
+    representation(:standard, :id, :addon_id, :addon_quantity, :addon_usage, :remaining, :purchase_date, :valid_to,
+                   :active, :status, :total_usage, :quantity_limit_free, :quantity_limit_type, :quantity_limit_charge)
+    representation(:minimal, :id, :addon_id, :addon_quantity, :addon_usage, :remaining, :purchase_date, :valid_to,
+                   :active, :status, :total_usage, :quantity_limit_free, :quantity_limit_type, :quantity_limit_charge)
   end
 end