violet/servers/https_test.go

63 lines
1.5 KiB
Go
Raw Normal View History

2023-06-04 22:28:48 +01:00
package servers
import (
"database/sql"
2023-07-22 01:11:47 +01:00
"github.com/1f349/violet/certs"
"github.com/1f349/violet/proxy"
2023-08-17 14:38:00 +01:00
"github.com/1f349/violet/proxy/websocket"
2023-07-22 01:11:47 +01:00
"github.com/1f349/violet/router"
"github.com/1f349/violet/servers/conf"
"github.com/1f349/violet/utils/fake"
2023-06-04 22:28:48 +01:00
_ "github.com/mattn/go-sqlite3"
"github.com/stretchr/testify/assert"
"net/http"
"net/http/httptest"
"sync"
"testing"
)
type fakeTransport struct{}
func (f *fakeTransport) RoundTrip(_ *http.Request) (*http.Response, error) {
rec := httptest.NewRecorder()
rec.WriteHeader(http.StatusOK)
return rec.Result(), nil
}
func TestNewHttpsServer_RateLimit(t *testing.T) {
db, err := sql.Open("sqlite3", "file::memory:?cache=shared")
assert.NoError(t, err)
ft := &fakeTransport{}
httpsConf := &conf.Conf{
2023-06-04 22:28:48 +01:00
RateLimit: 5,
Domains: &fake.Domains{},
2023-06-04 22:28:48 +01:00
Certs: certs.New(nil, nil, true),
Signer: fake.SnakeOilProv,
2023-08-17 14:38:00 +01:00
Router: router.NewManager(db, proxy.NewHybridTransportWithCalls(ft, ft, &websocket.Server{})),
2023-06-04 22:28:48 +01:00
}
srv := NewHttpsServer(httpsConf)
req, err := http.NewRequest(http.MethodGet, "https://example.com", nil)
req.RemoteAddr = "127.0.0.1:1447"
assert.NoError(t, err)
wg := &sync.WaitGroup{}
wg.Add(5)
for i := 0; i < 5; i++ {
go func() {
defer wg.Done()
rec := httptest.NewRecorder()
srv.Handler.ServeHTTP(rec, req)
res := rec.Result()
assert.Equal(t, http.StatusTeapot, res.StatusCode)
2023-06-04 22:28:48 +01:00
}()
}
wg.Wait()
rec := httptest.NewRecorder()
srv.Handler.ServeHTTP(rec, req)
res := rec.Result()
assert.Equal(t, http.StatusTooManyRequests, res.StatusCode)
}