public_git
/
orchard
mirror of https://github.com/cirruslabs/orchard.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

API(port forward endpoint): handle normal WebSocket closure gracefully (#108)

This commit is contained in:
Nikolay Edigaryev 2023-07-21 00:55:42 +04:00 committed by GitHub
parent c5c729c96f
commit a52c205c34
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 47 additions and 28 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

75
internal/controller/api_vms_portforward.go
View File

@ -9,6 +9,7 @@ import (
"github.com/cirruslabs/orchard/rpc"
"github.com/gin-gonic/gin"
"github.com/google/uuid"
"github.com/pkg/errors"
"net/http"
"nhooyr.io/websocket"
"strconv"
@ -41,34 +42,9 @@ func (controller *Controller) portForwardVM(ctx *gin.Context) responder.Responde
defer waitContextCancel()
// Look-up the VM
var vm *v1.VM
for {
if lookupResponder := controller.storeView(func(txn storepkg.Transaction) responder.Responder {
vm, err = txn.GetVM(name)
if err != nil {
return responder.Error(err)
}
return nil
}); lookupResponder != nil {
return lookupResponder
}
if vm.TerminalState() {
return responder.JSON(http.StatusExpectationFailed, NewErrorResponse("VM is in a terminal state '%s'", vm.Status))
}
if vm.Status == v1.VMStatusRunning {
// VM is running, proceed
break
}
select {
case <-waitContext.Done():
return responder.JSON(http.StatusRequestTimeout, NewErrorResponse("VM is not running on '%s' worker", vm.Worker))
case <-time.After(1 * time.Second):
// try again
continue
}
vm, responderImpl := controller.waitForVM(waitContext, name)
if responderImpl != nil {
return responderImpl
}
// Request and wait for a connection with a worker
@ -113,6 +89,14 @@ func (controller *Controller) portForwardVM(ctx *gin.Context) responder.Responde
wsConnAsNetConn := websocket.NetConn(ctx, wsConn, expectedMsgType)
if err := proxy.Connections(wsConnAsNetConn, fromWorkerConnection); err != nil {
var websocketCloseError websocket.CloseError
// Normal closure from the user
if errors.As(err, &websocketCloseError) &&
websocketCloseError.Code == websocket.StatusNormalClosure {
return responder.Empty()
}
controller.logger.Warnf("failed to port-forward: %v", err)
}
@ -121,3 +105,38 @@ func (controller *Controller) portForwardVM(ctx *gin.Context) responder.Responde
return responder.Error(ctx.Err())
}
}
func (controller *Controller) waitForVM(ctx context.Context, name string) (*v1.VM, responder.Responder) {
var vm *v1.VM
var err error
for {
if lookupResponder := controller.storeView(func(txn storepkg.Transaction) responder.Responder {
vm, err = txn.GetVM(name)
if err != nil {
return responder.Error(err)
}
return nil
}); lookupResponder != nil {
return nil, lookupResponder
}
if vm.TerminalState() {
return nil, responder.JSON(http.StatusExpectationFailed,
NewErrorResponse("VM is in a terminal state '%s'", vm.Status))
}
if vm.Status == v1.VMStatusRunning {
// VM is running, proceed
return vm, nil
}
select {
case <-ctx.Done():
return nil, responder.JSON(http.StatusRequestTimeout,
NewErrorResponse("VM is not running on '%s' worker", vm.Worker))
case <-time.After(1 * time.Second):
// try again
continue
}
}
}