Simplify exec on ACI

* both streams send messages to a `chan error`
* the main goroutine returns the error if one exists
This commit is contained in:
Djordje Lukic 2020-05-04 20:38:10 +02:00
parent 900d82ced0
commit d6417cb504
1 changed files with 22 additions and 23 deletions

View File

@ -8,15 +8,15 @@ import (
"net/http" "net/http"
"os" "os"
"github.com/docker/api/context/store"
"github.com/gobwas/ws"
"github.com/gobwas/ws/wsutil"
"github.com/Azure/azure-sdk-for-go/services/containerinstance/mgmt/2018-10-01/containerinstance" "github.com/Azure/azure-sdk-for-go/services/containerinstance/mgmt/2018-10-01/containerinstance"
"github.com/Azure/azure-sdk-for-go/services/keyvault/auth" "github.com/Azure/azure-sdk-for-go/services/keyvault/auth"
"github.com/Azure/go-autorest/autorest" "github.com/Azure/go-autorest/autorest"
"github.com/Azure/go-autorest/autorest/to" "github.com/Azure/go-autorest/autorest/to"
"github.com/gobwas/ws"
"github.com/gobwas/ws/wsutil"
"github.com/pkg/errors"
"github.com/docker/api/context/store"
tm "github.com/buger/goterm" tm "github.com/buger/goterm"
) )
@ -166,33 +166,33 @@ func execCommands(ctx context.Context, address string, password string, commands
} }
func exec(ctx context.Context, address string, password string, reader io.Reader, writer io.Writer) error { func exec(ctx context.Context, address string, password string, reader io.Reader, writer io.Writer) error {
ctx, cancel := context.WithCancel(ctx)
conn, _, _, err := ws.DefaultDialer.Dial(ctx, address) conn, _, _, err := ws.DefaultDialer.Dial(ctx, address)
if err != nil { if err != nil {
cancel()
return err return err
} }
err = wsutil.WriteClientMessage(conn, ws.OpText, []byte(password)) err = wsutil.WriteClientMessage(conn, ws.OpText, []byte(password))
if err != nil { if err != nil {
cancel()
return err return err
} }
done := make(chan struct{}) downstreamChannel := make(chan error, 10)
upstreamChannel := make(chan error, 10)
go func() { go func() {
defer close(done)
for { for {
msg, _, err := wsutil.ReadServerData(conn) msg, _, err := wsutil.ReadServerData(conn)
if err != nil { if err != nil {
if err == io.EOF {
downstreamChannel <- nil
return
}
downstreamChannel <- err
return return
} }
fmt.Fprint(writer, string(msg)) fmt.Fprint(writer, string(msg))
} }
}() }()
readChannel := make(chan []byte, 10)
go func() { go func() {
for { for {
// We send each byte, byte-per-byte over the // We send each byte, byte-per-byte over the
@ -200,26 +200,25 @@ func exec(ctx context.Context, address string, password string, reader io.Reader
buffer := make([]byte, 1) buffer := make([]byte, 1)
n, err := reader.Read(buffer) n, err := reader.Read(buffer)
if err != nil { if err != nil {
close(done) upstreamChannel <- err
cancel() return
break
} }
if n > 0 { if n > 0 {
readChannel <- buffer err := wsutil.WriteClientMessage(conn, ws.OpText, buffer)
if err != nil {
upstreamChannel <- err
}
} }
} }
}() }()
for { for {
select { select {
case <-done: case err := <-downstreamChannel:
return nil return errors.Wrap(err, "failed to read input from container")
case bytes := <-readChannel: case err := <-upstreamChannel:
err := wsutil.WriteClientMessage(conn, ws.OpText, bytes) return errors.Wrap(err, "failed to send input to container")
if err != nil {
return err
}
} }
} }
} }