pr-preview/backend/src/services/ec2.ts

import {
  EC2Client,
  RunInstancesCommand,
  TerminateInstancesCommand,
  DescribeInstancesCommand,
  CreateSecurityGroupCommand,
  DeleteSecurityGroupCommand,
  AuthorizeSecurityGroupIngressCommand,
  DescribeSecurityGroupsCommand,
  CreateKeyPairCommand,
  DeleteKeyPairCommand,
  CreateTagsCommand,
} from "@aws-sdk/client-ec2";
import { STSClient, GetCallerIdentityCommand } from "@aws-sdk/client-sts";
import { createLogger } from "../lib/logger";
import { decrypt } from "../lib/encryption";
import type { User } from "@prisma/client";

const log = createLogger("EC2");

const UBUNTU_22_04_AMI: Record<string, string> = {
  "us-east-1": "ami-0e86e20dae9224db8",
  "us-east-2": "ami-0a0d9cf81c479446a",
  "us-west-1": "ami-05c969369880fa2c2",
  "us-west-2": "ami-03f8acd418785369b",
  "eu-west-1": "ami-0694d931cee176e7d",
  "eu-west-2": "ami-0f3d9639a5674d559",
  "eu-west-3": "ami-022e307f4b9e39f45",
  "eu-central-1": "ami-0faab6bdbac9486fb",
  "ap-southeast-1": "ami-0823c236601fef765",
  "ap-southeast-2": "ami-07620139298af599e",
  "ap-northeast-1": "ami-0b7546e839d7ace12",
  "ap-northeast-2": "ami-042e76978adeb8c48",
  "ap-south-1": "ami-076e3a557efe1aa9c",
  "sa-east-1": "ami-0eed58016fbe42de3",
  "ca-central-1": "ami-024f768de9e73d4f4",
  "eu-north-1": "ami-00381a880aa48c6c6",
  "me-south-1": "ami-09574f34b8dcd2eac",
  "af-south-1": "ami-08fdcf06b39fe83ec",
};

export function makeEc2Client(user: User): EC2Client {
  const accessKeyId = user.awsAccessKeyId ? decrypt(user.awsAccessKeyId) : "";
  const secretAccessKey = user.awsSecretAccessKey ? decrypt(user.awsSecretAccessKey) : "";
  return new EC2Client({
    region: user.awsRegion!,
    credentials: { accessKeyId, secretAccessKey },
  });
}

export function makeStsClient(user: User): STSClient {
  const accessKeyId = user.awsAccessKeyId ? decrypt(user.awsAccessKeyId) : "";
  const secretAccessKey = user.awsSecretAccessKey ? decrypt(user.awsSecretAccessKey) : "";
  return new STSClient({
    region: user.awsRegion!,
    credentials: { accessKeyId, secretAccessKey },
  });
}

export async function validateAwsCredentials(user: User): Promise<{ success: boolean; arn?: string; error?: string }> {
  try {
    const sts = makeStsClient(user);
    const res = await sts.send(new GetCallerIdentityCommand({}));
    return { success: true, arn: res.Arn };
  } catch (e: any) {
    return { success: false, error: e.message };
  }
}

export async function createKeyPairAndGetPrivateKey(ec2: EC2Client, keyName: string): Promise<{ privateKey: string }> {
  const res = await ec2.send(new CreateKeyPairCommand({
    KeyName: keyName,
    KeyType: "rsa",
    KeyFormat: "pem",
  }));
  return { privateKey: res.KeyMaterial! };
}

export async function createPreviewSecurityGroup(ec2: EC2Client, groupName: string, port: number): Promise<string> {
  const describe = await ec2.send(new DescribeSecurityGroupsCommand({
    Filters: [{ Name: "group-name", Values: [groupName] }],
  }));
  if (describe.SecurityGroups && describe.SecurityGroups.length > 0) {
    return describe.SecurityGroups[0].GroupId!;
  }

  const res = await ec2.send(new CreateSecurityGroupCommand({
    GroupName: groupName,
    Description: `PP Preview security group: ${groupName}`,
  }));
  const groupId = res.GroupId!;

  const ingress: any[] = [
    {
      IpProtocol: "tcp",
      FromPort: 22,
      ToPort: 22,
      IpRanges: [{ CidrIp: "0.0.0.0/0" }],
    },
  ];
  if (port !== 22) {
    ingress.push({
      IpProtocol: "tcp",
      FromPort: port,
      ToPort: port,
      IpRanges: [{ CidrIp: "0.0.0.0/0" }],
    });
  }

  await ec2.send(new AuthorizeSecurityGroupIngressCommand({
    GroupId: groupId,
    IpPermissions: ingress,
  }));
  return groupId;
}

