From 2724336cad1509c212cc615394d2e4aa06234cbd Mon Sep 17 00:00:00 2001
From: Lorenzo Migliorero <lorenzo.migliorero@wetransfer.com>
Date: Thu, 25 Jul 2024 23:32:52 +0200
Subject: [PATCH] feat: ssh key type and whitespaces

---
 .../settings/ssh-keys/add-ssh-key.tsx         | 54 +++++++++++--------
 .../settings/ssh-keys/update-ssh-key.tsx      | 21 +++++++-
 server/api/routers/ssh-key.ts                 |  9 ++--
 server/db/schema/ssh-key.ts                   |  4 +-
 server/db/validations/index.ts                | 35 ++++++++----
 server/utils/filesystem/ssh.ts                | 20 +++----
 6 files changed, 94 insertions(+), 49 deletions(-)

diff --git a/components/dashboard/settings/ssh-keys/add-ssh-key.tsx b/components/dashboard/settings/ssh-keys/add-ssh-key.tsx
index cc6e03ab..903a8932 100644
--- a/components/dashboard/settings/ssh-keys/add-ssh-key.tsx
+++ b/components/dashboard/settings/ssh-keys/add-ssh-key.tsx
@@ -19,7 +19,7 @@ import {
 } from "@/components/ui/form";
 import { Input } from "@/components/ui/input";
 import { Textarea } from "@/components/ui/textarea";
-import { sshKeyCreate } from "@/server/db/validations";
+import { sshKeyCreate, type sshKeyType } from "@/server/db/validations";
 import { api } from "@/utils/api";
 import { zodResolver } from "@hookform/resolvers/zod";
 import { type ReactNode, useState } from "react";
@@ -65,6 +65,18 @@ export const AddSSHKey = ({ children }: Props) => {
 			});
 	};
 
