fix: generate firewall script

Author: cholical

internal/shadeform/v1/instance.go

@@ -77,7 +77,7 @@ func (c *ShadeformClient) CreateInstance(ctx context.Context, attrs v1.CreateIns
 		tags = append(tags, createdTag)
 	}
 
-	base64Script, err := c.generateFirewallScript(attrs.FirewallRules)
+	base64Script, err := c.GenerateFirewallScript(attrs.FirewallRules)
 	if err != nil {
 		return nil, err
 	}

internal/shadeform/v1/ufw.go

@@ -15,7 +15,7 @@ const (
 	ufwForceEnable          = "ufw --force enable"
 )
 
-func (c *ShadeformClient) generateFirewallScript(firewallRules v1.FirewallRules) (string, error) {
+func (c *ShadeformClient) GenerateFirewallScript(firewallRules v1.FirewallRules) (string, error) {
 	commands := []string{ufwForceReset, ufwDefaultDropIncoming, ufwDefaultAllowOutgoing, ufwDefaultAllowPort22, ufwDefaultAllowPort2222}
 
 	for _, firewallRule := range firewallRules.IngressRules {
@@ -40,38 +40,50 @@ func (c *ShadeformClient) generateFirewallScript(firewallRules v1.FirewallRules)
 
 func (c *ShadeformClient) convertIngressFirewallRuleToUfwCommand(firewallRule v1.FirewallRule) []string {
 	cmds := []string{}
-	portSpec := ""
+	portSpecs := []string{}
 	if firewallRule.FromPort == firewallRule.ToPort {
-		portSpec = fmt.Sprintf("port %d", firewallRule.FromPort)
+		portSpecs = append(portSpecs, fmt.Sprintf("port %d", firewallRule.FromPort))
 	} else {
-		portSpec = fmt.Sprintf("port %d:%d", firewallRule.FromPort, firewallRule.ToPort)
+		// port ranges require two separate rules for tcp and udp
+		portSpecs = append(portSpecs, fmt.Sprintf("port %d:%d proto tcp", firewallRule.FromPort, firewallRule.ToPort))
+		portSpecs = append(portSpecs, fmt.Sprintf("port %d:%d proto udp", firewallRule.FromPort, firewallRule.ToPort))
 	}
 
 	if len(firewallRule.IPRanges) == 0 {
-		cmds = append(cmds, fmt.Sprintf("ufw allow in from any to any port %s", portSpec))
+		for _, portSpec := range portSpecs {
+			cmds = append(cmds, fmt.Sprintf("ufw allow in from any to any %s", portSpec))
+		}
 	}
 
 	for _, ipRange := range firewallRule.IPRanges {
-		cmds = append(cmds, fmt.Sprintf("ufw allow in from %s to any port %s", ipRange, portSpec))
+		for _, portSpec := range portSpecs {
+			cmds = append(cmds, fmt.Sprintf("ufw allow in from %s to any %s", ipRange, portSpec))
+		}
 	}
 	return cmds
 }
 
 func (c *ShadeformClient) convertEgressFirewallRuleToUfwCommand(firewallRule v1.FirewallRule) []string {
 	cmds := []string{}
-	portSpec := ""
+	portSpecs := []string{}
 	if firewallRule.FromPort == firewallRule.ToPort {
-		portSpec = fmt.Sprintf("port %d", firewallRule.FromPort)
+		portSpecs = append(portSpecs, fmt.Sprintf("port %d", firewallRule.FromPort))
 	} else {
-		portSpec = fmt.Sprintf("port %d:%d", firewallRule.FromPort, firewallRule.ToPort)
+		// port ranges require two separate rules for tcp and udp
+		portSpecs = append(portSpecs, fmt.Sprintf("port %d:%d proto tcp", firewallRule.FromPort, firewallRule.ToPort))
+		portSpecs = append(portSpecs, fmt.Sprintf("port %d:%d proto udp", firewallRule.FromPort, firewallRule.ToPort))
 	}
 
 	if len(firewallRule.IPRanges) == 0 {
-		cmds = append(cmds, fmt.Sprintf("ufw allow out to any port %s", portSpec))
+		for _, portSpec := range portSpecs {
+			cmds = append(cmds, fmt.Sprintf("ufw allow out to any %s", portSpec))
+		}
 	}
 
 	for _, ipRange := range firewallRule.IPRanges {
-		cmds = append(cmds, fmt.Sprintf("ufw allow out to %s port %s", ipRange, portSpec))
+		for _, portSpec := range portSpecs {
+			cmds = append(cmds, fmt.Sprintf("ufw allow out to %s %s", ipRange, portSpec))
+		}
 	}
 	return cmds
 }

internal/shadeform/v1/validation_test.go

@@ -74,18 +74,30 @@ func TestInstanceTypeFilter(t *testing.T) {
 		FirewallRules: v1.FirewallRules{
 			EgressRules: []v1.FirewallRule{
 				{
-					ID:       "test-rule",
+					ID:       "test-rule1",
 					FromPort: 80,
 					ToPort:   8080,
-					IPRanges: []string{"127.0.0.1", "10.0.0.0/32"},
+					IPRanges: []string{"127.0.0.1", "10.0.0.0/24"},
+				},
+				{
+					ID:       "test-rule2",
+					FromPort: 5432,
+					ToPort:   5432,
+					IPRanges: []string{"127.0.0.1", "10.0.0.0/24"},
 				},
 			},
 			IngressRules: []v1.FirewallRule{
 				{
-					ID:       "test-rule",
+					ID:       "test-rule3",
+					FromPort: 80,
+					ToPort:   8080,
+					IPRanges: []string{"127.0.0.1", "10.0.0.0/24"},
+				},
+				{
+					ID:       "test-rule4",
 					FromPort: 5432,
 					ToPort:   5432,
-					IPRanges: []string{"127.0.0.1", "10.0.0.0/32"},
+					IPRanges: []string{"127.0.0.1", "10.0.0.0/24"},
 				},
 			},
 		},