Golang websocket.NewClient函数代码示例

func TestEcho(t *testing.T) {
	http.Handle("/echo", websocket.Handler(echoServer))
	once.Do(startCherryServer)

	client, err := net.Dial("tcp", serverAddr)
	if err != nil {
		t.Fatal("dialing", err)
	}

	conn, err := websocket.NewClient(newConfig(t, "/echo"), client)
	if err != nil {
		t.Errorf("WebSocket handshake error: %v", err)
		return
	}

	msg := []byte("hello, world\n")
	if _, err := conn.Write(msg); err != nil {
		t.Errorf("Write: %v", err)
	}
	var actual_msg = make([]byte, 512)
	n, err := conn.Read(actual_msg)
	if err != nil {
		t.Errorf("Read: %v", err)
	}
	actual_msg = actual_msg[0:n]
	if !bytes.Equal(msg, actual_msg) {
		t.Errorf("Echo: expected %q got %q", msg, actual_msg)
	}
	conn.Close()
}

func (st *serverTester) websocket(rp requestParam) (*serverResponse, error) {
	urlstring := st.url + "/echo"

	config, err := websocket.NewConfig(urlstring, st.url)
	if err != nil {
		st.t.Fatalf("websocket.NewConfig(%q, %q) returned error: %v", urlstring, st.url, err)
	}

	config.Header.Add("Test-Case", rp.name)
	for _, h := range rp.header {
		config.Header.Add(h.Name, h.Value)
	}

	ws, err := websocket.NewClient(config, st.conn)
	if err != nil {
		st.t.Fatalf("Error creating websocket client: %v", err)
	}

	if _, err := ws.Write(rp.body); err != nil {
		st.t.Fatalf("ws.Write() returned error: %v", err)
	}

	msg := make([]byte, 1024)
	var n int
	if n, err = ws.Read(msg); err != nil {
		st.t.Fatalf("ws.Read() returned error: %v", err)
	}

	res := &serverResponse{
		body: msg[:n],
	}

	return res, nil
}

func (s *S) TestServeHTTPWebSocket(c *check.C) {
	var servers []*httptest.Server
	for i := 0; i < 2; i++ {
		msg := fmt.Sprintf("server-%d", i)
		srv := httptest.NewServer(websocket.Handler(func(conn *websocket.Conn) {
			conn.Write([]byte(msg + conn.Request().URL.Path))
			buf := make([]byte, 5)
			conn.Read(buf)
			conn.Write(buf)
		}))
		defer srv.Close()
		servers = append(servers, srv)
	}
	var err error
	_, err = s.redis.Do("RPUSH", "frontend:myfrontend.com", "myfrontend", servers[0].URL, servers[1].URL)
	c.Assert(err, check.IsNil)
	router := Router{}
	err = router.Init()
	c.Assert(err, check.IsNil)
	proxyServer := httptest.NewServer(&router)
	defer proxyServer.Close()
	dialWS := func() *websocket.Conn {
		config, err := websocket.NewConfig("ws://myfrontend.com", "ws://localhost/")
		c.Assert(err, check.IsNil)
		url, _ := url.Parse(proxyServer.URL)
		client, err := net.Dial("tcp", url.Host)
		c.Assert(err, check.IsNil)
		conn, err := websocket.NewClient(config, client)
		c.Assert(err, check.IsNil)
		return conn
	}
	msgBuf := make([]byte, 100)
	conn := dialWS()
	defer conn.Close()
	n, err := conn.Read(msgBuf)
	c.Assert(err, check.IsNil)
	c.Assert(n, check.Equals, 9)
	c.Assert(string(msgBuf[:n]), check.Equals, "server-0/")
	n, err = conn.Write([]byte("12345"))
	c.Assert(err, check.IsNil)
	c.Assert(n, check.Equals, 5)
	n, err = conn.Read(msgBuf)
	c.Assert(err, check.IsNil)
	c.Assert(n, check.Equals, 5)
	c.Assert(string(msgBuf[:n]), check.Equals, "12345")
	conn = dialWS()
	defer conn.Close()
	n, err = conn.Read(msgBuf)
	c.Assert(err, check.IsNil)
	c.Assert(n, check.Equals, 9)
	c.Assert(string(msgBuf[:n]), check.Equals, "server-1/")
	n, err = conn.Write([]byte("12345"))
	c.Assert(err, check.IsNil)
	c.Assert(n, check.Equals, 5)
	n, err = conn.Read(msgBuf)
	c.Assert(err, check.IsNil)
	c.Assert(n, check.Equals, 5)
	c.Assert(string(msgBuf[:n]), check.Equals, "12345")
}

