From f8ff63e4d290d685a8b5710884daeb9141dcd38a Mon Sep 17 00:00:00 2001
From: Peter Cai <peter@typeblog.net>
Date: Tue, 24 Sep 2024 06:51:27 -0400
Subject: [PATCH] Add API calls for egress IPs (#104)

---
 resource_ip_addresses.go | 68 ++++++++++++++++++++++++++++++++++++++
 schema.graphql           | 70 ++++++++++++++++++++++++++++++++++++++++
 types.go                 | 24 ++++++++++++++
 3 files changed, 162 insertions(+)

diff --git a/resource_ip_addresses.go b/resource_ip_addresses.go
index 1d9f1b2..b4aa702 100644
--- a/resource_ip_addresses.go
+++ b/resource_ip_addresses.go
@@ -116,6 +116,74 @@ func (c *Client) AllocateSharedIPAddress(ctx context.Context, appName string) (n
 	return net.ParseIP(data.AllocateIPAddress.App.SharedIPAddress), nil
 }
 
+func (c *Client) AllocateEgressIPAddress(ctx context.Context, appName string, machineId string) (net.IP, net.IP, error) {
+	query := `
+		mutation($input: AllocateEgressIPAddressInput!) {
+			allocateEgressIpAddress(input: $input) {
+				v4,
+				v6
+			}
+		}
+	`
+
+	req := c.NewRequest(query)
+	ctx = ctxWithAction(ctx, "allocate_egress_ip_address")
+	req.Var("input", AllocateEgressIPAddressInput{AppID: appName, MachineID: machineId})
+
+	data, err := c.RunWithContext(ctx, req)
+	if err != nil {
+		return nil, nil, err
+	}
+
+	return net.ParseIP(data.AllocateEgressIPAddress.V4), net.ParseIP(data.AllocateEgressIPAddress.V6), nil
+}
+
+func (c *Client) GetEgressIPAddresses(ctx context.Context, appName string) (map[string][]EgressIPAddress, error) {
+	query := `
+		query ($appName: String!) {
+			app(name: $appName) {
+				machines {
+					nodes {
+						id
+						egressIpAddresses {
+							nodes {
+								id
+								ip
+								version
+								region
+							}
+						}
+					}
+				}
+			}
+		}
+	`
+
+	req := c.NewRequest(query)
+	ctx = ctxWithAction(ctx, "get_egress_ip_addresses")
+	req.Var("appName", appName)
+
+	data, err := c.RunWithContext(ctx, req)
+	if err != nil {
+		return nil, err
+	}
+
+	ret := make(map[string][]EgressIPAddress)
+	for _, m := range data.App.Machines.Nodes {
+		if m.EgressIpAddresses.Nodes == nil || len(m.EgressIpAddresses.Nodes) == 0 {
+			continue
+		}
+
+		ret[m.ID] = make([]EgressIPAddress, len(m.EgressIpAddresses.Nodes))
+
+		for i, ip := range m.EgressIpAddresses.Nodes {
+			ret[m.ID][i] = *ip
+		}
+	}
+
+	return ret, nil
+}
+
 func (c *Client) ReleaseIPAddress(ctx context.Context, appName string, ip string) error {
 	query := `
 		mutation($input: ReleaseIPAddressInput!) {
diff --git a/schema.graphql b/schema.graphql
index da41e76..de5f76e 100644
--- a/schema.graphql
+++ b/schema.graphql
@@ -1520,8 +1520,11 @@ scalar BigInt
 
 enum BillingStatus {
   CURRENT
+  DELINQUENT
   PAST_DUE
   SOURCE_REQUIRED
+  TRIAL_ACTIVE
+  TRIAL_ENDED
 }
 
 type Build implements Node {
@@ -5006,6 +5009,52 @@ type DummyWireGuardPeerPayload {
   pubkey: String!
 }
 
+type EgressIPAddress implements Node {
+  """
+  ID of the object.
+  """
+  id: ID!
+  ip: String!
+  region: String!
+  version: Int!
+}
+
+"""
+The connection type for EgressIPAddress.
+"""
+type EgressIPAddressConnection {
+  """
+  A list of edges.
+  """
+  edges: [EgressIPAddressEdge]
+
+  """
+  A list of nodes.
+  """
+  nodes: [EgressIPAddress]
+
+  """
+  Information to aid in pagination.
+  """
+  pageInfo: PageInfo!
+  totalCount: Int!
+}
+
+"""
+An edge in a connection.
+"""
+type EgressIPAddressEdge {
+  """
+  A cursor for use in pagination.
+  """
+  cursor: String!
+
+  """
+  The item at the end of the edge.
+  """
+  node: EgressIPAddress
+}
+
 type EmptyAppRole implements AppRole {
   """
   The name of this role
@@ -6321,6 +6370,27 @@ type Machine implements Node {
   app: App!
   config: JSON!
   createdAt: ISO8601DateTime!
+  egressIpAddresses(
+    """
+    Returns the elements in the list that come after the specified cursor.
+    """
+    after: String
+
+    """
+    Returns the elements in the list that come before the specified cursor.
+    """
+    before: String
+
+    """
+    Returns the first _n_ elements from the list.
+    """
+    first: Int
+
+    """
+    Returns the last _n_ elements from the list.
+    """
+    last: Int
+  ): EgressIPAddressConnection!
   events(
     """
     Returns the elements in the list that come after the specified cursor.
diff --git a/types.go b/types.go
index 97b2cb7..0aaf3e1 100644
--- a/types.go
+++ b/types.go
@@ -97,6 +97,10 @@ type Query struct {
 		App       App
 		IPAddress IPAddress
 	}
+	AllocateEgressIPAddress struct {
+		V4 string
+		V6 string
+	}
 	ReleaseIPAddress struct {
 		App App
 	}
@@ -276,6 +280,10 @@ type App struct {
 		Nodes []LimitedAccessToken
 	}
 
+	Machines struct {
+		Nodes []GqlMachine
+	}
+
 	CurrentLock *struct {
 		LockID     string
 		Expiration string
@@ -509,6 +517,13 @@ type IPAddress struct {
 	}
 }
 
+type EgressIPAddress struct {
+	ID      string
+	IP      string
+	Version int
+	Region  string
+}
+
 type VMSize struct {
 	Name        string
 	CPUCores    float32
@@ -680,6 +695,11 @@ type AllocateIPAddressInput struct {
 	Network        string `json:"network,omitempty"`
 }
 
+type AllocateEgressIPAddressInput struct {
+	AppID     string `json:"appId"`
+	MachineID string `json:"machineId"`
+}
+
 type ReleaseIPAddressInput struct {
 	AppID       *string `json:"appId"`
 	IPAddressID *string `json:"ipAddressId"`
@@ -856,6 +876,10 @@ type GqlMachine struct {
 	IPs struct {
 		Nodes []*MachineIP
 	}
+
+	EgressIpAddresses struct {
+		Nodes []*EgressIPAddress
+	}
 }
 
 type Logger interface {