Merge pull request #1313 from CXKswain/main

Updated ollama.go to fix a couple of potential DoS issues

pkgs/fabric/version.nix

@@ -1 +1 @@
-"1.4.139"
+"1.4.139"

restapi/chat.go

@@ -55,6 +55,7 @@ func (h *ChatHandler) HandleChat(c *gin.Context) {
 
 	if err := c.BindJSON(&request); err != nil {
 		log.Printf("Error binding JSON: %v", err)
+		c.Writer.Header().Set("Strict-Transport-Security", "max-age=63072000; includeSubDomains")
 		c.JSON(http.StatusBadRequest, gin.H{"error": fmt.Sprintf("Invalid request format: %v", err)})
 		return
 	}

restapi/configuration.go

@@ -90,7 +90,7 @@ func (h *ConfigHandler) UpdateConfig(c *gin.Context) {
 		LMStudioURL      string `json:"lm_studio_base_url"`
 	}
 
-	if err := c.BindJSON(&config); err != nil {
+	if err := c.ShouldBindJSON(&config); err != nil {
 		c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
 		return
 	}

restapi/ollama.go

@@ -5,18 +5,32 @@ import (
 	"context"
 	"encoding/json"
 	"fmt"
-	"github.com/danielmiessler/fabric/core"
-	"github.com/gin-gonic/gin"
 	"io"
 	"log"
 	"net/http"
 	"strings"
 	"time"
+
+	"github.com/danielmiessler/fabric/core"
+	"github.com/gin-gonic/gin"
 )
 
 type OllamaModel struct {
 	Models []Model `json:"models"`
 }
+
+func isValidOllamaRequestBody(body OllamaRequestBody) bool {
+	if body.Model == "" || len(body.Messages) == 0 {
+		return false
+	}
+	for _, msg := range body.Messages {
+		if msg.Content == "" || msg.Role == "" {
+			return false
+		}
+	}
+	return true
+}
+
 type Model struct {
 	Details    ModelDetails `json:"details"`
 	Digest     string       `json:"digest"`
@@ -142,12 +156,16 @@ func (f APIConvert) ollamaTags(c *gin.Context) {
 		})
 	}
 
-	c.JSON(200, response)
-
-}
-
-func (f APIConvert) ollamaChat(c *gin.Context) {
-	body, err := io.ReadAll(c.Request.Body)
+	if !isValidOllamaRequestBody(prompt) {
+		log.Printf("Invalid request body: %v", prompt)
+		c.JSON(http.StatusBadRequest, gin.H{"error": "invalid request body"})
+		return
+	}
+	if !isValidOllamaRequestBody(prompt) {
+		log.Printf("Invalid request body: %v", prompt)
+		c.JSON(http.StatusBadRequest, gin.H{"error": "invalid request body"})
+		return
+	}
 	if err != nil {
 		log.Printf("Error reading body: %v", err)
 		c.JSON(http.StatusInternalServerError, gin.H{"error": "testing endpoint"})

version.go

@@ -1,3 +1,3 @@
 package main
 
-var version = "v1.4.139"
+var version = "v1.4.139"