traefik/integration/websocket_test.go
Ludovic Fernandez 759c269dee Code cleaning.
2017-08-18 02:18:02 +02:00

234 lines
5.7 KiB
Go

package integration
import (
"net"
"net/http"
"net/http/httptest"
"os"
"time"
"github.com/containous/traefik/integration/try"
"github.com/go-check/check"
gorillawebsocket "github.com/gorilla/websocket"
checker "github.com/vdemeester/shakers"
"golang.org/x/net/websocket"
)
// WebsocketSuite
type WebsocketSuite struct{ BaseSuite }
func (suite *WebsocketSuite) TestBase(c *check.C) {
var upgrader = gorillawebsocket.Upgrader{} // use default options
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
c, err := upgrader.Upgrade(w, r, nil)
if err != nil {
return
}
defer c.Close()
for {
mt, message, err := c.ReadMessage()
if err != nil {
break
}
err = c.WriteMessage(mt, message)
if err != nil {
break
}
}
}))
file := suite.adaptFile(c, "fixtures/websocket/config.toml", struct {
WebsocketServer string
}{
WebsocketServer: srv.URL,
})
defer os.Remove(file)
cmd, _ := suite.cmdTraefik(withConfigFile(file), "--debug")
err := cmd.Start()
c.Assert(err, check.IsNil)
defer cmd.Process.Kill()
// wait for traefik
err = try.GetRequest("http://127.0.0.1:8080/api/providers", 10*time.Second, try.BodyContains("127.0.0.1"))
c.Assert(err, checker.IsNil)
conn, _, err := gorillawebsocket.DefaultDialer.Dial("ws://127.0.0.1:8000/ws", nil)
c.Assert(err, checker.IsNil)
err = conn.WriteMessage(gorillawebsocket.TextMessage, []byte("OK"))
c.Assert(err, checker.IsNil)
_, msg, err := conn.ReadMessage()
c.Assert(err, checker.IsNil)
c.Assert(string(msg), checker.Equals, "OK")
}
func (suite *WebsocketSuite) TestWrongOrigin(c *check.C) {
var upgrader = gorillawebsocket.Upgrader{} // use default options
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
c, err := upgrader.Upgrade(w, r, nil)
if err != nil {
return
}
defer c.Close()
for {
mt, message, err := c.ReadMessage()
if err != nil {
break
}
err = c.WriteMessage(mt, message)
if err != nil {
break
}
}
}))
file := suite.adaptFile(c, "fixtures/websocket/config.toml", struct {
WebsocketServer string
}{
WebsocketServer: srv.URL,
})
defer os.Remove(file)
cmd, _ := suite.cmdTraefik(withConfigFile(file), "--debug")
err := cmd.Start()
c.Assert(err, check.IsNil)
defer cmd.Process.Kill()
// wait for traefik
err = try.GetRequest("http://127.0.0.1:8080/api/providers", 10*time.Second, try.BodyContains("127.0.0.1"))
c.Assert(err, checker.IsNil)
config, err := websocket.NewConfig("ws://127.0.0.1:8000/ws", "ws://127.0.0.1:800")
c.Assert(err, check.IsNil)
conn, err := net.DialTimeout("tcp", "127.0.0.1:8000", time.Second)
c.Assert(err, checker.IsNil)
_, err = websocket.NewClient(config, conn)
c.Assert(err, checker.NotNil)
c.Assert(err, checker.ErrorMatches, "bad status")
}
func (suite *WebsocketSuite) TestOrigin(c *check.C) {
// use default options
var upgrader = gorillawebsocket.Upgrader{}
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
c, err := upgrader.Upgrade(w, r, nil)
if err != nil {
return
}
defer c.Close()
for {
mt, message, err := c.ReadMessage()
if err != nil {
break
}
err = c.WriteMessage(mt, message)
if err != nil {
break
}
}
}))
file := suite.adaptFile(c, "fixtures/websocket/config.toml", struct {
WebsocketServer string
}{
WebsocketServer: srv.URL,
})
defer os.Remove(file)
cmd, _ := suite.cmdTraefik(withConfigFile(file), "--debug")
err := cmd.Start()
c.Assert(err, check.IsNil)
defer cmd.Process.Kill()
// wait for traefik
err = try.GetRequest("http://127.0.0.1:8080/api/providers", 10*time.Second, try.BodyContains("127.0.0.1"))
c.Assert(err, checker.IsNil)
config, err := websocket.NewConfig("ws://127.0.0.1:8000/ws", "ws://127.0.0.1:8000")
c.Assert(err, check.IsNil)
conn, err := net.DialTimeout("tcp", "127.0.0.1:8000", time.Second)
c.Assert(err, check.IsNil)
client, err := websocket.NewClient(config, conn)
c.Assert(err, checker.IsNil)
n, err := client.Write([]byte("OK"))
c.Assert(err, checker.IsNil)
c.Assert(n, checker.Equals, 2)
msg := make([]byte, 2)
n, err = client.Read(msg)
c.Assert(err, checker.IsNil)
c.Assert(n, checker.Equals, 2)
c.Assert(string(msg), checker.Equals, "OK")
}
func (suite *WebsocketSuite) TestWrongOriginIgnoredByServer(c *check.C) {
var upgrader = gorillawebsocket.Upgrader{CheckOrigin: func(r *http.Request) bool {
return true
}}
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
c, err := upgrader.Upgrade(w, r, nil)
if err != nil {
return
}
defer c.Close()
for {
mt, message, err := c.ReadMessage()
if err != nil {
break
}
err = c.WriteMessage(mt, message)
if err != nil {
break
}
}
}))
file := suite.adaptFile(c, "fixtures/websocket/config.toml", struct {
WebsocketServer string
}{
WebsocketServer: srv.URL,
})
defer os.Remove(file)
cmd, _ := suite.cmdTraefik(withConfigFile(file), "--debug")
err := cmd.Start()
c.Assert(err, check.IsNil)
defer cmd.Process.Kill()
// wait for traefik
err = try.GetRequest("http://127.0.0.1:8080/api/providers", 10*time.Second, try.BodyContains("127.0.0.1"))
c.Assert(err, checker.IsNil)
config, err := websocket.NewConfig("ws://127.0.0.1:8000/ws", "ws://127.0.0.1:80")
c.Assert(err, check.IsNil)
conn, err := net.DialTimeout("tcp", "127.0.0.1:8000", time.Second)
c.Assert(err, checker.IsNil)
client, err := websocket.NewClient(config, conn)
c.Assert(err, checker.IsNil)
n, err := client.Write([]byte("OK"))
c.Assert(err, checker.IsNil)
c.Assert(n, checker.Equals, 2)
msg := make([]byte, 2)
n, err = client.Read(msg)
c.Assert(err, checker.IsNil)
c.Assert(n, checker.Equals, 2)
c.Assert(string(msg), checker.Equals, "OK")
}