How to know if Go WebSocket has closed and then exit from this function?

Question

I have written a Go function to attach to a Docker container. I want to access this bash shell remotely, so I am forwarding the output stream to the web socket and whatever I receive from the web socket I send it to stdin for the Docker container. I have infinite loops inside the function to continuously perform these read/write operations.

This works fine. But how do I exit from the infinite loops when the web socket is closed? In .NET we have something called CancellationToken and in every iteration of the loop, we check if IsCancellationRequested == true and stop executing. I am using Go for the first time. Is there a similar thing in Go to know if the web socket is closed and how to identify this in the infinite loops?

func handleDockerShell(c *websocket.Conn, containerId string) {
    cli, _ := client.NewClientWithOpts(client.FromEnv)
    execConfig := types.ExecConfig{
        Tty: true,
        Detach: false,
        AttachStdin: true,
        AttachStderr: true,
        AttachStdout: true,
        Cmd: []string { "bash" },
    }

    idResponse, err := cli.ContainerExecCreate(context.Background(), containerId, execConfig)
    if(err != nil) {
        log.Println(err)
        return
    }

    hijackedResponse, err := cli.ContainerExecAttach(context.Background(), idResponse.ID, types.ExecStartCheck{Tty: true})
    if(err != nil) {
        log.Println(err)
        return
    }

    go func () {
        for {
            _, inputBytes, _  := c.ReadMessage()
            hijackedResponse.Conn.Write(inputBytes)
        }
    }()

    go func ()  {
        for {
            outputBytesBuffer := make([]byte, 1024)
            n, _ := hijackedResponse.Conn.Read(outputBytesBuffer)
            c.WriteMessage(websocket.BinaryMessage, outputBytesBuffer[0:n])
        }
    }()

    for {
        time.Sleep(1 * time.Second)
    }
}

Answer 1

The websocket connections methods return an error when the connection failed. The application should exit the loop and close the connection.

go func () {
    defer c.Close()
    for {
        _, inputBytes, err  := c.ReadMessage()
        if err != nil { return }
        hijackedResponse.Conn.Write(inputBytes)
    }
}()

go func ()  {
    defer c.Close()
    for {
        outputBytesBuffer := make([]byte, 1024)
        n, _ := hijackedResponse.Conn.Read(outputBytesBuffer)
        err := c.WriteMessage(websocket.BinaryMessage, outputBytesBuffer[0:n])
        if err != nil { return }
    }
}()

Delete the for loop with sleep at the end of the function.

I answered the question regarding the websocket connection. A similar change should be made to handle errors on the hijacked connection.