Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

DXCDT-550: Propagate errors through websocket connection #879

Merged
Merged
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
51 changes: 42 additions & 9 deletions internal/cli/universal_login_customize.go
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,7 @@ const (
loadBrandingMessageType = "LOAD_BRANDING"
fetchPromptMessageType = "FETCH_PROMPT"
saveBrandingMessageType = "SAVE_BRANDING"
errorMessageType = "ERROR"
)

type (
Expand Down Expand Up @@ -342,16 +343,21 @@ func (h *webSocketHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {

connection, err := upgrader.Upgrade(w, r, nil)
if err != nil {
h.display.Errorf("failed to upgrade the connection to the WebSocket protocol: %v", err)
h.display.Errorf("Failed to upgrade the connection to the WebSocket protocol: %v", err)
h.display.Warnf("Try restarting the command.")
h.shutdown()
return
}
defer func() {
_ = connection.Close()
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We weren't properly closing the connection before. This is really important to do.

}()

connection.SetReadLimit(1e+6) // 1 MB.

payload, err := json.Marshal(&h.brandingData)
if err != nil {
h.display.Errorf("failed to encode the branding data to json: %v", err)
h.display.Errorf("Failed to encode the branding data to json: %v", err)
h.display.Warnf("Try restarting the command.")
h.shutdown()
return
}
Expand All @@ -362,23 +368,30 @@ func (h *webSocketHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
}

if err = connection.WriteJSON(loadBrandingMsg); err != nil {
h.display.Errorf("failed to send branding data message: %v", err)
h.display.Errorf("Failed to send branding data message: %v", err)
h.display.Warnf("Try restarting the command.")
h.shutdown()
return
}

for {
var message webSocketMessage
if err := connection.ReadJSON(&message); err != nil {
h.display.Errorf("failed to read WebSocket message: %v", err)
if websocket.IsCloseError(err, websocket.CloseNormalClosure, websocket.CloseGoingAway) ||
websocket.IsUnexpectedCloseError(err, websocket.CloseAbnormalClosure) {
// The connection was closed.
break
}
Comment on lines +380 to +384
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This logic is needed in order to prevent a panic where we keep trying to read from a closed connection due to the for loop.


h.display.Errorf("Failed to read WebSocket message: %v", err)
continue
}

switch message.Type {
case fetchPromptMessageType:
var promptToFetch promptData
if err := json.Unmarshal(message.Payload, &promptToFetch); err != nil {
h.display.Errorf("failed to unmarshal fetch prompt payload: %v", err)
h.display.Errorf("Failed to unmarshal fetch prompt payload: %v", err)
continue
}

Expand All @@ -389,7 +402,17 @@ func (h *webSocketHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
promptToFetch.Language,
)
if err != nil {
h.display.Errorf("failed to fetch custom text for prompt: %v", err)
h.display.Errorf("Failed to fetch custom text for prompt: %v", err)

errorMsg := webSocketMessage{
Type: errorMessageType,
Payload: []byte(fmt.Sprintf(`{"error":%q}`, err.Error())),
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Should we be using a struct to represent the error, and then JSON-encode it?

Copy link
Contributor Author

@sergiught sergiught Oct 16, 2023

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yep, we should. I plan to do that in a follow up PR and refactor the entire webSocketMessage so we can add custom JSON marshalling and unmarshalling and reuse the payload accordingly. This will prevent us from needing to unmarshal twice the message when we received it through the websocket as we're doing now.

The solution would be similar to what we do within the Go SDK for Connection Options.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Here it is @Widcket -> #880

}

if err := connection.WriteJSON(errorMsg); err != nil {
h.display.Errorf("Failed to send error message: %v", err)
}

continue
}

Expand All @@ -406,18 +429,28 @@ func (h *webSocketHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
}

if err = connection.WriteJSON(fetchPromptMsg); err != nil {
h.display.Errorf("failed to send prompt data message: %v", err)
h.display.Errorf("Failed to send prompt data message: %v", err)
continue
}
case saveBrandingMessageType:
var saveBrandingMsg universalLoginBrandingData
if err := json.Unmarshal(message.Payload, &saveBrandingMsg); err != nil {
h.display.Errorf("failed to unmarshal save branding data payload: %v", err)
h.display.Errorf("Failed to unmarshal save branding data payload: %v", err)
continue
}

if err := saveUniversalLoginBrandingData(r.Context(), h.api, &saveBrandingMsg); err != nil {
h.display.Errorf("failed to save branding data: %v", err)
h.display.Errorf("Failed to save branding data: %v", err)

errorMsg := webSocketMessage{
Type: errorMessageType,
Payload: []byte(fmt.Sprintf(`{"error":%q}`, err.Error())),
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Same as above.

}

if err := connection.WriteJSON(errorMsg); err != nil {
h.display.Errorf("Failed to send error message: %v", err)
}

continue
}
}
Expand Down