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