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) }