From 8215389604982fd7635f412a4f56a4f9e85091f8 Mon Sep 17 00:00:00 2001
From: Chintan Raval <ChintanRaval@users.noreply.github.com>
Date: Wed, 3 Oct 2018 20:21:51 +1000
Subject: [PATCH] feat(aws-ec2): support UDP port ranges in SecurityGroups
 (#835)

Add support for UDP to ec2.SecurityGroupRule
---
 .../aws-ec2/lib/security-group-rule.ts        | 87 ++++++++++++++++++-
 .../@aws-cdk/aws-ec2/test/test.connections.ts |  8 +-
 2 files changed, 92 insertions(+), 3 deletions(-)

diff --git a/packages/@aws-cdk/aws-ec2/lib/security-group-rule.ts b/packages/@aws-cdk/aws-ec2/lib/security-group-rule.ts
index 9b707fd82f29b..c0541da365afa 100644
--- a/packages/@aws-cdk/aws-ec2/lib/security-group-rule.ts
+++ b/packages/@aws-cdk/aws-ec2/lib/security-group-rule.ts
@@ -234,7 +234,92 @@ export class TcpAllPorts implements IPortRange {
 }
 
 /**
- * All TCP Ports
+ * A single UDP port
+ */
+export class UdpPort implements IPortRange {
+  public readonly canInlineRule = true;
+
+  constructor(private readonly port: number) {
+  }
+
+  public toRuleJSON(): any {
+    return {
+      ipProtocol: Protocol.Udp,
+      fromPort: this.port,
+      toPort: this.port
+    };
+  }
+
+  public toString() {
+    return `UDP ${this.port}`;
+  }
+}
+
+/**
+ * A single UDP port that is provided by a resource attribute
+ */
+export class UdpPortFromAttribute implements IPortRange {
+  public readonly canInlineRule = false;
+
+  constructor(private readonly port: string) {
+  }
+
+  public toRuleJSON(): any {
+    return {
+      ipProtocol: Protocol.Udp,
+      fromPort: this.port,
+      toPort: this.port
+    };
+  }
+
+  public toString() {
+    return 'UDP {IndirectPort}';
+  }
+}
+
+/**
+ * A UDP port range
+ */
+export class UdpPortRange implements IPortRange {
+  public readonly canInlineRule = true;
+
+  constructor(private readonly startPort: number, private readonly endPort: number) {
+  }
+
+  public toRuleJSON(): any {
+    return {
+      ipProtocol: Protocol.Udp,
+      fromPort: this.startPort,
+      toPort: this.endPort
+    };
+  }
+
+  public toString() {
+    return `UDP ${this.startPort}-${this.endPort}`;
+  }
+}
+
+/**
+ * All UDP Ports
+ */
+export class UdpAllPorts implements IPortRange {
+  public readonly canInlineRule = true;
+
+  public toRuleJSON(): any {
+    return {
+      ipProtocol: Protocol.Udp,
+      fromPort: 0,
+      toPort: 65535
+    };
+  }
+
+  public toString() {
+    return 'UDP ALL PORTS';
+  }
+}
+
+/**
+ * All Traffic
  */
 export class AllConnections implements IPortRange {
   public readonly canInlineRule = true;
diff --git a/packages/@aws-cdk/aws-ec2/test/test.connections.ts b/packages/@aws-cdk/aws-ec2/test/test.connections.ts
index 34bbf485b06bb..de7087f425b64 100644
--- a/packages/@aws-cdk/aws-ec2/test/test.connections.ts
+++ b/packages/@aws-cdk/aws-ec2/test/test.connections.ts
@@ -2,7 +2,7 @@ import { expect, haveResource } from '@aws-cdk/assert';
 import { Stack } from '@aws-cdk/cdk';
 import { Test } from 'nodeunit';
 import { AllConnections, AnyIPv4, AnyIPv6, Connections, IConnectable, PrefixList, SecurityGroup, SecurityGroupRef,
-  TcpAllPorts, TcpPort, TcpPortFromAttribute, TcpPortRange, VpcNetwork } from '../lib';
+  TcpAllPorts, TcpPort, TcpPortFromAttribute, TcpPortRange, UdpAllPorts, UdpPort, UdpPortFromAttribute, UdpPortRange, VpcNetwork } from '../lib';
 
 export = {
   'peering between two security groups does not recursive infinitely'(test: Test) {
@@ -73,9 +73,13 @@ export = {
 
     const ports = [
       new TcpPort(1234),
-      new TcpPortFromAttribute("port!"),
+      new TcpPortFromAttribute("tcp-test-port!"),
       new TcpAllPorts(),
       new TcpPortRange(80, 90),
+      new UdpPort(2345),
+      new UdpPortFromAttribute("udp-test-port!"),
+      new UdpAllPorts(),
+      new UdpPortRange(85, 95),
       new AllConnections()
     ];