Skip to content
GitLab
Menu
Projects
Groups
Snippets
Loading...
Help
Help
Support
Community forum
Keyboard shortcuts
?
Submit feedback
Contribute to GitLab
Sign in / Register
Toggle navigation
Menu
Open sidebar
orangecat
ollama
Commits
1f9078d6
"torchvision/vscode:/vscode.git/clone" did not exist on "a00a72b1ee41483407717379fb5cafe992de2f82"
Unverified
Commit
1f9078d6
authored
Feb 12, 2024
by
Jeffrey Morgan
Committed by
GitHub
Feb 12, 2024
Browse files
Check image filetype in api handlers (#2467)
parent
26b13fc3
Changes
2
Hide whitespace changes
Inline
Side-by-side
Showing
2 changed files
with
24 additions
and
1 deletion
+24
-1
cmd/interactive.go
cmd/interactive.go
+1
-1
server/routes.go
server/routes.go
+23
-0
No files found.
cmd/interactive.go
View file @
1f9078d6
...
@@ -625,7 +625,7 @@ func getImageData(filePath string) ([]byte, error) {
...
@@ -625,7 +625,7 @@ func getImageData(filePath string) ([]byte, error) {
}
}
contentType
:=
http
.
DetectContentType
(
buf
)
contentType
:=
http
.
DetectContentType
(
buf
)
allowedTypes
:=
[]
string
{
"image/jpeg"
,
"image/jpg"
,
"image/svg+xml"
,
"image/png"
}
allowedTypes
:=
[]
string
{
"image/jpeg"
,
"image/jpg"
,
"image/png"
}
if
!
slices
.
Contains
(
allowedTypes
,
contentType
)
{
if
!
slices
.
Contains
(
allowedTypes
,
contentType
)
{
return
nil
,
fmt
.
Errorf
(
"invalid image type: %s"
,
contentType
)
return
nil
,
fmt
.
Errorf
(
"invalid image type: %s"
,
contentType
)
}
}
...
...
server/routes.go
View file @
1f9078d6
...
@@ -22,6 +22,7 @@ import (
...
@@ -22,6 +22,7 @@ import (
"github.com/gin-contrib/cors"
"github.com/gin-contrib/cors"
"github.com/gin-gonic/gin"
"github.com/gin-gonic/gin"
"golang.org/x/exp/slices"
"github.com/jmorganca/ollama/api"
"github.com/jmorganca/ollama/api"
"github.com/jmorganca/ollama/gpu"
"github.com/jmorganca/ollama/gpu"
...
@@ -136,6 +137,12 @@ func modelOptions(model *Model, requestOpts map[string]interface{}) (api.Options
...
@@ -136,6 +137,12 @@ func modelOptions(model *Model, requestOpts map[string]interface{}) (api.Options
return
opts
,
nil
return
opts
,
nil
}
}
func
isSupportedImageType
(
image
[]
byte
)
bool
{
contentType
:=
http
.
DetectContentType
(
image
)
allowedTypes
:=
[]
string
{
"image/jpeg"
,
"image/jpg"
,
"image/png"
}
return
slices
.
Contains
(
allowedTypes
,
contentType
)
}
func
GenerateHandler
(
c
*
gin
.
Context
)
{
func
GenerateHandler
(
c
*
gin
.
Context
)
{
loaded
.
mu
.
Lock
()
loaded
.
mu
.
Lock
()
defer
loaded
.
mu
.
Unlock
()
defer
loaded
.
mu
.
Unlock
()
...
@@ -166,6 +173,13 @@ func GenerateHandler(c *gin.Context) {
...
@@ -166,6 +173,13 @@ func GenerateHandler(c *gin.Context) {
return
return
}
}
for
_
,
img
:=
range
req
.
Images
{
if
!
isSupportedImageType
(
img
)
{
c
.
AbortWithStatusJSON
(
http
.
StatusBadRequest
,
gin
.
H
{
"error"
:
"unsupported image format"
})
return
}
}
model
,
err
:=
GetModel
(
req
.
Model
)
model
,
err
:=
GetModel
(
req
.
Model
)
if
err
!=
nil
{
if
err
!=
nil
{
var
pErr
*
fs
.
PathError
var
pErr
*
fs
.
PathError
...
@@ -1103,6 +1117,15 @@ func ChatHandler(c *gin.Context) {
...
@@ -1103,6 +1117,15 @@ func ChatHandler(c *gin.Context) {
return
return
}
}
for
_
,
msg
:=
range
req
.
Messages
{
for
_
,
img
:=
range
msg
.
Images
{
if
!
isSupportedImageType
(
img
)
{
c
.
AbortWithStatusJSON
(
http
.
StatusBadRequest
,
gin
.
H
{
"error"
:
"unsupported image format"
})
return
}
}
}
model
,
err
:=
GetModel
(
req
.
Model
)
model
,
err
:=
GetModel
(
req
.
Model
)
if
err
!=
nil
{
if
err
!=
nil
{
var
pErr
*
fs
.
PathError
var
pErr
*
fs
.
PathError
...
...
Write
Preview
Markdown
is supported
0%
Try again
or
attach a new file
.
Attach a file
Cancel
You are about to add
0
people
to the discussion. Proceed with caution.
Finish editing this message first!
Cancel
Please
register
or
sign in
to comment