35 lines
1.0 KiB
Go
35 lines
1.0 KiB
Go
|
package chat
|
||
|
|
|
||
|
|
import (
|
||
|
|
"net/http"
|
||
|
|
"net/http/httptest"
|
||
|
|
"testing"
|
||
|
|
|
||
|
|
secutils "github.com/Tencent/WeKnora/internal/utils"
|
||
|
|
)
|
||
|
|
|
||
|
|
func TestResolveImageForOllamaRejectsInternalURL(t *testing.T) {
|
||
|
|
t.Setenv("SSRF_WHITELIST", "")
|
||
|
|
secutils.ResetSSRFWhitelistForTest()
|
||
|
|
t.Cleanup(secutils.ResetSSRFWhitelistForTest)
|
||
|
|
|
||
|
|
if data := resolveImageForOllama("http://169.254.169.254/latest/meta-data/"); data != nil {
|
||
|
|
t.Fatalf("resolveImageForOllama returned data for blocked internal URL")
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestResolveImageForOllamaBlocksRedirectToInternalURL(t *testing.T) {
|
||
|
|
t.Setenv("SSRF_WHITELIST", "127.0.0.1")
|
||
|
|
secutils.ResetSSRFWhitelistForTest()
|
||
|
|
t.Cleanup(secutils.ResetSSRFWhitelistForTest)
|
||
|
|
|
||
|
|
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||
|
|
http.Redirect(w, r, "http://169.254.169.254/latest/meta-data/", http.StatusFound)
|
||
|
|
}))
|
||
|
|
defer server.Close()
|
||
|
|
|
||
|
|
if data := resolveImageForOllama(server.URL); data != nil {
|
||
|
|
t.Fatalf("resolveImageForOllama returned data after redirect to blocked internal URL")
|
||
|
|
}
|
||
|
|
}
|