mirror of
https://github.com/danielmiessler/fabric.git
synced 2026-09-10 07:36:44 -04:00
fix: validate Fabric chat URL host and tidy Ollama responses
## CHANGES - Set NDJSON content type before checking upstream errors - Reject parsed URLs that omit a hostname - Remove hardcoded eval count placeholders from responses - Add unit tests for Fabric chat URL builder - Cover colon-port, host:port, and IP address inputs
This commit is contained in:
parent
ae6d4d1fb3
commit
29a32a8439
|
|
@ -217,6 +217,10 @@ func (f APIConvert) ollamaChat(c *gin.Context) {
|
|||
}
|
||||
defer fabricRes.Body.Close()
|
||||
|
||||
if prompt.Stream {
|
||||
c.Header("Content-Type", "application/x-ndjson")
|
||||
}
|
||||
|
||||
if fabricRes.StatusCode < http.StatusOK || fabricRes.StatusCode >= http.StatusMultipleChoices {
|
||||
bodyBytes, readErr := io.ReadAll(fabricRes.Body)
|
||||
if readErr != nil {
|
||||
|
|
@ -234,10 +238,6 @@ func (f APIConvert) ollamaChat(c *gin.Context) {
|
|||
return
|
||||
}
|
||||
|
||||
if prompt.Stream {
|
||||
c.Header("Content-Type", "application/x-ndjson")
|
||||
}
|
||||
|
||||
var contentBuilder strings.Builder
|
||||
scanner := bufio.NewScanner(fabricRes.Body)
|
||||
scanner.Buffer(make([]byte, 0, 64*1024), 1024*1024)
|
||||
|
|
@ -304,9 +304,7 @@ func (f APIConvert) ollamaChat(c *gin.Context) {
|
|||
Done: true,
|
||||
TotalDuration: time.Since(now).Nanoseconds(),
|
||||
LoadDuration: int(time.Since(now).Nanoseconds()),
|
||||
PromptEvalCount: 42,
|
||||
PromptEvalDuration: int(time.Since(now).Nanoseconds()),
|
||||
EvalCount: 420,
|
||||
EvalDuration: time.Since(now).Nanoseconds(),
|
||||
}
|
||||
c.JSON(200, response)
|
||||
|
|
@ -327,9 +325,7 @@ func (f APIConvert) ollamaChat(c *gin.Context) {
|
|||
Done: true,
|
||||
TotalDuration: time.Since(now).Nanoseconds(),
|
||||
LoadDuration: int(time.Since(now).Nanoseconds()),
|
||||
PromptEvalCount: 42,
|
||||
PromptEvalDuration: int(time.Since(now).Nanoseconds()),
|
||||
EvalCount: 420,
|
||||
EvalDuration: time.Since(now).Nanoseconds(),
|
||||
}
|
||||
if err := writeOllamaResponseStruct(c, finalResponse); err != nil {
|
||||
|
|
@ -346,6 +342,9 @@ func buildFabricChatURL(addr string) (string, error) {
|
|||
if err != nil {
|
||||
return "", fmt.Errorf("invalid address: %w", err)
|
||||
}
|
||||
if parsed.Host == "" {
|
||||
return "", fmt.Errorf("invalid address: missing host")
|
||||
}
|
||||
return strings.TrimRight(parsed.String(), "/"), nil
|
||||
}
|
||||
if strings.HasPrefix(addr, ":") {
|
||||
|
|
|
|||
88
internal/server/ollama_test.go
Normal file
88
internal/server/ollama_test.go
Normal file
|
|
@ -0,0 +1,88 @@
|
|||
package restapi
|
||||
|
||||
import (
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestBuildFabricChatURL(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
addr string
|
||||
want string
|
||||
wantErr bool
|
||||
}{
|
||||
{
|
||||
name: "empty address",
|
||||
addr: "",
|
||||
want: "",
|
||||
wantErr: true,
|
||||
},
|
||||
{
|
||||
name: "valid http URL",
|
||||
addr: "http://localhost:8080",
|
||||
want: "http://localhost:8080",
|
||||
wantErr: false,
|
||||
},
|
||||
{
|
||||
name: "valid https URL",
|
||||
addr: "https://api.example.com",
|
||||
want: "https://api.example.com",
|
||||
wantErr: false,
|
||||
},
|
||||
{
|
||||
name: "http URL with trailing slash",
|
||||
addr: "http://localhost:8080/",
|
||||
want: "http://localhost:8080",
|
||||
wantErr: false,
|
||||
},
|
||||
{
|
||||
name: "malformed URL - missing host",
|
||||
addr: "http://",
|
||||
want: "",
|
||||
wantErr: true,
|
||||
},
|
||||
{
|
||||
name: "malformed URL - port only with http",
|
||||
addr: "https://:8080",
|
||||
want: "https://:8080",
|
||||
wantErr: false,
|
||||
},
|
||||
{
|
||||
name: "colon-prefixed port",
|
||||
addr: ":8080",
|
||||
want: "http://127.0.0.1:8080",
|
||||
wantErr: false,
|
||||
},
|
||||
{
|
||||
name: "bare host:port",
|
||||
addr: "localhost:8080",
|
||||
want: "http://localhost:8080",
|
||||
wantErr: false,
|
||||
},
|
||||
{
|
||||
name: "bare hostname",
|
||||
addr: "localhost",
|
||||
want: "http://localhost",
|
||||
wantErr: false,
|
||||
},
|
||||
{
|
||||
name: "IP address with port",
|
||||
addr: "192.168.1.1:3000",
|
||||
want: "http://192.168.1.1:3000",
|
||||
wantErr: false,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
got, err := buildFabricChatURL(tt.addr)
|
||||
if (err != nil) != tt.wantErr {
|
||||
t.Errorf("buildFabricChatURL() error = %v, wantErr %v", err, tt.wantErr)
|
||||
return
|
||||
}
|
||||
if got != tt.want {
|
||||
t.Errorf("buildFabricChatURL() = %v, want %v", got, tt.want)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
Loading…
Reference in a new issue