func (c *Client) proxyWebsocket(config *websocket.Config, proxy string) (*websocket.Conn, error) {
	u, err := url.Parse(proxy)

	if err != nil {
		return nil, err
	}

	host := u.Host

	if !strings.Contains(host, ":") {
		host += ":443"
	}

	conn, err := net.DialTimeout("tcp", u.Host, 3*time.Second)

	if err != nil {
		return nil, err
	}

	if _, err = conn.Write([]byte(fmt.Sprintf("CONNECT %s:443 HTTP/1.1\r\n", c.Host))); err != nil {
		return nil, err
	}

	if _, err = conn.Write([]byte(fmt.Sprintf("Host: %s:443\r\n", c.Host))); err != nil {
		return nil, err
	}

	if auth := u.User; auth != nil {
		enc := base64.StdEncoding.EncodeToString([]byte(auth.String()))

		if _, err = conn.Write([]byte(fmt.Sprintf("Proxy-Authorization: Basic %s\r\n", enc))); err != nil {
			return nil, err
		}
	}

	if _, err = conn.Write([]byte("Proxy-Connection: Keep-Alive\r\n\r\n")); err != nil {
		return nil, err
	}

	data, err := bufio.NewReader(conn).ReadString('\n')

	if err != nil {
		return nil, err
	}

	// need an http 200 response
	if !strings.Contains(string(data), " 200 ") {
		return nil, fmt.Errorf("proxy error: %s", strings.TrimSpace(string(data)))
	}

	return websocket.NewClient(config, tls.Client(conn, config.TlsConfig))
}

func (h DockerHandler) GetWebSocket(w http.ResponseWriter, r *http.Request) (*websocket.Conn, error) {
	client, err := net.Dial("unix", h.dockerEndpoint)
	if err != nil {
		helpers.ErrorLog("DOCKER", fmt.Sprintf("Following error occured when trying to make request to Docker endpoint %s : %s", h.dockerEndpoint, err.Error()))
		return nil, err
	}
	config := &websocket.Config{}
	config.Version = websocket.ProtocolVersionHybi13
	config.Location = r.URL
	config.Origin = r.URL
	config.Header = r.Header
	return websocket.NewClient(config, client)
}

// dialSocketListener opens a WebSocket client connection to pl with config.
// config.Location, config.Origin, and config.Version must be set.
func dialSocketListener(pl *pipeListener, config *websocket.Config) (
	conn *websocket.Conn, err error) {

	socket, err := pl.Dial("", "")
	if err != nil {
		return nil, err
	}
	if conn, err = websocket.NewClient(config, socket); err != nil {
		socket.Close()
		return nil, err
	}
	return conn, nil
}

func (w WS) Select(netcon net.Conn, url string) (*websocket.Conn, error) {
	conf, err := websocket.NewConfig(url, url)
	if err != nil {
		return nil, err
	}

	wcon, err := websocket.NewClient(conf, netcon)
	if err != nil {
		return nil, err
	}

	return wcon, nil
}

func (s *DockerSuite) TestGetContainersAttachWebsocket(c *check.C) {
	testRequires(c, DaemonIsLinux)
	out, _ := dockerCmd(c, "run", "-dit", "busybox", "cat")

	rwc, err := sockConn(time.Duration(10*time.Second), "")
	c.Assert(err, checker.IsNil)

	cleanedContainerID := strings.TrimSpace(out)
	config, err := websocket.NewConfig(
		"/containers/"+cleanedContainerID+"/attach/ws?stream=1&stdin=1&stdout=1&stderr=1",
		"http://localhost",
	)
	c.Assert(err, checker.IsNil)

	ws, err := websocket.NewClient(config, rwc)
	c.Assert(err, checker.IsNil)
	defer ws.Close()

	expected := []byte("hello")
	actual := make([]byte, len(expected))

	outChan := make(chan error)
	go func() {
		_, err := io.ReadFull(ws, actual)
		outChan <- err
		close(outChan)
	}()

	inChan := make(chan error)
	go func() {
		_, err := ws.Write(expected)
		inChan <- err
		close(inChan)
	}()

	select {
	case err := <-inChan:
		c.Assert(err, checker.IsNil)
	case <-time.After(5 * time.Second):
		c.Fatal("Timeout writing to ws")
	}

	select {
	case err := <-outChan:
		c.Assert(err, checker.IsNil)
	case <-time.After(5 * time.Second):
		c.Fatal("Timeout reading from ws")
	}

	c.Assert(actual, checker.DeepEquals, expected, check.Commentf("Websocket didn't return the expected data"))
}

// Connects to the tab/process for sending/recv'ing debug events
func wsConnection(addr, url string) *websocket.Conn {
	conn, err := net.Dial("tcp", addr)
	if err != nil {
		log.Fatal("error dialing ChromeTarget websocket: ", err)
	}
	config, errConfig := websocket.NewConfig(url, "http://localhost")
	if errConfig != nil {
		log.Fatalf("error building websocket config: addr: %s url: %s %v\n", addr, url, err)
	}
	client, errWS := websocket.NewClient(config, conn)
	if errWS != nil {
		log.Fatalf("error during websocket handshake: %v\n", errWS)
	}
	return client
}

func newClient() (*websocket.Conn, error) {
	//todo
L:
	client, err := net.Dial("tcp", cfg.AgentAddr())
	if err != nil {
		log.Warning("not connected to agent server, try again ...")
		time.Sleep(1 * time.Second)
		goto L
	}
	conn, err := websocket.NewClient(newConfig_("/"), client)
	if err != nil {
		log.Errorf("WebSocket handshake error: %v", err)
		return nil, err
	}
	return conn, nil
}

