darkwing/server/api/

start_controller.rs

use axum::extract::{Json, Path};
use axum::routing::{get, post};
use axum::Router;
use tracing::debug;

use crate::measure_time;
use crate::server::dtos::start_dto::{
  DatadirResponse, PreliminaryData, PreliminaryResponse, StartRequest,
  StartResponse,
};
use crate::server::error::AppResult;
use crate::server::error::Error;
use crate::server::extractors::RequiredAuthentication;
use crate::server::extractors::ValidationExtractor;
use crate::server::services::browser_profile_services::args::Args;
use crate::server::services::browser_profile_services::{
  BrowserProfileService, BrowserProfileServiceTrait,
};
use crate::server::services::cloud_file_services::S3ClientType;
use crate::server::services::Services;

pub struct StartController;

impl StartController {
  pub fn app() -> Router {
    Router::new()
      .route("/", post(Self::main))
      .route("/preliminary/:browser_profile_id", get(Self::preliminary))
      .route("/datadir/:browser_profile_id", get(Self::datadir))
  }

  pub async fn main(
    RequiredAuthentication {
      user,
      team,
      services,
      token,
    }: RequiredAuthentication,
    ValidationExtractor(request): ValidationExtractor<StartRequest>,
  ) -> AppResult<Json<StartResponse>> {
    if !request.use_mock_profile {
      let access = measure_time!(
        "start_access_checks",
        services
          .browser_profile_service
          .check_access(user, request.browser_profile_id as u64)
          .await
      )?;

      if !access {
        return Err(Error::Forbidden);
      }
    } else if !services.users.does_plan_allows_scenarios(team.plan) {
      return Err(Error::PlanDoesNotAllowAccess);
    }

    let profile = match request.use_mock_profile {
      true => BrowserProfileService::get_mock_profile(request.clone().os),
      false => measure_time!(
        "start_get_profile",
        services
          .browser_profile_service
          .get_by_id_with_fingerprint(request.clone().browser_profile_id as i64)
          .await
      )?,
    };

    let config = measure_time!(
      "start_create_config",
      BrowserProfileService::create_config(&profile, request.clone(), token)
    )?;

    debug!("config: {:?}", config);

    let encrypted_config = measure_time!(
      "start_encrypt_config",
      services
        .config_encryption
        .encrypt_config(config.clone())
        .await
    )?;

    let storage_path = measure_time!(
      "start_generate_storage_path",
      services
        .browser_profile_service
        .generate_storage_path(profile.clone().into())
    )?;

    let patch_url = match team.is_fully_free_plan() {
      true => None,
      false => measure_time!(
        "start_maybe_prefetch_patch",
        maybe_prefetch_patch(
          &services,
          &storage_path,
          request.present_datadir_hash.clone(),
          profile.datadir_hash.clone(),
        )
        .await
      )?,
    };

    let datadir_url = match team.is_fully_free_plan() {
      true => String::new(),
      false => measure_time!(
        "start_generate_datadir_url",
        services
          .cloud_files
          .generate_presigned_url(&storage_path)
          .await
      )?,
    };

    let datadir_hash = profile.datadir_hash.clone();

    let command = measure_time!(
      "start_generate_command",
      Args::new(config, profile, request).output()
    );

    debug!("command: {:?}", command);

    let response = StartResponse {
      hash: datadir_hash,
      datadir_url,
      patch_url,
      config: encrypted_config,
      command,
    };

    Ok(Json(response))
  }

  pub async fn preliminary(
    RequiredAuthentication {
      user,
      team,
      services,
      token: _,
    }: RequiredAuthentication,
    Path(browser_profile_id): Path<u64>,
  ) -> AppResult<Json<PreliminaryResponse>> {
    if browser_profile_id == 0 {
      if !services.users.does_plan_allows_scenarios(team.plan) {
        return Err(Error::PlanDoesNotAllowAccess);
      }

      return Ok(Json(PreliminaryResponse::from(PreliminaryData::get_mock())));
    } else {
      let access = measure_time!(
        "preliminary_check_access",
        services
          .browser_profile_service
          .check_access(user.clone(), browser_profile_id)
          .await
      )?;

      if !access {
        return Err(Error::Forbidden);
      }
    }

    let data = measure_time!(
      "preliminary_get_preliminary_data",
      services
        .browser_profile_service
        .get_preliminary_data(
          user.id as i64,
          team.id as i64,
          browser_profile_id as i64,
        )
        .await
    )?;

    let response = PreliminaryResponse::from(data);

    Ok(Json(response))
  }

  pub async fn datadir(
    RequiredAuthentication {
      user,
      team,
      services,
      token: _,
    }: RequiredAuthentication,
    Path(browser_profile_id): Path<u64>,
  ) -> AppResult<Json<DatadirResponse>> {
    let access = measure_time!(
      "datadir_check_access",
      services
        .browser_profile_service
        .check_access(user, browser_profile_id)
        .await
    )?;

    if !access {
      return Err(Error::Forbidden);
    }

    let browser_profile = measure_time!(
      "datadir_get_mini_profile",
      services
        .browser_profile_service
        .get_mini_by_id(browser_profile_id as i64)
        .await
    )?;

    let storage_path = measure_time!(
      "datadir_generate_storage_path",
      services
        .browser_profile_service
        .generate_storage_path(browser_profile.clone())
    )?;

    let datadir_url = match team.is_fully_free_plan() {
      true => String::new(),
      false => measure_time!(
        "datadir_generate_presigned_url",
        services
          .cloud_files
          .generate_presigned_url(&storage_path)
          .await
      )?,
    };

    let response = DatadirResponse {
      hash: browser_profile.datadir_hash,
      url: datadir_url,
    };

    Ok(Json(response))
  }
}

async fn maybe_prefetch_patch(
  services: &Services,
  storage_path: &str,
  present_datadir_hash: Option<String>,
  datadir_hash: Option<String>,
) -> AppResult<Option<String>> {
  match (present_datadir_hash, datadir_hash) {
    (Some(present_datadir_hash), Some(_datadir_hash)) => {
      let patch_key =
        format!("{}/{}.patch", storage_path, present_datadir_hash);
      if services
        .cloud_files
        .does_file_exist(&patch_key, S3ClientType::Primary)
        .await?
      {
        metrics::counter!("darkwing_start_requests_total", "type" => "raw_plus_patch")
          .increment(1);
        Ok(Some(
          services
            .cloud_files
            .generate_presigned_url(&patch_key)
            .await?,
        ))
      } else {
        metrics::counter!("darkwing_start_requests_total", "type" => "raw")
          .increment(1);
        Ok(None)
      }
    }
    _ => {
      metrics::counter!("darkwing_start_requests_total", "type" => "raw")
        .increment(1);
      Ok(None)
    }
  }
}