From fc5abbc5894cf67716756298da082c28cafa4288 Mon Sep 17 00:00:00 2001
From: starryskymeow <fzc_study@163.com>
Date: Sat, 5 Jul 2025 20:12:17 +0800
Subject: [PATCH] add key to support wsrx

---
 internal/api/handlers.go  |  4 ++--
 internal/k8s/workspace.go | 36 ++++++++++++++++++++++++++++--------
 2 files changed, 30 insertions(+), 10 deletions(-)

diff --git a/internal/api/handlers.go b/internal/api/handlers.go
index dfeb696..55eb608 100644
--- a/internal/api/handlers.go
+++ b/internal/api/handlers.go
@@ -56,13 +56,13 @@ func (h *Handler) createWorkspace(c *gin.Context) {
 		Namespace:      Namespace,
 	}
 
-	err := k8s.CreateWorkspace(k8sReq)
+	keyword, err := k8s.CreateWorkspace(k8sReq)
 	if err != nil {
 		c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
 		return
 	}
 
-	c.JSON(http.StatusAccepted, gin.H{"status": "creating", "workspaceId": req.WorkspaceID})
+	c.JSON(http.StatusAccepted, gin.H{"status": "creating", "workspaceId": req.WorkspaceID, "keyword": keyword})
 }
 
 func (h *Handler) deleteWorkspace(c *gin.Context) {
diff --git a/internal/k8s/workspace.go b/internal/k8s/workspace.go
index 7fa8f78..04ed9bf 100644
--- a/internal/k8s/workspace.go
+++ b/internal/k8s/workspace.go
@@ -2,6 +2,8 @@ package k8s
 
 import (
 	"context"
+	"crypto/rand"
+	"encoding/hex"
 	"fmt"
 	"time"
 
@@ -38,12 +40,27 @@ type ResourceLimits struct {
 	Storage string `json:"storage"`
 }
 
-func CreateWorkspace(req *WorkspaceRequest) error {
+func generateRandomString(length int) (string, error) {
+	bytes := make([]byte, length)
+	if _, err := rand.Read(bytes); err != nil {
+		return "", err
+	}
+	return hex.EncodeToString(bytes), nil
+}
+
+func CreateWorkspace(req *WorkspaceRequest) (string, error) {
+	// Generate a random keyword for wsrx
+	wsrxKeyword, err := generateRandomString(8)
+	if err != nil {
+		return "", fmt.Errorf("failed to generate random string for wsrx keyword: %w", err)
+	}
+
 	// Define StatefulSet
 	sts := &appsv1.StatefulSet{
 		ObjectMeta: metav1.ObjectMeta{
 			Name:      req.WorkspaceID,
 			Namespace: req.Namespace,
+			Labels:    map[string]string{"educode/autodelete": "true"},
 			Annotations: map[string]string{
 				"educode/expires-at": time.Now().Add(1 * time.Hour).Format(time.RFC3339),
 			},
@@ -93,16 +110,18 @@ func CreateWorkspace(req *WorkspaceRequest) error {
 	}
 
 	// Create resources
-	sts, err := req.Clientset.AppsV1().StatefulSets(req.Namespace).Create(context.TODO(), sts, metav1.CreateOptions{})
+	sts, err = req.Clientset.AppsV1().StatefulSets(req.Namespace).Create(context.TODO(), sts, metav1.CreateOptions{})
 	if err != nil {
-		return fmt.Errorf("failed to create statefulset: %w", err)
+		return "", fmt.Errorf("failed to create statefulset: %w", err)
 	}
 
 	// Define Service
 	svc := &corev1.Service{
 		ObjectMeta: metav1.ObjectMeta{
-			Name:      req.WorkspaceID,
-			Namespace: req.Namespace,
+			Name:        req.WorkspaceID,
+			Namespace:   req.Namespace,
+			Labels:      map[string]string{"wsrx/enabled": "true"},
+			Annotations: map[string]string{"wsrx/keyword": wsrxKeyword},
 			OwnerReferences: []metav1.OwnerReference{
 				{
 					APIVersion: "apps/v1",
@@ -119,7 +138,8 @@ func CreateWorkspace(req *WorkspaceRequest) error {
 				{Name: "code-server", Port: 8080, TargetPort: intstr.FromInt(8080)},
 				{Name: "ssh", Port: 22, TargetPort: intstr.FromInt(22)},
 			},
-			Type: corev1.ServiceTypeClusterIP,
+			Type:      corev1.ServiceTypeClusterIP,
+			ClusterIP: "None",
 		},
 	}
 
@@ -130,10 +150,10 @@ func CreateWorkspace(req *WorkspaceRequest) error {
 		req.Clientset.AppsV1().StatefulSets(req.Namespace).Delete(context.TODO(), req.WorkspaceID, metav1.DeleteOptions{
 			PropagationPolicy: &deletePolicy,
 		})
-		return fmt.Errorf("failed to create service: %w", err)
+		return "", fmt.Errorf("failed to create service: %w", err)
 	}
 
-	return nil
+	return wsrxKeyword, nil
 }
 
 func DeleteWorkspace(clientset *kubernetes.Clientset, namespace, workspaceID string) error {