dev day #99

internal/controllers/shortlink_controller.go

@@ -136,6 +136,23 @@ func codeFromHash(s string, n int) string {
 	return string(out)
 }
 
+func sanitizeCode(in string) string {
+    s := strings.TrimSpace(in)
+    if s == "" { return "" }
+    // filter allowed runes
+    rb := make([]rune, 0, len(s))
+    for _, ch := range s {
+        if (ch >= 'a' && ch <= 'z') || (ch >= 'A' && ch <= 'Z') || (ch >= '0' && ch <= '9') || ch == '-' || ch == '_' {
+            rb = append(rb, ch)
+        }
+    }
+    if len(rb) == 0 { return "" }
+    if len(rb) > 16 {
+        rb = rb[:16]
+    }
+    return string(rb)
+}
+
 func getScheme(c *gin.Context) string {
 	if p := c.GetHeader("X-Forwarded-Proto"); p != "" {
 		return p
@@ -256,18 +273,18 @@ func (s *ShortLinkController) CreateShortLink(c *gin.Context) {
 		c.JSON(http.StatusBadRequest, gin.H{"error": "invalid target_url"})
 		return
 	}
-	code := strings.TrimSpace(body.Code)
-	if code == "" {
-		for i := 0; i < 5; i++ {
-			cnd, _ := randCode(7)
-			var cnt int64
-			s.DB.Model(&models.ShortLink{}).Where("code = ?", cnd).Count(&cnt)
-			if cnt == 0 {
-				code = cnd
-				break
-			}
-		}
-	}
+	code := sanitizeCode(strings.TrimSpace(body.Code))
+    if code == "" {
+        for i := 0; i < 5; i++ {
+            cnd, _ := randCode(7)
+            var cnt int64
+            s.DB.Model(&models.ShortLink{}).Where("code = ?", cnd).Count(&cnt)
+            if cnt == 0 {
+                code = cnd
+                break
+            }
+        }
+    }
 	if code == "" {
 		c.JSON(http.StatusInternalServerError, gin.H{"error": "cannot generate code"})
 		return