hacking

saucelabs · Nov 7, 2023 · c61a8da · c61a8da
1 parent eeb7a02
commit c61a8da
Show file tree

Hide file tree

Showing 2 changed files with 54 additions and 0 deletions.
diff --git a/internal/http/imagerunner.go b/internal/http/imagerunner.go
@@ -6,10 +6,12 @@ import (
 	"encoding/json"
 	"fmt"
 	"io"
+	"log"
 	"net/http"
 	"strings"
 	"time"
 
+	"github.com/gorilla/websocket"
 	"github.com/hashicorp/go-retryablehttp"
 	"github.com/saucelabs/saucectl/internal/iam"
 	"github.com/saucelabs/saucectl/internal/imagerunner"
@@ -197,6 +199,29 @@ func (c *ImageRunner) GetLogs(ctx context.Context, id string) (string, error) {
 	return c.doGetStr(ctx, urlResponse.URL)
 }
 
+func (c *ImageRunner) getWebsocketUrl() string {
+	wsUrl := c.URL
+	wsUrl = strings.Replace(wsUrl, "https://", "wss://", 1)
+	wsUrl = strings.Replace(wsUrl, "http://", "ws://", 1)
+	return wsUrl
+}
+
+func (c *ImageRunner) OpenAsyncEventsWebsocket(ctx context.Context, id string) (*websocket.Conn, error) {
+	// dummy request so that we build basic auth header consistently
+	req := http.Request{}
+	req.SetBasicAuth(c.Creds.Username, c.Creds.AccessKey)
+
+	url := fmt.Sprintf("%s/-/livelogs?r=%s", c.getWebsocketUrl(), id)
+
+	header := http.Header{}
+	header.Add("Authorization", req.Header.Get("Authorization"))
+	ws, _, err := websocket.DefaultDialer.Dial(url, header)
+	if err != nil {
+		log.Fatal("dial:", err)
+	}
+	return ws, nil
+}
+
 func (c *ImageRunner) doGetStr(ctx context.Context, url string) (string, error) {
 	urlReq, err := NewRetryableRequestWithContext(ctx, http.MethodGet, url, nil)
 	if err != nil {

diff --git a/internal/saucecloud/imagerunner.go b/internal/saucecloud/imagerunner.go
@@ -13,6 +13,7 @@ import (
 	"reflect"
 	"time"
 
+	"github.com/gorilla/websocket"
 	"github.com/rs/zerolog/log"
 	"github.com/ryanuber/go-glob"
 	szip "github.com/saucelabs/saucectl/internal/archive/zip"
@@ -30,6 +31,7 @@ type ImageRunner interface {
 	StopRun(ctx context.Context, id string) error
 	DownloadArtifacts(ctx context.Context, id string) (io.ReadCloser, error)
 	GetLogs(ctx context.Context, id string) (string, error)
+	OpenAsyncEventsWebsocket(ctx context.Context, id string) (*websocket.Conn, error)
 }
 
 type SuiteTimeoutError struct {
@@ -269,6 +271,12 @@ func (r *ImgRunner) runSuite(suite imagerunner.Suite) (imagerunner.Runner, error
 		return runner, nil
 	}
 
+	go func() {
+		err := r.HandleAsyncEvents(ctx, runner.ID)
+		// TODO: handle error better
+		log.Err(err).Msg("Async event handler failed.")
+	}()
+
 	var run imagerunner.Runner
 	run, err = r.PollRun(ctx, runner.ID, runner.Status)
 	if errors.Is(err, context.DeadlineExceeded) && ctx.Err() != nil {
@@ -399,6 +407,27 @@ func (r *ImgRunner) PollRun(ctx context.Context, id string, lastStatus string) (
 	}
 }
 
+func (r *ImgRunner) HandleAsyncEvents(ctx context.Context, id string) error {
+	conn, err := r.RunnerService.OpenAsyncEventsWebsocket(ctx, id)
+	if err != nil {
+		return err
+	}
+	defer conn.Close()
+
+	for {
+		select {
+		case <-ctx.Done():
+			return ctx.Err()
+		default:
+			_, msg, err := conn.ReadMessage()
+			if err != nil {
+				return err
+			}
+			log.Info().Msgf("Received async event: %s", msg)
+		}
+	}
+}
+
 // DownloadArtifact downloads a zipped archive of artifacts and extracts the required files.
 func (r *ImgRunner) DownloadArtifacts(runnerID, suiteName, status string, passed bool) []string {
 	if r.Async ||