wehub-resource-sync
133 lines
4.2 KiB
Go
133 lines
4.2 KiB
Go
package httpclient_test
|
|
|
|
import (
|
|
"crypto/tls"
|
|
"errors"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/mudler/LocalAI/pkg/httpclient"
|
|
|
|
. "github.com/onsi/ginkgo/v2"
|
|
. "github.com/onsi/gomega"
|
|
)
|
|
|
|
func TestHTTPClient(t *testing.T) {
|
|
RegisterFailHandler(Fail)
|
|
RunSpecs(t, "httpclient suite")
|
|
}
|
|
|
|
var _ = Describe("httpclient", func() {
|
|
Describe("New (default)", func() {
|
|
It("refuses to follow redirects and never reaches the redirect target", func() {
|
|
sinkHit := make(chan string, 1)
|
|
sink := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
sinkHit <- r.Header.Get("X-Api-Key")
|
|
w.WriteHeader(http.StatusOK)
|
|
}))
|
|
defer sink.Close()
|
|
|
|
redirector := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
http.Redirect(w, r, sink.URL, http.StatusFound)
|
|
}))
|
|
defer redirector.Close()
|
|
|
|
req, _ := http.NewRequest(http.MethodGet, redirector.URL, nil)
|
|
req.Header.Set("X-Api-Key", "secret")
|
|
|
|
_, err := httpclient.New().Do(req)
|
|
Expect(err).To(HaveOccurred(), "redirect must surface as an error")
|
|
Expect(errors.Is(err, httpclient.ErrRedirectBlocked)).To(BeTrue(), "error should wrap ErrRedirectBlocked")
|
|
Expect(sinkHit).NotTo(Receive(), "the redirect target must never be contacted")
|
|
})
|
|
|
|
It("sets no overall timeout (streaming-safe) by default", func() {
|
|
Expect(httpclient.New().Timeout).To(BeZero())
|
|
})
|
|
|
|
It("sets a TLS 1.2 floor on the default transport", func() {
|
|
c := httpclient.New()
|
|
t, ok := c.Transport.(*http.Transport)
|
|
Expect(ok).To(BeTrue())
|
|
Expect(t.TLSClientConfig).NotTo(BeNil())
|
|
Expect(t.TLSClientConfig.MinVersion).To(Equal(uint16(tls.VersionTLS12)))
|
|
})
|
|
})
|
|
|
|
Describe("NewWithTimeout", func() {
|
|
It("applies the overall timeout", func() {
|
|
Expect(httpclient.NewWithTimeout(5 * time.Second).Timeout).To(Equal(5 * time.Second))
|
|
})
|
|
})
|
|
|
|
Describe("WithFollowRedirects", func() {
|
|
It("follows same-host redirects keeping the credential header", func() {
|
|
got := make(chan string, 2)
|
|
var srv *httptest.Server
|
|
srv = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
if r.URL.Path == "/start" {
|
|
http.Redirect(w, r, srv.URL+"/end", http.StatusFound)
|
|
return
|
|
}
|
|
got <- r.Header.Get("X-Api-Key")
|
|
w.WriteHeader(http.StatusOK)
|
|
}))
|
|
defer srv.Close()
|
|
|
|
req, _ := http.NewRequest(http.MethodGet, srv.URL+"/start", nil)
|
|
req.Header.Set("X-Api-Key", "secret")
|
|
|
|
resp, err := httpclient.New(httpclient.WithFollowRedirects()).Do(req)
|
|
Expect(err).NotTo(HaveOccurred())
|
|
_ = resp.Body.Close()
|
|
Expect(<-got).To(Equal("secret"), "same-host redirect should preserve the header")
|
|
})
|
|
|
|
It("strips credential headers on a cross-host redirect", func() {
|
|
sinkKey := make(chan string, 1)
|
|
sink := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
sinkKey <- r.Header.Get("X-Api-Key")
|
|
w.WriteHeader(http.StatusOK)
|
|
}))
|
|
defer sink.Close()
|
|
|
|
redirector := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
http.Redirect(w, r, sink.URL, http.StatusFound)
|
|
}))
|
|
defer redirector.Close()
|
|
|
|
req, _ := http.NewRequest(http.MethodGet, redirector.URL, nil)
|
|
req.Header.Set("X-Api-Key", "secret")
|
|
|
|
resp, err := httpclient.New(httpclient.WithFollowRedirects()).Do(req)
|
|
Expect(err).NotTo(HaveOccurred())
|
|
_ = resp.Body.Close()
|
|
Expect(<-sinkKey).To(BeEmpty(), "x-api-key must be stripped crossing to a different host")
|
|
})
|
|
})
|
|
|
|
Describe("Harden", func() {
|
|
It("adds NoRedirect and a TLS floor to a bare client without clobbering existing config", func() {
|
|
c := httpclient.Harden(&http.Client{})
|
|
Expect(c.CheckRedirect).NotTo(BeNil())
|
|
t, ok := c.Transport.(*http.Transport)
|
|
Expect(ok).To(BeTrue())
|
|
Expect(t.TLSClientConfig.MinVersion).To(Equal(uint16(tls.VersionTLS12)))
|
|
})
|
|
|
|
It("returns nil for a nil client", func() {
|
|
Expect(httpclient.Harden(nil)).To(BeNil())
|
|
})
|
|
|
|
It("preserves a caller-supplied CheckRedirect", func() {
|
|
sentinel := errors.New("mine")
|
|
c := httpclient.Harden(&http.Client{
|
|
CheckRedirect: func(*http.Request, []*http.Request) error { return sentinel },
|
|
})
|
|
Expect(c.CheckRedirect(nil, nil)).To(Equal(sentinel))
|
|
})
|
|
})
|
|
})
|