func TestBackendGoesAway(t *testing.T) {
dialer := &websocket.Dialer{}
headers := http.Header{}
headers.Add("X-Cattle-HostId", "1")
backendWs, _, err := dialer.Dial("ws://127.0.0.1:2222/connectbackend", headers)
if err != nil {
t.Fatal("Failed to connect to proxy.", err)
}
handlers := make(map[string]Handler)
handlers["/v1/echo"] = &echoHandler{}
go connectToProxyWS(backendWs, handlers)
signedToken := test_utils.CreateToken("1", privateKey)
url := "ws://localhost:2222/v1/echo?token=" + signedToken
ws := getClientConnection(url, t)
if err := ws.WriteMessage(1, []byte("a message")); err != nil {
t.Fatal(err)
}
backendWs.Close()
if _, _, err := ws.ReadMessage(); err != io.EOF {
t.Fatal("Expected error indicating websocket was closed.")
}
dialer = &websocket.Dialer{}
ws, _, err = dialer.Dial(url, http.Header{})
if ws != nil || err != websocket.ErrBadHandshake {
t.Fatal("Should not have been able to connect.")
}
}
func TestHostStatsLegacy(t *testing.T) {
dialer := &websocket.Dialer{}
headers := http.Header{}
token := wsp_utils.CreateToken("1", privateKey)
url := "ws://localhost:1111/v1/stats?token=" + token
ws, _, err := dialer.Dial(url, headers)
if err != nil {
t.Fatal(err)
}
defer ws.Close()
count := 0
for {
if count > 3 {
break
}
_, msg, err := ws.ReadMessage()
if err != nil {
t.Fatal(err)
}
stats := string(msg)
if !strings.Contains(stats, "cpu") {
t.Fatalf("Stats are not working. Output: [%s]", stats)
}
count++
}
}
func TestEndToEnd(t *testing.T) {
signedToken := test_utils.CreateToken("1", privateKey)
ws := getClientConnection("ws://localhost:1111/v1/echo?token="+signedToken, t)
sendAndAssertReply(ws, strconv.FormatInt(time.Now().UnixNano()/int64(time.Millisecond), 10), t)
time.Sleep(1 * time.Millisecond) // Ensure different timestamp
sendAndAssertReply(ws, strconv.FormatInt(time.Now().UnixNano()/int64(time.Millisecond), 10), t)
}
func TestFrontendClosesConnection(t *testing.T) {
signedToken := test_utils.CreateToken("1", privateKey)
ws := getClientConnection("ws://localhost:1111/v1/oneanddone?token="+signedToken, t)
if err := ws.WriteControl(websocket.CloseMessage, nil, time.Now().Add(time.Second)); err != nil {
t.Fatal(err)
}
if _, _, err := ws.ReadMessage(); err == nil {
t.Fatal("Expected error indicating websocket was closed.")
}
}
func TestAuthHeaderBearerToken(t *testing.T) {
signedToken := test_utils.CreateToken("1", privateKey)
dialer := &websocket.Dialer{}
headers := http.Header{}
headers.Add("Authorization", "Bearer "+signedToken)
ws, _, err := dialer.Dial("ws://localhost:1111/v1/echo", headers)
if err != nil {
t.Fatal(err)
}
sendAndAssertReply(ws, strconv.FormatInt(time.Now().UnixNano()/int64(time.Millisecond), 10), t)
time.Sleep(1 * time.Millisecond) // Ensure different timestamp
sendAndAssertReply(ws, strconv.FormatInt(time.Now().UnixNano()/int64(time.Millisecond), 10), t)
}
func TestMultiHostStats(t *testing.T) {
payload := map[string]interface{}{
"project": []map[string]string{
{
"url": "ws://localhost:1111/v1/hostStats/project",
"token": test_utils.CreateToken("1", privateKey),
},
{
"url": "ws://localhost:1111/v1/hostStats/project",
"token": test_utils.CreateToken("2", privateKey),
},
},
}
signedToken := test_utils.CreateTokenWithPayload(payload, privateKey)
ws := getClientConnection("ws://localhost:1111/v1/hostStats/project?token="+signedToken, t)
one := false
two := false
for i := 0; i < 100; i++ {
err := ws.WriteMessage(1, []byte("x"))
if err != nil {
t.Fatal("Error talking to host")
}
_, msgBytes, err := ws.ReadMessage()
if err != nil {
t.Fatal("Error reading response from various containers")
}
if string(msgBytes) == "1" {
one = true
}
if string(msgBytes) == "2" {
two = true
}
if one && two {
return
}
}
t.Fatal("Did not get container stats from two hosts")
}
func TestBackendClosesConnection(t *testing.T) {
signedToken := test_utils.CreateToken("1", privateKey)
ws := getClientConnection("ws://localhost:1111/v1/oneanddone?token="+signedToken, t)
if err := ws.WriteMessage(1, []byte("a message")); err != nil {
t.Fatal(err)
}
if _, _, err := ws.ReadMessage(); err != nil {
t.Fatal(err)
}
if msgType, msgBytes, err := ws.ReadMessage(); err != io.EOF {
t.Fatalf("Expected an EOF error to indicate connection was closed. [%v] [%s] [%v]", msgType, msgBytes, err)
}
}
func TestBackendSendAfterClose(t *testing.T) {
signedToken := test_utils.CreateToken("1", privateKey)
ws := getClientConnection("ws://localhost:1111/v1/sendafterclose?token="+signedToken, t)
go func() {
for {
_, _, err := ws.ReadMessage()
if err != nil {
return
}
}
}()
ws2 := getClientConnection("ws://localhost:1111/v1/echo?token="+signedToken, t)
// If deadlock occurs, the read deadline will be hit and the sendAndAssertReply will fail the test
ws2.SetReadDeadline(time.Now().Add(2 * time.Second))
sendAndAssertReply(ws2, strconv.FormatInt(time.Now().UnixNano()/int64(time.Millisecond), 10), t)
}
func TestManyChattyConnections(t *testing.T) {
// Spin up a hundred connections. The repeat handler will send a new message to each one
// every 10 milliseconds. Stop after 5 seconds. This is just to prove that the proxy can handle a little load.
for i := 1; i <= 100; i++ {
signedToken := test_utils.CreateToken("1", privateKey)
msg := strconv.FormatInt(time.Now().UnixNano()/int64(time.Millisecond), 10)
ws := getClientConnection("ws://localhost:1111/v1/repeat?token="+signedToken+"&msg="+msg, t)
go func(expectedPrefix string) {
for {
_, reply, err := ws.ReadMessage()
if err != nil {
t.Fatal(err)
}
if !strings.HasPrefix(string(reply), expectedPrefix) {
t.Fatalf("Unexpected repsonse: [%v]", reply)
}
}
}(msg)
time.Sleep(1 * time.Millisecond) // Ensure different timestamp
}
time.Sleep(5 * time.Second)
}