Get AI to fix vunerabilities in site

backend/handlers/handle_auth.go

@@ -1,7 +1,6 @@
 package handlers
 
 import (
-	"fmt"
 	"net/http"
 
 	"adam-french.co.uk/backend/models"
@@ -68,10 +67,9 @@ func (store *Store) RefreshToken(ctx *gin.Context) {
 	claims, err := store.Auth.VerifyJWT(refreshToken)
 	if err != nil {
 		ctx.JSON(http.StatusUnauthorized, err.Error())
+		return
 	}
 
-	fmt.Printf("claims: %v\n", claims)
-
 	userIDF, ok := (*claims)["id"].(float64)
 	if !ok {
 		ctx.JSON(http.StatusInternalServerError, gin.H{"error": "invalid token claims"})
@@ -93,6 +91,7 @@ func (store *Store) RefreshToken(ctx *gin.Context) {
 		return
 	}
 
+	ctx.SetSameSite(http.SameSiteLaxMode)
 	ctx.SetCookie(
 		"access_token",
 		tokens.AccessToken,
@@ -122,12 +121,12 @@ func (store *Store) Login(ctx *gin.Context) {
 
 	user := models.User{}
 	if err := store.DB.Where("username = ?", input.Username).First(&user).Error; err != nil {
-		ctx.JSON(http.StatusNotFound, err.Error())
+		ctx.JSON(http.StatusUnauthorized, gin.H{"error": "invalid credentials"})
 		return
 	}
 
 	if err := bcrypt.CompareHashAndPassword(user.Password, []byte(input.Password)); err != nil {
-		ctx.JSON(http.StatusUnauthorized, err.Error())
+		ctx.JSON(http.StatusUnauthorized, gin.H{"error": "invalid credentials"})
 		return
 	}
 
@@ -137,6 +136,7 @@ func (store *Store) Login(ctx *gin.Context) {
 		return
 	}
 
+	ctx.SetSameSite(http.SameSiteLaxMode)
 	ctx.SetCookie(
 		"access_token",
 		tokens.AccessToken,
@@ -164,6 +164,7 @@ func (store *Store) Logout(ctx *gin.Context) {
 }
 
 func removeCookies(ctx *gin.Context) {
+	ctx.SetSameSite(http.SameSiteLaxMode)
 	ctx.SetCookie(
 		"access_token",
 		"",

backend/handlers/handle_message_upload.go

@@ -17,6 +17,12 @@ var allowedExtensions = map[string]bool{
 	".pdf": true, ".txt": true,
 }
 
+var extensionToMIMEPrefix = map[string]string{
+	".jpg": "image/", ".jpeg": "image/", ".png": "image/png", ".gif": "image/gif", ".webp": "image/webp",
+	".mp4": "video/", ".webm": "video/webm", ".mp3": "audio/", ".ogg": "audio/",
+	".pdf": "application/pdf", ".txt": "text/",
+}
+
 func (store *Store) UploadMessageFile(ctx *gin.Context) {
 	file, err := ctx.FormFile("file")
 	if err != nil {
@@ -36,6 +42,23 @@ func (store *Store) UploadMessageFile(ctx *gin.Context) {
 		return
 	}
 
+	// Validate actual content type matches extension
+	f, err := file.Open()
+	if err != nil {
+		ctx.JSON(http.StatusInternalServerError, gin.H{"error": "failed to read file"})
+		return
+	}
+	buf := make([]byte, 512)
+	n, _ := f.Read(buf)
+	f.Close()
+	detectedType := http.DetectContentType(buf[:n])
+
+	expectedPrefix, ok := extensionToMIMEPrefix[ext]
+	if ok && !strings.HasPrefix(detectedType, expectedPrefix) {
+		ctx.JSON(http.StatusBadRequest, gin.H{"error": "file content does not match extension"})
+		return
+	}
+
 	b := make([]byte, 16)
 	if _, err := rand.Read(b); err != nil {
 		ctx.JSON(http.StatusInternalServerError, gin.H{"error": "failed to generate filename"})

backend/handlers/handle_user.go

@@ -15,6 +15,21 @@ type UserCredentials struct {
 }
 
 func (store *Store) CreateUser(ctx *gin.Context) {
+	claimsVal, ok := ctx.Get("userClaims")
+	if !ok {
+		ctx.JSON(http.StatusUnauthorized, gin.H{"error": "user claims could not be found"})
+		return
+	}
+	claims, ok := claimsVal.(*jwt.MapClaims)
+	if !ok {
+		ctx.JSON(http.StatusInternalServerError, gin.H{"error": "invalid claims"})
+		return
+	}
+	if !(*claims)["admin"].(bool) {
+		ctx.JSON(http.StatusForbidden, gin.H{"error": "admin access required"})
+		return
+	}
+
 	var input UserCredentials
 	if err := ctx.ShouldBindBodyWithJSON(&input); err != nil {
 		ctx.JSON(http.StatusBadRequest, err.Error())
@@ -31,32 +46,9 @@ func (store *Store) CreateUser(ctx *gin.Context) {
 	tx := store.DB.Create(&user)
 	if tx.Error != nil {
 		ctx.JSON(http.StatusInternalServerError, tx.Error.Error())
-	}
-
-	// Generate JWT token
-	tokens, err := store.Auth.GenerateJWT(&user)
-	if err != nil {
-		ctx.JSON(http.StatusInternalServerError, err.Error())
 		return
 	}
 
-	ctx.SetCookie(
-		"access_token",
-		tokens.AccessToken,
-		int(store.Auth.Config.AccessTokenLifetime.Seconds()),
-		store.Auth.Config.Endpoint,
-		store.Auth.Config.Domain,
-		true, true,
-	)
-	ctx.SetCookie(
-		"refresh_token",
-		tokens.RefreshToken,
-		int(store.Auth.Config.RefreshTokenLifetime.Seconds()),
-		store.Auth.Config.Endpoint,
-		store.Auth.Config.Domain,
-		true, true,
-	)
-
 	ctx.JSON(http.StatusOK, user)
 }
 
@@ -141,6 +133,7 @@ func (store *Store) DeleteUser(ctx *gin.Context) {
 		return
 	}
 
+	ctx.SetSameSite(http.SameSiteLaxMode)
 	ctx.SetCookie(
 		"access_token",
 		"",

backend/main.go

@@ -41,7 +41,8 @@ func main() {
 	if os.Getenv("SEED_DB") == "true" {
 		services.SeedDatabase(db)
 	}
-	services.InitWebSocket(db)
+	domainName := os.Getenv("DOMAIN")
+	services.InitWebSocket(db, domainName)
 
 	// SPOTIFY
 	spotifyAuthState := os.Getenv("SPOTIFY_AUTH_STATE")
@@ -57,7 +58,6 @@ func main() {
 	claudeClient := services.InitClaude(&claudeConfig)
 
 	authSecret := os.Getenv("BACKEND_SECRET")
-	domainName := os.Getenv("DOMAIN")
 	backendEndpoint := os.Getenv("BACKEND_ENDPOINT")
 	accessTokenLifetime := 24 * time.Hour
 	refreshTokenLifetime := 365 * 24 * time.Hour
@@ -96,7 +96,7 @@ func main() {
 	protected.PUT("/user/:id", store.UpdateUser)
 	protected.DELETE("/user/:id", store.DeleteUser)
 	r.GET("/user", store.GetUsers)
-	r.POST("/user", store.CreateUser)
+	protected.POST("/user", store.CreateUser)
 
 	// AUTH
 	r.POST("/auth/login", store.Login)
@@ -112,7 +112,7 @@ func main() {
 
 	// MESSAGES
 	r.GET("/ws", store.ConnectWebSocket)
-	r.POST("/messages/upload", store.UploadMessageFile)
+	protected.POST("/messages/upload", store.UploadMessageFile)
 
 	// NOTES
 	r.GET("/notes/*path", store.GetNoteFile)

backend/services/websocket.go

@@ -2,7 +2,9 @@ package services
 
 import (
 	"net/http"
+	"strings"
 	"sync"
+	"time"
 
 	"adam-french.co.uk/backend/models"
 	"gorm.io/gorm"
@@ -12,11 +14,19 @@ import (
 
 const maxMessages = 50
 
+var allowedDomain string
+
 var Upgrader = websocket.Upgrader{
 	ReadBufferSize:  1024,
 	WriteBufferSize: 1024,
 	CheckOrigin: func(r *http.Request) bool {
-		return true
+		origin := r.Header.Get("Origin")
+		if origin == "" {
+			return false
+		}
+		origin = strings.TrimPrefix(origin, "https://")
+		origin = strings.TrimPrefix(origin, "http://")
+		return origin == allowedDomain || origin == "www."+allowedDomain
 	},
 }
 
@@ -27,8 +37,14 @@ var (
 	nextAuthorID uint
 )
 
-func InitWebSocket(database *gorm.DB) {
+const (
+	rateLimitWindow   = time.Second
+	rateLimitMaxMsgs  = 10
+)
+
+func InitWebSocket(database *gorm.DB, domain string) {
 	wsDB = database
+	allowedDomain = domain
 }
 
 func HandleWebSocket(conn *websocket.Conn) {
@@ -50,12 +66,25 @@ func HandleWebSocket(conn *websocket.Conn) {
 	}
 	mu.Unlock()
 
+	msgCount := 0
+	windowStart := time.Now()
+
 	for {
 		var incoming models.Message
 		if err := conn.ReadJSON(&incoming); err != nil {
 			break
 		}
 
+		now := time.Now()
+		if now.Sub(windowStart) > rateLimitWindow {
+			msgCount = 0
+			windowStart = now
+		}
+		msgCount++
+		if msgCount > rateLimitMaxMsgs {
+			continue
+		}
+
 		incoming.AuthorID = authorID
 
 		mu.Lock()