+	const onGenerateSSHKey = (type: z.infer<typeof sshKeyType>) =>
+		generateMutation
+			.mutateAsync(type)
+			.then(async (data) => {
+				toast.success("SSH Key Generated");
+				form.setValue("privateKey", data.privateKey);
+				form.setValue("publicKey", data.publicKey);
+			})
+			.catch(() => {
+				toast.error("Error to generate the SSH Key");
+			});
+
 	return (
 		<Dialog open={isOpen} onOpenChange={setIsOpen}>
 			<DialogTrigger className="" asChild>
@@ -78,26 +90,26 @@ export const AddSSHKey = ({ children }: Props) => {
 							In this section you can add one of your keys or generate a new
 							one.
 						</div>
-						<Button
-							variant={"secondary"}
-							isLoading={generateMutation.isLoading}
-							className="max-sm:w-full"
-							onClick={async () => {
-								await generateMutation
-									.mutateAsync()
-									.then(async (data) => {
-										toast.success("SSH Key Generated");
-										form.setValue("privateKey", data.privateKey);
-										form.setValue("publicKey", data.publicKey);
-									})
-									.catch(() => {
-										toast.error("Error to generate the SSH Key");
-									});
-							}}
-							type="button"
-						>
-							Generate SSH Key
-						</Button>
+						<div className="flex gap-4">
+							<Button
+								variant={"secondary"}
+								disabled={generateMutation.isLoading}
+								className="max-sm:w-full"
+								onClick={() => onGenerateSSHKey("rsa")}
+								type="button"
+							>
+								Generate RSA SSH Key
+							</Button>
+							<Button
+								variant={"secondary"}
+								disabled={generateMutation.isLoading}
+								className="max-sm:w-full"
+								onClick={() => onGenerateSSHKey("ed25519")}
+								type="button"
+							>
+								Generate ED25519 SSH Key
+							</Button>
+						</div>
 					</DialogDescription>
 				</DialogHeader>
 				{isError && <AlertBlock type="error">{error?.message}</AlertBlock>}
diff --git a/components/dashboard/settings/ssh-keys/update-ssh-key.tsx b/components/dashboard/settings/ssh-keys/update-ssh-key.tsx
index 95f2ec1d..e763ee92 100644
--- a/components/dashboard/settings/ssh-keys/update-ssh-key.tsx
+++ b/components/dashboard/settings/ssh-keys/update-ssh-key.tsx
@@ -23,6 +23,8 @@ import { Textarea } from "@/components/ui/textarea";
 import { sshKeyUpdate } from "@/server/db/validations";
 import { api } from "@/utils/api";
 import { zodResolver } from "@hookform/resolvers/zod";
+import copy from "copy-to-clipboard";
+import { CopyIcon } from "lucide-react";
 import { type ReactNode, useEffect, useState } from "react";
 import { useForm } from "react-hook-form";
 import { toast } from "sonner";
@@ -131,7 +133,24 @@ export const UpdateSSHKey = ({ children, sshKeyId = "" }: Props) => {
 						<FormItem>
 							<FormLabel>Public Key</FormLabel>
 							<FormControl>
-								<Textarea rows={7} readOnly disabled value={data?.publicKey} />
+								<div className="relative">
+									<Textarea
+										rows={7}
+										readOnly
+										disabled
+										value={data?.publicKey}
+									/>
+									<button
+										type="button"
+										className="absolute right-2 top-2"
+										onClick={() => {
+											copy(data?.publicKey || "Generate a SSH Key");
+											toast.success("SSH Copied to clipboard");
+										}}
+									>
+										<CopyIcon className="size-4" />
+									</button>
+								</div>
 							</FormControl>
 							<FormMessage />
 						</FormItem>
diff --git a/server/api/routers/ssh-key.ts b/server/api/routers/ssh-key.ts
index d51e31c1..77054fdd 100644
--- a/server/api/routers/ssh-key.ts
+++ b/server/api/routers/ssh-key.ts
@@ -7,6 +7,7 @@ import { db } from "@/server/db";
 import {
 	apiCreateSshKey,
 	apiFindOneSshKey,
+	apiGenerateSSHKey,
 	apiRemoveSshKey,
 	apiUpdateSshKey,
 } from "@/server/db/schema";
@@ -50,9 +51,11 @@ export const sshRouter = createTRPCRouter({
 	all: adminProcedure.query(async () => {
 		return await db.query.sshKeys.findMany({});
 	}),
-	generate: protectedProcedure.mutation(async () => {
-		return await generateSSHKey();
-	}),
+	generate: protectedProcedure
+		.input(apiGenerateSSHKey)
+		.mutation(async ({ input }) => {
+			return await generateSSHKey(input);
+		}),
 	update: adminProcedure.input(apiUpdateSshKey).mutation(async ({ input }) => {
 		try {
 			return await updateSSHKeyById(input);
diff --git a/server/db/schema/ssh-key.ts b/server/db/schema/ssh-key.ts
index 0ef4554e..2251c0ae 100644
--- a/server/db/schema/ssh-key.ts
+++ b/server/db/schema/ssh-key.ts
@@ -1,6 +1,6 @@
 import { applications } from "@/server/db/schema/application";
 import { compose } from "@/server/db/schema/compose";
-import { sshKeyCreate } from "@/server/db/validations";
+import { sshKeyCreate, sshKeyType } from "@/server/db/validations";
 import { relations } from "drizzle-orm";
 import { pgTable, text, time } from "drizzle-orm/pg-core";
 import { createInsertSchema } from "drizzle-zod";
@@ -45,6 +45,8 @@ export const apiFindOneSshKey = createSchema
 	})
 	.required();
 
+export const apiGenerateSSHKey = sshKeyType;
+
 export const apiRemoveSshKey = createSchema
 	.pick({
 		sshKeyId: true,
diff --git a/server/db/validations/index.ts b/server/db/validations/index.ts
index 930ce5e7..568e2f46 100644
--- a/server/db/validations/index.ts
+++ b/server/db/validations/index.ts
@@ -3,20 +3,33 @@ import { z } from "zod";
 export const sshKeyCreate = z.object({
 	name: z.string().min(1),
 	description: z.string().optional(),
-	publicKey: z.string().regex(/^ssh-rsa\s+([A-Za-z0-9+/=]+)\s*(.*)?$/, {
-		message: "Invalid format",
-	}),
-	privateKey: z
-		.string()
-		.regex(
-			/^-----BEGIN RSA PRIVATE KEY-----\n([A-Za-z0-9+/=\n]+)-----END RSA PRIVATE KEY-----$/,
-			{
-				message: "Invalid format",
-			},
-		),
+	publicKey: z.string().refine(
+		(key) => {
+			const rsaPubPattern = /^ssh-rsa\s+([A-Za-z0-9+/=]+)\s*(.*)?\s*$/;
+			const ed25519PubPattern = /^ssh-ed25519\s+([A-Za-z0-9+/=]+)\s*(.*)?\s*$/;
+			return rsaPubPattern.test(key) || ed25519PubPattern.test(key);
+		},
+		{
+			message: "Invalid public key format",
+		},
+	),
+	privateKey: z.string().refine(
+		(key) => {
+			const rsaPrivPattern =
+				/^-----BEGIN RSA PRIVATE KEY-----\n([A-Za-z0-9+/=\n]+)-----END RSA PRIVATE KEY-----\s*$/;
+			const ed25519PrivPattern =
+				/^-----BEGIN OPENSSH PRIVATE KEY-----\n([A-Za-z0-9+/=\n]+)-----END OPENSSH PRIVATE KEY-----\s*$/;
+			return rsaPrivPattern.test(key) || ed25519PrivPattern.test(key);
+		},
+		{
+			message: "Invalid private key format",
+		},
+	),
 });
 
 export const sshKeyUpdate = sshKeyCreate.pick({
 	name: true,
 	description: true,
 });
+
+export const sshKeyType = z.enum(["rsa", "ed25519"]).optional();
diff --git a/server/utils/filesystem/ssh.ts b/server/utils/filesystem/ssh.ts
index 2d3b63d0..fff0eba0 100644
--- a/server/utils/filesystem/ssh.ts
+++ b/server/utils/filesystem/ssh.ts
@@ -10,16 +10,12 @@ const readSSHKey = async (id: string) => {
 		}
 
 		return {
-			privateKey: fs
-				.readFileSync(path.join(SSH_PATH, `${id}_rsa`), {
-					encoding: "utf-8",
-				})
-				.trim(),
-			publicKey: fs
-				.readFileSync(path.join(SSH_PATH, `${id}_rsa.pub`), {
-					encoding: "utf-8",
-				})
-				.trim(),
+			privateKey: fs.readFileSync(path.join(SSH_PATH, `${id}_rsa`), {
+				encoding: "utf-8",
+			}),
+			publicKey: fs.readFileSync(path.join(SSH_PATH, `${id}_rsa.pub`), {
+				encoding: "utf-8",
+			}),
 		};
 	} catch (error) {
 		throw error;
@@ -47,7 +43,7 @@ export const saveSSHKey = async (
 	publicKeyStream.end();
 };
 
-export const generateSSHKey = async () => {
+export const generateSSHKey = async (type: "rsa" | "ed25519" = "rsa") => {
 	const applicationDirectory = SSH_PATH;
 
 	if (!fs.existsSync(applicationDirectory)) {
@@ -66,7 +62,7 @@ export const generateSSHKey = async () => {
 
 	const args = [
 		"-t",
-		"rsa",
+		type,
 		"-b",
 		"4096",
 		"-C",