const BOOTSTRAP_SCRIPT = `#!/bin/bash
set -e
apt-get update -y
apt-get install -y curl git unzip build-essential
curl -fsSL https://get.docker.com | sh
systemctl enable docker
systemctl start docker
apt-get install -y docker-compose-plugin
curl -o- https://raw.githubusercontent.com/nvm-sh/nvm/v0.39.7/install.sh | bash
export NVM_DIR="/root/.nvm"
source "$NVM_DIR/nvm.sh"
nvm install --lts
nvm alias default lts/*
`;

export async function launchInstance(opts: {
  ec2: EC2Client;
  region: string;
  instanceType: string;
  keyName: string;
  securityGroupId: string;
  tags: Record<string, string>;
}): Promise<string> {
  const ami = UBUNTU_22_04_AMI[opts.region] ?? UBUNTU_22_04_AMI["us-east-1"];
  const tagSpecs = Object.entries(opts.tags).map(([k, v]) => ({ Key: k, Value: v }));
  tagSpecs.push({ Key: "Name", Value: `pp-preview-${opts.tags["pp:previewId"]}` });

  const res = await opts.ec2.send(new RunInstancesCommand({
    ImageId: ami,
    InstanceType: opts.instanceType as any,
    MinCount: 1,
    MaxCount: 1,
    KeyName: opts.keyName,
    SecurityGroupIds: [opts.securityGroupId],
    UserData: Buffer.from(BOOTSTRAP_SCRIPT).toString("base64"),
    TagSpecifications: [
      { ResourceType: "instance", Tags: tagSpecs },
    ],
  }));

  return res.Instances![0].InstanceId!;
}

export async function waitForInstanceRunning(ec2: EC2Client, instanceId: string, maxWaitMs = 300_000): Promise<string> {
  const start = Date.now();
  while (Date.now() - start < maxWaitMs) {
    const res = await ec2.send(new DescribeInstancesCommand({
      InstanceIds: [instanceId],
    }));
    const inst = res.Reservations?.[0]?.Instances?.[0];
    if (inst?.State?.Name === "running" && inst.PublicIpAddress) {
      return inst.PublicIpAddress;
    }
    await sleep(5000);
  }
  throw new Error(`Instance ${instanceId} did not reach running state within timeout`);
}

export async function terminateInstance(ec2: EC2Client, instanceId: string): Promise<void> {
  await ec2.send(new TerminateInstancesCommand({ InstanceIds: [instanceId] }));
}

export async function deleteKeyPairAws(ec2: EC2Client, keyName: string): Promise<void> {
  try {
    await ec2.send(new DeleteKeyPairCommand({ KeyName: keyName }));
  } catch (e) {
    log.warn({ e, keyName }, "Failed to delete key pair");
  }
}

export async function deleteSecurityGroupAws(ec2: EC2Client, groupName: string): Promise<void> {
  try {
    const describe = await ec2.send(new DescribeSecurityGroupsCommand({
      Filters: [{ Name: "group-name", Values: [groupName] }],
    }));
    const groupId = describe.SecurityGroups?.[0]?.GroupId;
    if (groupId) {
      await sleep(5000);
      await ec2.send(new DeleteSecurityGroupCommand({ GroupId: groupId }));
    }
  } catch (e) {
    log.warn({ e, groupName }, "Failed to delete security group");
  }
}

export async function describeAllManagedInstances(ec2: EC2Client): Promise<any[]> {
  const res = await ec2.send(new DescribeInstancesCommand({
    Filters: [
      { Name: "tag:pp:managed", Values: ["true"] },
      { Name: "instance-state-name", Values: ["running", "pending", "stopping"] },
    ],
  }));
  return (res.Reservations ?? []).flatMap(r => r.Instances ?? []);
}

function sleep(ms: number) {
  return new Promise(r => setTimeout(r, ms));
}