github.com/chenbh/concourse/v6@v6.4.2/fly/commands/internal/hijacker/hijacker_test.go (about)

     1  package hijacker_test
     2  
     3  import (
     4  	"crypto/tls"
     5  	"fmt"
     6  	"net/http"
     7  	"time"
     8  
     9  	. "github.com/onsi/ginkgo"
    10  	. "github.com/onsi/gomega"
    11  	"github.com/tedsuo/rata"
    12  
    13  	"github.com/gorilla/websocket"
    14  	"github.com/onsi/gomega/gbytes"
    15  	"github.com/onsi/gomega/ghttp"
    16  
    17  	"github.com/chenbh/concourse/v6/atc"
    18  	"github.com/chenbh/concourse/v6/fly/commands/internal/hijacker"
    19  )
    20  
    21  var _ = Describe("Hijacker", func() {
    22  	// Other functionality tested through the hijack command integration test.
    23  
    24  	upgrader := websocket.Upgrader{}
    25  
    26  	wasPingedHandler := func(id string, didHijack chan<- struct{}, didGetPinged chan<- struct{}) http.HandlerFunc {
    27  		return ghttp.CombineHandlers(
    28  			ghttp.VerifyRequest("GET", fmt.Sprintf("/api/v1/teams/some-team/containers/%s/hijack", id)),
    29  			func(w http.ResponseWriter, r *http.Request) {
    30  				defer GinkgoRecover()
    31  
    32  				close(didHijack)
    33  				conn, err := upgrader.Upgrade(w, r, nil)
    34  				Expect(err).NotTo(HaveOccurred())
    35  
    36  				defer conn.Close()
    37  
    38  				conn.SetPingHandler(func(data string) error {
    39  					conn.WriteControl(websocket.CloseMessage, []byte{}, time.Now().Add(time.Second))
    40  					close(didGetPinged)
    41  					return nil
    42  				})
    43  
    44  				for {
    45  					_, _, err := conn.ReadMessage()
    46  					if err != nil {
    47  						break
    48  					}
    49  				}
    50  
    51  			},
    52  		)
    53  	}
    54  
    55  	var (
    56  		server *ghttp.Server
    57  
    58  		didHijack  chan struct{}
    59  		didGetPing chan struct{}
    60  	)
    61  
    62  	BeforeEach(func() {
    63  		didHijack = make(chan struct{})
    64  		didGetPing = make(chan struct{})
    65  
    66  		server = ghttp.NewServer()
    67  	})
    68  
    69  	AfterEach(func() {
    70  		server.Close()
    71  	})
    72  
    73  	Describe("keeping the connection alive", func() {
    74  		BeforeEach(func() {
    75  			server.AppendHandlers(wasPingedHandler("hello", didHijack, didGetPing))
    76  		})
    77  
    78  		It("sends the occasional ping", func() {
    79  			tlsConfig := &tls.Config{
    80  				InsecureSkipVerify: true,
    81  			}
    82  
    83  			reqGenerator := rata.NewRequestGenerator(server.URL(), atc.Routes)
    84  
    85  			stdin := gbytes.NewBuffer()
    86  			stdout := gbytes.NewBuffer()
    87  			stderr := gbytes.NewBuffer()
    88  
    89  			h := hijacker.New(tlsConfig, reqGenerator, nil)
    90  			_, err := h.Hijack("some-team", "hello", atc.HijackProcessSpec{
    91  				Path: "/bin/echo",
    92  				Args: []string{"hello", "world"},
    93  			}, hijacker.ProcessIO{
    94  				In:  stdin,
    95  				Out: stdout,
    96  				Err: stderr,
    97  			})
    98  
    99  			h.SetHeartbeatInterval(100 * time.Millisecond)
   100  
   101  			Expect(err).ShouldNot(HaveOccurred())
   102  			Expect(didHijack).To(BeClosed())
   103  			Eventually(didGetPing).Should(BeClosed())
   104  		})
   105  	})
   106  })