ws_test.go

Documentation: github.com/moby/spdystream/ws

     1  /*
     2     Copyright 2014-2021 Docker Inc.
     3  
     4     Licensed under the Apache License, Version 2.0 (the "License");
     5     you may not use this file except in compliance with the License.
     6     You may obtain a copy of the License at
     7  
     8         http://www.apache.org/licenses/LICENSE-2.0
     9  
    10     Unless required by applicable law or agreed to in writing, software
    11     distributed under the License is distributed on an "AS IS" BASIS,
    12     WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    13     See the License for the specific language governing permissions and
    14     limitations under the License.
    15  */
    16  
    17  package ws
    18  
    19  import (
    20  	"bytes"
    21  	"io"
    22  	"log"
    23  	"net/http"
    24  	"net/http/httptest"
    25  	"strings"
    26  	"testing"
    27  
    28  	"github.com/moby/spdystream"
    29  	"github.com/gorilla/websocket"
    30  )
    31  
    32  var upgrader = websocket.Upgrader{
    33  	ReadBufferSize:  1024,
    34  	WriteBufferSize: 1024,
    35  }
    36  
    37  var serverSpdyConn *spdystream.Connection
    38  
    39  // Connect to the Websocket endpoint at ws://localhost
    40  // using SPDY over Websockets framing.
    41  func ExampleConn() {
    42  	wsconn, _, _ := websocket.DefaultDialer.Dial("ws://localhost/", http.Header{"Origin": {"http://localhost/"}})
    43  	conn, _ := spdystream.NewConnection(NewConnection(wsconn), false)
    44  	go conn.Serve(spdystream.NoOpStreamHandler)
    45  	stream, _ := conn.CreateStream(http.Header{}, nil, false)
    46  	stream.Wait()
    47  }
    48  
    49  func serveWs(w http.ResponseWriter, r *http.Request) {
    50  	if r.Method != "GET" {
    51  		http.Error(w, "Method not allowed", 405)
    52  		return
    53  	}
    54  
    55  	ws, err := upgrader.Upgrade(w, r, nil)
    56  	if err != nil {
    57  		if _, ok := err.(websocket.HandshakeError); !ok {
    58  			log.Println(err)
    59  		}
    60  		return
    61  	}
    62  
    63  	wrap := NewConnection(ws)
    64  	spdyConn, err := spdystream.NewConnection(wrap, true)
    65  	if err != nil {
    66  		log.Fatal(err)
    67  		return
    68  	}
    69  	serverSpdyConn = spdyConn
    70  	spdyConn.Serve(authStreamHandler)
    71  }
    72  
    73  func TestSpdyStreamOverWs(t *testing.T) {
    74  	server := httptest.NewServer(http.HandlerFunc(serveWs))
    75  	defer server.Close()
    76  	defer func() {
    77  		if serverSpdyConn != nil {
    78  			serverSpdyConn.Close()
    79  		}
    80  	}()
    81  
    82  	wsconn, _, err := websocket.DefaultDialer.Dial(strings.Replace(server.URL, "http://", "ws://", 1), http.Header{"Origin": {server.URL}})
    83  	if err != nil {
    84  		t.Fatal(err)
    85  	}
    86  
    87  	wrap := NewConnection(wsconn)
    88  	spdyConn, err := spdystream.NewConnection(wrap, false)
    89  	if err != nil {
    90  		defer wsconn.Close()
    91  		t.Fatal(err)
    92  	}
    93  	defer spdyConn.Close()
    94  	authenticated = true
    95  	go spdyConn.Serve(spdystream.NoOpStreamHandler)
    96  
    97  	stream, streamErr := spdyConn.CreateStream(http.Header{}, nil, false)
    98  	if streamErr != nil {
    99  		t.Fatalf("Error creating stream: %s", streamErr)
   100  	}
   101  
   102  	waitErr := stream.Wait()
   103  	if waitErr != nil {
   104  		t.Fatalf("Error waiting for stream: %s", waitErr)
   105  	}
   106  
   107  	message := []byte("hello")
   108  	writeErr := stream.WriteData(message, false)
   109  	if writeErr != nil {
   110  		t.Fatalf("Error writing data")
   111  	}
   112  
   113  	buf := make([]byte, 10)
   114  	n, readErr := stream.Read(buf)
   115  	if readErr != nil {
   116  		t.Fatalf("Error reading data from stream: %s", readErr)
   117  	}
   118  	if n != 5 {
   119  		t.Fatalf("Unexpected number of bytes read:\nActual: %d\nExpected: 5", n)
   120  	}
   121  	if !bytes.Equal(buf[:n], message) {
   122  		t.Fatalf("Did not receive expected message:\nActual: %s\nExpectd: %s", buf, message)
   123  	}
   124  
   125  	writeErr = stream.WriteData(message, true)
   126  	if writeErr != nil {
   127  		t.Fatalf("Error writing data")
   128  	}
   129  
   130  	smallBuf := make([]byte, 3)
   131  	n, readErr = stream.Read(smallBuf)
   132  	if readErr != nil {
   133  		t.Fatalf("Error reading data from stream: %s", readErr)
   134  	}
   135  	if n != 3 {
   136  		t.Fatalf("Unexpected number of bytes read:\nActual: %d\nExpected: 3", n)
   137  	}
   138  	if !bytes.Equal(smallBuf[:n], []byte("hel")) {
   139  		t.Fatalf("Did not receive expected message:\nActual: %s\nExpectd: %s", smallBuf[:n], message)
   140  	}
   141  	n, readErr = stream.Read(smallBuf)
   142  	if readErr != nil {
   143  		t.Fatalf("Error reading data from stream: %s", readErr)
   144  	}
   145  	if n != 2 {
   146  		t.Fatalf("Unexpected number of bytes read:\nActual: %d\nExpected: 2", n)
   147  	}
   148  	if !bytes.Equal(smallBuf[:n], []byte("lo")) {
   149  		t.Fatalf("Did not receive expected message:\nActual: %s\nExpected: lo", smallBuf[:n])
   150  	}
   151  
   152  	n, readErr = stream.Read(buf)
   153  	if readErr != io.EOF {
   154  		t.Fatalf("Expected EOF reading from finished stream, read %d bytes", n)
   155  	}
   156  
   157  	// Closing again should return error since the stream is already closed
   158  	streamCloseErr := stream.Close()
   159  	if streamCloseErr == nil {
   160  		t.Fatalf("No error closing finished stream")
   161  	}
   162  	if streamCloseErr != spdystream.ErrWriteClosedStream {
   163  		t.Fatalf("Unexpected error closing stream: %s", streamCloseErr)
   164  	}
   165  
   166  	streamResetErr := stream.Reset()
   167  	if streamResetErr != nil {
   168  		t.Fatalf("Error reseting stream: %s", streamResetErr)
   169  	}
   170  
   171  	authenticated = false
   172  	badStream, badStreamErr := spdyConn.CreateStream(http.Header{}, nil, false)
   173  	if badStreamErr != nil {
   174  		t.Fatalf("Error creating stream: %s", badStreamErr)
   175  	}
   176  
   177  	waitErr = badStream.Wait()
   178  	if waitErr == nil {
   179  		t.Fatalf("Did not receive error creating stream")
   180  	}
   181  	if waitErr != spdystream.ErrReset {
   182  		t.Fatalf("Unexpected error creating stream: %s", waitErr)
   183  	}
   184  
   185  	spdyCloseErr := spdyConn.Close()
   186  	if spdyCloseErr != nil {
   187  		t.Fatalf("Error closing spdy connection: %s", spdyCloseErr)
   188  	}
   189  }
   190  
   191  var authenticated bool
   192  
   193  func authStreamHandler(stream *spdystream.Stream) {
   194  	if !authenticated {
   195  		stream.Refuse()
   196  		return
   197  	}
   198  	spdystream.MirrorStreamHandler(stream)
   199  }
   200
View as plain text