github.com/pf-qiu/concourse/v6@v6.7.3-0.20201207032516-1f455d73275f/fly/commands/internal/hijacker/hijacker_test.go (about)

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