1use std::any::Any;
2use std::collections::HashMap;
3use std::fmt::Debug;
4use std::sync::{Arc, Mutex, OnceLock};
5
6use anyhow::Result;
7use nanoid::nanoid;
8use serde_json::json;
9
10use super::terraform::{TERRAFORM_ALPHABET, TerraformOutput, TerraformProvider};
11use super::{ClientStrategy, Host, HostTargetType, LaunchedHost, ResourceBatch, ResourceResult};
12use crate::ssh::LaunchedSshHost;
13use crate::{BaseServerStrategy, HostStrategyGetter, PortNetworkHint};
14
15pub struct LaunchedEc2Instance {
16 resource_result: Arc<ResourceResult>,
17 user: String,
18 pub internal_ip: String,
19 pub external_ip: Option<String>,
20}
21
22impl LaunchedSshHost for LaunchedEc2Instance {
23 fn get_external_ip(&self) -> Option<String> {
24 self.external_ip.clone()
25 }
26
27 fn get_internal_ip(&self) -> String {
28 self.internal_ip.clone()
29 }
30
31 fn get_cloud_provider(&self) -> String {
32 "AWS".to_string()
33 }
34
35 fn resource_result(&self) -> &Arc<ResourceResult> {
36 &self.resource_result
37 }
38
39 fn ssh_user(&self) -> &str {
40 self.user.as_str()
41 }
42}
43
44#[derive(Debug)]
45pub struct AwsNetwork {
46 pub region: String,
47 pub existing_vpc: OnceLock<String>,
48 id: String,
49}
50
51impl AwsNetwork {
52 pub fn new(region: impl Into<String>, existing_vpc: Option<String>) -> Arc<Self> {
53 Arc::new(Self {
54 region: region.into(),
55 existing_vpc: existing_vpc.map(From::from).unwrap_or_default(),
56 id: nanoid!(8, &TERRAFORM_ALPHABET),
57 })
58 }
59
60 fn collect_resources(&self, resource_batch: &mut ResourceBatch) -> String {
61 resource_batch
62 .terraform
63 .terraform
64 .required_providers
65 .insert(
66 "aws".to_string(),
67 TerraformProvider {
68 source: "hashicorp/aws".to_string(),
69 version: "5.0.0".to_string(),
70 },
71 );
72
73 resource_batch.terraform.provider.insert(
74 "aws".to_string(),
75 json!({
76 "region": self.region
77 }),
78 );
79
80 let vpc_network = format!("hydro-vpc-network-{}", self.id);
81
82 if let Some(existing) = self.existing_vpc.get() {
83 if resource_batch
84 .terraform
85 .resource
86 .get("aws_vpc")
87 .unwrap_or(&HashMap::new())
88 .contains_key(existing)
89 {
90 format!("aws_vpc.{existing}")
91 } else {
92 resource_batch
93 .terraform
94 .data
95 .entry("aws_vpc".to_string())
96 .or_default()
97 .insert(
98 vpc_network.clone(),
99 json!({
100 "id": existing,
101 }),
102 );
103
104 format!("data.aws_vpc.{vpc_network}")
105 }
106 } else {
107 resource_batch
108 .terraform
109 .resource
110 .entry("aws_vpc".to_string())
111 .or_default()
112 .insert(
113 vpc_network.clone(),
114 json!({
115 "cidr_block": "10.0.0.0/16",
116 "enable_dns_hostnames": true,
117 "enable_dns_support": true,
118 "tags": {
119 "Name": vpc_network
120 }
121 }),
122 );
123
124 let igw_key = format!("{vpc_network}-igw");
126 resource_batch
127 .terraform
128 .resource
129 .entry("aws_internet_gateway".to_string())
130 .or_default()
131 .insert(
132 igw_key.clone(),
133 json!({
134 "vpc_id": format!("${{aws_vpc.{}.id}}", vpc_network),
135 "tags": {
136 "Name": igw_key
137 }
138 }),
139 );
140
141 let subnet_key = format!("{vpc_network}-subnet");
143 resource_batch
144 .terraform
145 .resource
146 .entry("aws_subnet".to_string())
147 .or_default()
148 .insert(
149 subnet_key.clone(),
150 json!({
151 "vpc_id": format!("${{aws_vpc.{}.id}}", vpc_network),
152 "cidr_block": "10.0.1.0/24",
153 "availability_zone": format!("{}a", self.region),
154 "map_public_ip_on_launch": true,
155 "tags": {
156 "Name": subnet_key
157 }
158 }),
159 );
160
161 let rt_key = format!("{vpc_network}-rt");
163 resource_batch
164 .terraform
165 .resource
166 .entry("aws_route_table".to_string())
167 .or_default()
168 .insert(
169 rt_key.clone(),
170 json!({
171 "vpc_id": format!("${{aws_vpc.{}.id}}", vpc_network),
172 "tags": {
173 "Name": rt_key
174 }
175 }),
176 );
177
178 resource_batch
180 .terraform
181 .resource
182 .entry("aws_route".to_string())
183 .or_default()
184 .insert(
185 format!("{vpc_network}-route"),
186 json!({
187 "route_table_id": format!("${{aws_route_table.{}.id}}", rt_key),
188 "destination_cidr_block": "0.0.0.0/0",
189 "gateway_id": format!("${{aws_internet_gateway.{}.id}}", igw_key)
190 }),
191 );
192
193 resource_batch
194 .terraform
195 .resource
196 .entry("aws_route_table_association".to_string())
197 .or_default()
198 .insert(
199 format!("{vpc_network}-rta"),
200 json!({
201 "subnet_id": format!("${{aws_subnet.{}.id}}", subnet_key),
202 "route_table_id": format!("${{aws_route_table.{}.id}}", rt_key)
203 }),
204 );
205
206 let sg_key = format!("{vpc_network}-default-sg");
208 resource_batch
209 .terraform
210 .resource
211 .entry("aws_security_group".to_string())
212 .or_default()
213 .insert(
214 sg_key.clone(),
215 json!({
216 "name": format!("{vpc_network}-default-allow-internal"),
217 "description": "Allow internal communication between instances",
218 "vpc_id": format!("${{aws_vpc.{}.id}}", vpc_network),
219 "ingress": [
220 {
221 "from_port": 0,
222 "to_port": 65535,
223 "protocol": "tcp",
224 "cidr_blocks": ["10.0.0.0/16"],
225 "description": "Allow all TCP traffic within VPC",
226 "ipv6_cidr_blocks": [],
227 "prefix_list_ids": [],
228 "security_groups": [],
229 "self": false
230 },
231 {
232 "from_port": 0,
233 "to_port": 65535,
234 "protocol": "udp",
235 "cidr_blocks": ["10.0.0.0/16"],
236 "description": "Allow all UDP traffic within VPC",
237 "ipv6_cidr_blocks": [],
238 "prefix_list_ids": [],
239 "security_groups": [],
240 "self": false
241 },
242 {
243 "from_port": -1,
244 "to_port": -1,
245 "protocol": "icmp",
246 "cidr_blocks": ["10.0.0.0/16"],
247 "description": "Allow ICMP within VPC",
248 "ipv6_cidr_blocks": [],
249 "prefix_list_ids": [],
250 "security_groups": [],
251 "self": false
252 }
253 ],
254 "egress": [
255 {
256 "from_port": 0,
257 "to_port": 0,
258 "protocol": "-1",
259 "cidr_blocks": ["0.0.0.0/0"],
260 "description": "Allow all outbound traffic",
261 "ipv6_cidr_blocks": [],
262 "prefix_list_ids": [],
263 "security_groups": [],
264 "self": false
265 }
266 ]
267 }),
268 );
269
270 let out = format!("aws_vpc.{vpc_network}");
271 self.existing_vpc.set(vpc_network).unwrap();
272 out
273 }
274 }
275}
276
277pub struct AwsEc2Host {
278 id: usize,
280
281 region: String,
282 instance_type: String,
283 target_type: HostTargetType,
284 ami: String,
285 network: Arc<AwsNetwork>,
286 user: Option<String>,
287 display_name: Option<String>,
288 pub launched: OnceLock<Arc<LaunchedEc2Instance>>,
289 external_ports: Mutex<Vec<u16>>,
290}
291
292impl Debug for AwsEc2Host {
293 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
294 f.write_fmt(format_args!(
295 "AwsEc2Host({} ({:?}))",
296 self.id, &self.display_name
297 ))
298 }
299}
300
301impl AwsEc2Host {
302 #[expect(clippy::too_many_arguments, reason = "used via builder pattern")]
303 pub fn new(
304 id: usize,
305 region: impl Into<String>,
306 instance_type: impl Into<String>,
307 target_type: HostTargetType,
308 ami: impl Into<String>,
309 network: Arc<AwsNetwork>,
310 user: Option<String>,
311 display_name: Option<String>,
312 ) -> Self {
313 Self {
314 id,
315 region: region.into(),
316 instance_type: instance_type.into(),
317 target_type,
318 ami: ami.into(),
319 network,
320 user,
321 display_name,
322 launched: OnceLock::new(),
323 external_ports: Mutex::new(Vec::new()),
324 }
325 }
326}
327
328impl Host for AwsEc2Host {
329 fn target_type(&self) -> HostTargetType {
330 self.target_type
331 }
332
333 fn request_port_base(&self, bind_type: &BaseServerStrategy) {
334 match bind_type {
335 BaseServerStrategy::UnixSocket => {}
336 BaseServerStrategy::InternalTcpPort(_) => {}
337 BaseServerStrategy::ExternalTcpPort(port) => {
338 let mut external_ports = self.external_ports.lock().unwrap();
339 if !external_ports.contains(port) {
340 if self.launched.get().is_some() {
341 todo!("Cannot adjust security group after host has been launched");
342 }
343 external_ports.push(*port);
344 }
345 }
346 }
347 }
348
349 fn request_custom_binary(&self) {
350 self.request_port_base(&BaseServerStrategy::ExternalTcpPort(22));
351 }
352
353 fn id(&self) -> usize {
354 self.id
355 }
356
357 fn collect_resources(&self, resource_batch: &mut ResourceBatch) {
358 if self.launched.get().is_some() {
359 return;
360 }
361
362 let vpc_path = self.network.collect_resources(resource_batch);
363
364 resource_batch
366 .terraform
367 .terraform
368 .required_providers
369 .insert(
370 "local".to_string(),
371 TerraformProvider {
372 source: "hashicorp/local".to_string(),
373 version: "2.3.0".to_string(),
374 },
375 );
376
377 resource_batch
378 .terraform
379 .terraform
380 .required_providers
381 .insert(
382 "tls".to_string(),
383 TerraformProvider {
384 source: "hashicorp/tls".to_string(),
385 version: "4.0.4".to_string(),
386 },
387 );
388
389 resource_batch
391 .terraform
392 .resource
393 .entry("tls_private_key".to_string())
394 .or_default()
395 .insert(
396 "vm_instance_ssh_key".to_string(),
397 json!({
398 "algorithm": "RSA",
399 "rsa_bits": 4096
400 }),
401 );
402
403 resource_batch
404 .terraform
405 .resource
406 .entry("local_file".to_string())
407 .or_default()
408 .insert(
409 "vm_instance_ssh_key_pem".to_string(),
410 json!({
411 "content": "${tls_private_key.vm_instance_ssh_key.private_key_pem}",
412 "filename": ".ssh/vm_instance_ssh_key_pem",
413 "file_permission": "0600",
414 "directory_permission": "0700"
415 }),
416 );
417
418 resource_batch
419 .terraform
420 .resource
421 .entry("aws_key_pair".to_string())
422 .or_default()
423 .insert(
424 "ec2_key_pair".to_string(),
425 json!({
426 "key_name": format!("hydro-key-{}", nanoid!(8, &TERRAFORM_ALPHABET)),
427 "public_key": "${tls_private_key.vm_instance_ssh_key.public_key_openssh}"
428 }),
429 );
430
431 let instance_key = format!("ec2-instance-{}", self.id);
432 let mut instance_name = format!("hydro-ec2-instance-{}", nanoid!(8, &TERRAFORM_ALPHABET));
433
434 if let Some(mut display_name) = self.display_name.clone() {
435 instance_name.push('-');
436 display_name = display_name.replace("_", "-").to_lowercase();
437
438 let num_chars_to_cut = instance_name.len() + display_name.len() - 63;
439 if num_chars_to_cut > 0 {
440 display_name.drain(0..num_chars_to_cut);
441 }
442 instance_name.push_str(&display_name);
443 }
444
445 let network_id = self.network.id.clone();
446 let vpc_ref = format!("${{{}.id}}", vpc_path);
447 let subnet_ref = format!("${{aws_subnet.hydro-vpc-network-{}-subnet.id}}", network_id);
448 let default_sg_ref = format!(
449 "${{aws_security_group.hydro-vpc-network-{}-default-sg.id}}",
450 network_id
451 );
452
453 let mut security_groups = vec![default_sg_ref.clone()];
455 let external_ports = self.external_ports.lock().unwrap();
456
457 if !external_ports.is_empty() {
458 let sg_key = format!("sg-{}", self.id);
459 let mut sg_rules = vec![];
460
461 for port in external_ports.iter() {
462 sg_rules.push(json!({
463 "from_port": port,
464 "to_port": port,
465 "protocol": "tcp",
466 "cidr_blocks": ["0.0.0.0/0"],
467 "description": format!("External port {}", port),
468 "ipv6_cidr_blocks": [],
469 "prefix_list_ids": [],
470 "security_groups": [],
471 "self": false
472 }));
473 }
474
475 resource_batch
476 .terraform
477 .resource
478 .entry("aws_security_group".to_string())
479 .or_default()
480 .insert(
481 sg_key.clone(),
482 json!({
483 "name": format!("hydro-sg-{}", nanoid!(8, &TERRAFORM_ALPHABET)),
484 "description": "Hydro external ports security group",
485 "vpc_id": vpc_ref,
486 "ingress": sg_rules,
487 "egress": [{
488 "from_port": 0,
489 "to_port": 0,
490 "protocol": "-1",
491 "cidr_blocks": ["0.0.0.0/0"],
492 "description": "All outbound traffic",
493 "ipv6_cidr_blocks": [],
494 "prefix_list_ids": [],
495 "security_groups": [],
496 "self": false
497 }]
498 }),
499 );
500
501 security_groups.push(format!("${{aws_security_group.{}.id}}", sg_key));
502 }
503 drop(external_ports);
504
505 resource_batch
507 .terraform
508 .resource
509 .entry("aws_instance".to_string())
510 .or_default()
511 .insert(
512 instance_key.clone(),
513 json!({
514 "ami": self.ami,
515 "instance_type": self.instance_type,
516 "key_name": "${aws_key_pair.ec2_key_pair.key_name}",
517 "vpc_security_group_ids": security_groups,
518 "subnet_id": subnet_ref,
519 "associate_public_ip_address": true,
520 "tags": {
521 "Name": instance_name
522 }
523 }),
524 );
525
526 resource_batch.terraform.output.insert(
527 format!("{}-private-ip", instance_key),
528 TerraformOutput {
529 value: format!("${{aws_instance.{}.private_ip}}", instance_key),
530 },
531 );
532
533 resource_batch.terraform.output.insert(
534 format!("{}-public-ip", instance_key),
535 TerraformOutput {
536 value: format!("${{aws_instance.{}.public_ip}}", instance_key),
537 },
538 );
539 }
540
541 fn launched(&self) -> Option<Arc<dyn LaunchedHost>> {
542 self.launched
543 .get()
544 .map(|a| a.clone() as Arc<dyn LaunchedHost>)
545 }
546
547 fn provision(&self, resource_result: &Arc<ResourceResult>) -> Arc<dyn LaunchedHost> {
548 self.launched
549 .get_or_init(|| {
550 let id = self.id;
551
552 let internal_ip = resource_result
553 .terraform
554 .outputs
555 .get(&format!("ec2-instance-{id}-private-ip"))
556 .unwrap()
557 .value
558 .clone();
559
560 let external_ip = resource_result
561 .terraform
562 .outputs
563 .get(&format!("ec2-instance-{id}-public-ip"))
564 .map(|v| v.value.clone());
565
566 Arc::new(LaunchedEc2Instance {
567 resource_result: resource_result.clone(),
568 user: self
569 .user
570 .as_ref()
571 .cloned()
572 .unwrap_or("ec2-user".to_string()),
573 internal_ip,
574 external_ip,
575 })
576 })
577 .clone()
578 }
579
580 fn strategy_as_server<'a>(
581 &'a self,
582 client_host: &dyn Host,
583 network_hint: PortNetworkHint,
584 ) -> Result<(ClientStrategy<'a>, HostStrategyGetter)> {
585 if matches!(network_hint, PortNetworkHint::Auto)
586 && client_host.can_connect_to(ClientStrategy::UnixSocket(self.id))
587 {
588 Ok((
589 ClientStrategy::UnixSocket(self.id),
590 Box::new(|_| BaseServerStrategy::UnixSocket),
591 ))
592 } else if matches!(
593 network_hint,
594 PortNetworkHint::Auto | PortNetworkHint::TcpPort(_)
595 ) && client_host.can_connect_to(ClientStrategy::InternalTcpPort(self))
596 {
597 Ok((
598 ClientStrategy::InternalTcpPort(self),
599 Box::new(move |_| {
600 BaseServerStrategy::InternalTcpPort(match network_hint {
601 PortNetworkHint::Auto => None,
602 PortNetworkHint::TcpPort(port) => port,
603 })
604 }),
605 ))
606 } else if matches!(network_hint, PortNetworkHint::Auto)
607 && client_host.can_connect_to(ClientStrategy::ForwardedTcpPort(self))
608 {
609 Ok((
610 ClientStrategy::ForwardedTcpPort(self),
611 Box::new(|me| {
612 me.downcast_ref::<AwsEc2Host>()
613 .unwrap()
614 .request_port_base(&BaseServerStrategy::ExternalTcpPort(22));
615 BaseServerStrategy::InternalTcpPort(None)
616 }),
617 ))
618 } else {
619 anyhow::bail!("Could not find a strategy to connect to AWS EC2 instance")
620 }
621 }
622
623 fn can_connect_to(&self, typ: ClientStrategy) -> bool {
624 match typ {
625 ClientStrategy::UnixSocket(id) => {
626 #[cfg(unix)]
627 {
628 self.id == id
629 }
630
631 #[cfg(not(unix))]
632 {
633 let _ = id;
634 false
635 }
636 }
637 ClientStrategy::InternalTcpPort(target_host) => {
638 if let Some(aws_target) = <dyn Any>::downcast_ref::<AwsEc2Host>(target_host) {
639 self.region == aws_target.region
640 && Arc::ptr_eq(&self.network, &aws_target.network)
641 } else {
642 false
643 }
644 }
645 ClientStrategy::ForwardedTcpPort(_) => false,
646 }
647 }
648}