func TestCount(t *testing.T) {
	http.Handle("/count", websocket.Handler(countServer))
	once.Do(startCherryServer)

	// websocket.Dial()
	client, err := net.Dial("tcp", serverAddr)
	if err != nil {
		t.Fatal("dialing", err)
	}
	conn, err := websocket.NewClient(newConfig(t, "/count"), client)
	if err != nil {
		t.Errorf("WebSocket handshake error: %v", err)
		return
	}

	var count Count
	count.S = "hello"
	if err := websocket.JSON.Send(conn, count); err != nil {
		t.Errorf("Write: %v", err)
	}

	if err := websocket.JSON.Receive(conn, &count); err != nil {
		t.Errorf("Read: %v", err)
	}
	if count.N != 1 {
		t.Errorf("count: expected %d got %d", 1, count.N)
	}
	if count.S != "hello" {
		t.Errorf("count: expected %q got %q", "hello", count.S)
	}
	if err := websocket.JSON.Send(conn, count); err != nil {
		t.Errorf("Write: %v", err)
	}
	if err := websocket.JSON.Receive(conn, &count); err != nil {
		t.Errorf("Read: %v", err)
	}
	if count.N != 2 {
		t.Errorf("count: expected %d got %d", 2, count.N)
	}
	if count.S != "hellohello" {
		t.Errorf("count: expected %q got %q", "hellohello", count.S)
	}

	conn.Close()
}

func newWSConn(originURL, user, pass string) (*wsConn, error) {
	v := url.Values{}
	v.Set("login", user)
	v.Set("password", pass)
	params := v.Encode()
	wsURL := strings.Replace(originURL, "http", "ws", 1) + "?" + params

	wsConfig, err := websocket.NewConfig(wsURL, originURL)
	if err != nil {
		return nil, err
	}
	wsConfig.TlsConfig = &tls.Config{RootCAs: rootCA}

	var wsc *websocket.Conn

	req := http.Request{
		URL: &url.URL{},
	}

	proxyURL, err := getProxy(&req)
	if err != nil {
		return nil, err
	}
	if proxyURL == nil {
		wsc, err = websocket.DialConfig(wsConfig)
		if err != nil {
			return nil, err
		}
	} else {
		conn, err := httpConnect(proxyURL.Host, wsConfig)
		if err != nil {
			return nil, err
		}
		if wsConfig.Location.Scheme == "wss" {
			conn = tls.Client(conn, wsConfig.TlsConfig)
		}

		wsc, err = websocket.NewClient(wsConfig, conn)
		if err != nil {
			return nil, err
		}
	}
	return &wsConn{conn: wsc}, nil
}

func NewClient(addr string, m ClientModel, out chan string) (*Client, error) {
L:
	client, err := net.Dial("tcp", addr)
	if err != nil {
		log.Warningf("not connected to %+v, try again ...", addr)
		time.Sleep(1 * time.Second)
		goto L
	}
	conn, err := websocket.NewClient(newConfig("/"), client)
	if err != nil {
		log.Errorf("WebSocket handshake error: %v", err)
		return nil, err
	}
	rwc := NewWsReadWriter(conn)
	return &Client{
		ReadWriteCloser: rwc,
		ClientModel:     m,
		send:            make(chan string, 10),
		ctrl:            out,
	}, nil
}

func addWsClient(t *testing.T, roomID string, clientID string) *websocket.Conn {
	c, err := net.Dial("tcp", serverAddr)
	if err != nil {
		t.Fatalf("net.Dial(tcp, %q) got error: %s, want nil", serverAddr, err.Error())
	}
	config := newConfig(t, "/ws")
	conn, err := websocket.NewClient(config, c)
	if err != nil {
		t.Fatalf("websocket.NewClient(%v, %v) got error: %s, want nil", config, c, err.Error())
	}

	// Registers the client.
	m := wsClientMsg{
		Cmd:      "register",
		ClientID: clientID,
		RoomID:   roomID,
	}
	write(t, conn, m)

	return conn
}

func websocketProxyDial(urlString, origin string) (ws *websocket.Conn, err error) {
	if os.Getenv("HTTP_PROXY") == "" {
		return websocket.Dial(urlString, "", origin)
	}

	purl, err := url.Parse(os.Getenv("HTTP_PROXY"))
	if err != nil {
		return nil, err
	}

	config, err := websocket.NewConfig(urlString, origin)
	if err != nil {
		return nil, err
	}

	client, err := websocketHTTPConnect(purl.Host, urlString)
	if err != nil {
		return nil, err
	}

	switch config.Location.Scheme {
	case "ws":
	case "wss":
		tlsClient := tls.Client(client, &tls.Config{
			ServerName: strings.Split(config.Location.Host, ":")[0],
		})
		err := tlsClient.Handshake()
		if err != nil {
			tlsClient.Close()
			return nil, err
		}
		client = tlsClient

	default:
		return nil, errors.New("invalid websocket schema")
	}

	return websocket.NewClient(config, client)
}