feat: allow granting rds_replication

cdk-postgresql/lib/handler.ts

@@ -3,13 +3,16 @@ import { CloudFormationCustomResourceEvent } from "aws-lambda/trigger/cloudforma
 import { VError } from "verror";
 import { handler as dbHandler } from "./database.handler";
 import { handler as roleHandler } from "./role.handler";
+import { handler as roleMembershipHandler } from "./role-membership.handler";
 
 export const handler = async (event: CloudFormationCustomResourceEvent) => {
   switch (event.ResourceType) {
     case "Custom::Postgresql-Role":
       return roleHandler(event);
     case "Custom::Postgresql-Database":
       return dbHandler(event);
+    case "Custom::Postgresql-RoleMembership":
+      return roleMembershipHandler(event);
     default:
       throw new VError(`unexpected ResourceType: ${event.ResourceType}`);
   }

cdk-postgresql/lib/index.ts

@@ -1,3 +1,4 @@
 export * from "./database";
 export * from "./role";
+export * from "./role-membership";
 export * from "./provider";

cdk-postgresql/lib/postgres.ts

@@ -6,6 +6,12 @@ const isDatabaseError = (e: any): e is DatabaseError => {
   return typeof e.name === "string" && typeof e.length === "number";
 };
 
+/**
+ * Postgres error code for a statement naming an object that does not exist,
+ * such as a role that has already been dropped.
+ */
+const UNDEFINED_OBJECT_ERROR_CODE = "42704";
+
 export const createRole = async (props: {
   client: Client;
   name: string;
@@ -55,3 +61,47 @@ export const createDatabase = async (props: {
     `CREATE DATABASE ${escapeIdentifier(name)} WITH OWNER ${escapeIdentifier(owner)}`
   );
 };
+
+export const grantRoleMembership = async (props: {
+  client: Client;
+  role: string;
+  member: string;
+}) => {
+  const { client, role, member } = props;
+
+  await client.query(
+    `GRANT ${escapeIdentifier(role)} TO ${escapeIdentifier(member)}`
+  );
+};
+
+/**
+ * Revoking tolerates a role that no longer exists, because a membership whose
+ * role or member has already been dropped is in the wanted state. Postgres
+ * raises an error for it, which would otherwise leave the custom resource in
+ * DELETE_FAILED whenever the roles are dropped before the membership.
+ */
+export const revokeRoleMembership = async (props: {
+  client: Client;
+  role: string;
+  member: string;
+}) => {
+  const { client, role, member } = props;
+
+  try {
+    await client.query(
+      `REVOKE ${escapeIdentifier(role)} FROM ${escapeIdentifier(member)}`
+    );
+  } catch (thrown: unknown) {
+    if (!util.types.isNativeError(thrown)) {
+      throw thrown;
+    }
+    if (
+      !isDatabaseError(thrown) ||
+      thrown.code !== UNDEFINED_OBJECT_ERROR_CODE
+    ) {
+      throw new VError(thrown, "unexpected error while revoking role membership");
+    }
+
+    console.warn(thrown.message);
+  }
+};

cdk-postgresql/lib/role-membership.handler.ts

@@ -0,0 +1,115 @@
+import {
+  CloudFormationCustomResourceEvent,
+  CloudFormationCustomResourceCreateEvent,
+  CloudFormationCustomResourceUpdateEvent,
+  CloudFormationCustomResourceDeleteEvent,
+} from "aws-lambda/trigger/cloudformation-custom-resource";
+
+import { validateConnection, hashCode, getConnectedClient } from "./util";
+import { Connection } from "./lambda.types";
+import * as postgres from "./postgres";
+
+interface Props {
+  ServiceToken: string;
+  Connection: Connection;
+  Role: string;
+  Member: string;
+}
+
+export const handler = async (event: CloudFormationCustomResourceEvent) => {
+  switch (event.RequestType) {
+    case "Create":
+      return handleCreate(event);
+    case "Update":
+      return handleUpdate(event);
+    case "Delete":
+      return handleDelete(event);
+  }
+};
+
+const handleCreate = async (event: CloudFormationCustomResourceCreateEvent) => {
+  const props = event.ResourceProperties as Props;
+  validateProps(props);
+  await grantRoleMembership(props.Connection, props.Role, props.Member);
+  return {
+    PhysicalResourceId: generatePhysicalId(props),
+  };
+};
+
+/**
+ * Returning a physical id that encodes both the role and the member makes
+ * CloudFormation grant the new membership and then revoke the old one whenever
+ * either of them changes, which is safe because the two are never the same grant.
+ */
+const handleUpdate = async (event: CloudFormationCustomResourceUpdateEvent) => {
+  const props = event.ResourceProperties as Props;
+  validateProps(props);
+
+  const physicalResourceId = generatePhysicalId(props);
+  const isDifferentMembership = physicalResourceId != event.PhysicalResourceId;
+
+  if (isDifferentMembership) {
+    await grantRoleMembership(props.Connection, props.Role, props.Member);
+  }
+
+  return { PhysicalResourceId: physicalResourceId };
+};
+
+const handleDelete = async (event: CloudFormationCustomResourceDeleteEvent) => {
+  const props = event.ResourceProperties as Props;
+  validateProps(props);
+  await revokeRoleMembership(props.Connection, props.Role, props.Member);
+  return {};
+};
+
+const validateProps = (props: Props) => {
+  if (!("Connection" in props)) {
+    throw "Connection property is required";
+  }
+  validateConnection(props.Connection);
+
+  if (!("Role" in props)) {
+    throw "Role property is required";
+  }
+  if (!("Member" in props)) {
+    throw "Member property is required";
+  }
+};
+
+const generatePhysicalId = (props: Props): string => {
+  const { Host, Port } = props.Connection;
+  const suffix = Math.abs(
+    hashCode(JSON.stringify([Host, Port, props.Role, props.Member]))
+  );
+  return `role-membership-${suffix}`;
+};
+
+export const grantRoleMembership = async (
+  connection: Connection,
+  role: string,
+  member: string
+) => {
+  console.log(`Granting ${role} to ${member}`);
+  const client = await getConnectedClient(connection);
+
+  try {
+    await postgres.grantRoleMembership({ client, role, member });
+  } finally {
+    await client.end();
+  }
+};
+
+export const revokeRoleMembership = async (
+  connection: Connection,
+  role: string,
+  member: string
+) => {
+  console.log(`Revoking ${role} from ${member}`);
+  const client = await getConnectedClient(connection);
+
+  try {
+    await postgres.revokeRoleMembership({ client, role, member });
+  } finally {
+    await client.end();
+  }
+};

cdk-postgresql/lib/role-membership.ts

@@ -0,0 +1,80 @@
+import { Construct } from "constructs";
+import * as cdk from "aws-cdk-lib";
+import { RemovalPolicy } from "aws-cdk-lib";
+import { Provider } from "./provider";
+
+/**
+ * Predefined roles, listed so they auto-complete in vscode.
+ */
+export type PredefinedRoleName =
+  | "rds_replication"
+  | "rds_superuser"
+  | "rds_iam"
+  | "rds_password"
+  | "pg_monitor"
+  | "pg_read_all_data"
+  | "pg_write_all_data"
+  | "pg_read_all_settings"
+  | "pg_read_all_stats"
+  | "pg_stat_scan_tables"
+  | "pg_signal_backend";
+
+export interface RoleMembershipProps {
+  /**
+   * Provider required to connect to the Postgresql server
+   */
+  provider: Provider;
+
+  /**
+   * The role whose membership is granted, as in `GRANT <role> TO <member>`. It
+   * has to exist already, whether it is predefined or was created elsewhere in
+   * the stack.
+   */
+  role: PredefinedRoleName | (string & {});
+
+  /**
+   * The role receiving the membership, as in `GRANT <role> TO <member>`. It
+   * has to exist already, like the role it is made a member of.
+   */
+  member: string;
+
+  /**
+   * Policy to apply when the membership is removed from this stack.
+   *
+   * @default - The membership will be revoked.
+   */
+  removalPolicy?: RemovalPolicy;
+}
+
+/**
+ * Membership of one Postgresql role in another, which is how a role receives
+ * privileges it cannot be granted directly. On RDS and Aurora the master role
+ * cannot grant the replication attribute, so a role that has to read the
+ * write-ahead log is made a member of `rds_replication` instead.
+ *
+ * Both roles are named by plain strings, which creates no CloudFormation
+ * dependency. When either of them is created by a `Role` construct in the same
+ * stack, call `addDependency` on this construct so that the grant does not run
+ * first and fail with `role "..." does not exist`.
+ */
+export class RoleMembership extends Construct {
+  constructor(scope: Construct, id: string, props: RoleMembershipProps) {
+    super(scope, id);
+
+    const { provider, role, member, removalPolicy } = props;
+
+    const cr = new cdk.CustomResource(this, "CustomResource", {
+      serviceToken: provider.serviceToken,
+      resourceType: "Custom::Postgresql-RoleMembership",
+      properties: {
+        connection: provider.buildConnectionProperty(),
+        role,
+        member,
+      },
+      pascalCaseProperties: true,
+    });
+
+    cr.applyRemovalPolicy(removalPolicy || cdk.RemovalPolicy.DESTROY);
+    cr.node.addDependency(provider);
+  }
+}

cdk-postgresql/test/constructs.test.ts

@@ -3,7 +3,7 @@ import { Template } from "aws-cdk-lib/assertions";
 import * as secretsmanager from "aws-cdk-lib/aws-secretsmanager";
 import * as cdk from "aws-cdk-lib";
 import { Construct } from "constructs";
-import { Database, Role, Provider } from "../lib";
+import { Database, Role, Provider, RoleMembership } from "../lib";
 
 class TestStack extends cdk.Stack {
   constructor(scope: Construct, id: string, props?: cdk.StackProps) {
@@ -176,3 +176,83 @@ describe("role", () => {
     template.resourceCountIs("AWS::Lambda::Function", 3);
   });
 });
+
+describe(RoleMembership, () => {
+  const buildStackWithProvider = () => {
+    const stack = new TestStack(new cdk.App(), "Stack");
+    const connectionPassword = new secretsmanager.Secret(stack, "ConnectionPassword");
+    const provider = new Provider(stack, "provider", {
+      host: "somedb.com",
+      username: "theusername",
+      password: connectionPassword,
+    });
+    return { stack, provider, connectionPassword };
+  };
+
+  test("creates a custom resource carrying the connection, role and member", () => {
+    // Arrange
+    const { stack, provider, connectionPassword } = buildStackWithProvider();
+
+    // Act
+    new RoleMembership(stack, "RoleMembership", {
+      provider,
+      role: "rds_replication",
+      member: "myrole",
+    });
+
+    // Assert
+    const template = Template.fromStack(stack);
+    template.resourceCountIs("Custom::Postgresql-RoleMembership", 1);
+    template.hasResourceProperties("Custom::Postgresql-RoleMembership", {
+      Connection: {
+        Host: "somedb.com",
+        Port: 5432,
+        Database: "postgres",
+        Username: "theusername",
+        PasswordArn: {
+          Ref: getLogicalId(connectionPassword),
+        },
+        SSLMode: "require",
+      },
+      Role: "rds_replication",
+      Member: "myrole",
+    });
+  });
+
+  test("revokes the membership when it is removed from the stack", () => {
+    // Arrange
+    const { stack, provider } = buildStackWithProvider();
+
+    // Act
+    new RoleMembership(stack, "RoleMembership", {
+      provider,
+      role: "rds_replication",
+      member: "myrole",
+    });
+
+    // Assert
+    const template = Template.fromStack(stack);
+    template.hasResource("Custom::Postgresql-RoleMembership", {
+      DeletionPolicy: "Delete",
+    });
+  });
+
+  test("keeps the membership when the removal policy is retain", () => {
+    // Arrange
+    const { stack, provider } = buildStackWithProvider();
+
+    // Act
+    new RoleMembership(stack, "RoleMembership", {
+      provider,
+      role: "rds_replication",
+      member: "myrole",
+      removalPolicy: cdk.RemovalPolicy.RETAIN,
+    });
+
+    // Assert
+    const template = Template.fromStack(stack);
+    template.hasResource("Custom::Postgresql-RoleMembership", {
+      DeletionPolicy: "Retain",
+    });
+  });
+});

cdk-postgresql/test/helpers.ts

@@ -68,3 +68,16 @@ export const getDbOwner = async (
   }
   return dbRow.owner;
 };
+
+export const isMemberOf = async (props: {
+  client: Client;
+  member: string;
+  role: string;
+}): Promise<boolean> => {
+  const { client, member, role } = props;
+  const { rows } = await client.query(
+    "SELECT pg_has_role($1, $2, 'member') AS is_member",
+    [member, role]
+  );
+  return rows[0].is_member;
+};