Skip to content

Commit

Permalink
handle interrupts
Browse files Browse the repository at this point in the history
  • Loading branch information
niqdev committed Oct 19, 2023
1 parent ce1e463 commit 40bd907
Show file tree
Hide file tree
Showing 11 changed files with 95 additions and 41 deletions.
8 changes: 2 additions & 6 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -302,14 +302,12 @@ Credit should go to all the authors and maintainers for their open source tools,

<!--
fix distroless kube
solve the machine and add how to after docker https://github.com/juice-shop/juice-shop#docker-container
* test all catalog
* discord + social links
* replace task/htb example with thm
* verify/support kube config relative path
* verify/support kube config relative path + remote cluster
* update cloud pkg
* update platform prs
* verify network connectivity between boxes/tasks i.e. kube.svc
Expand All @@ -331,8 +329,7 @@ TODO
- debug `htb-postman`
- lab inputs
- add flow example
- verify kube/cloud distroless support
- verify kube/cloud no-shell support
- verify cloud no-shell support
- play htb: linux/win
- RELEASE example https://github.com/boz/kail#homebrew
- docker release and gh-action
Expand Down Expand Up @@ -377,7 +374,6 @@ TODO
- print/event shared directory, same as envs, ports etc.
- review tty resize
- expose copy from/to ???
- kube: add distroless support
- kube: verify if `close()` is needed or `return nil`
- kube: `execBox` deployment always check/scale replica to 1 before exec (test with replica=0)
- kube: update resources sizes + comparison
Expand Down
22 changes: 7 additions & 15 deletions pkg/box/docker/docker.go
Original file line number Diff line number Diff line change
Expand Up @@ -2,10 +2,6 @@ package docker

import (
"fmt"
"os"
"os/signal"
"syscall"

"github.com/pkg/errors"
"golang.org/x/exp/maps"
"golang.org/x/exp/slices"
Expand All @@ -16,6 +12,7 @@ import (
commonDocker "github.com/hckops/hckctl/pkg/common/docker"
commonModel "github.com/hckops/hckctl/pkg/common/model"
"github.com/hckops/hckctl/pkg/schema"
"github.com/hckops/hckctl/pkg/util"
)

func newDockerBoxClient(commonOpts *boxModel.CommonBoxOptions, dockerOpts *commonModel.DockerOptions) (*DockerBoxClient, error) {
Expand Down Expand Up @@ -189,7 +186,6 @@ func (box *DockerBoxClient) searchBox(name string) (*boxModel.BoxInfo, error) {
}

func (box *DockerBoxClient) execBox(template *boxModel.BoxV1, info boxModel.BoxInfo, streamOpts *commonModel.StreamOptions, deleteOnExit bool) error {
box.eventBus.Publish(newContainerExecDockerEvent(info.Id, info.Name, template.Shell))

// attempt to restart all associated sidecars
sidecars, err := box.dockerCommon.SidecarList(info.Name)
Expand Down Expand Up @@ -283,7 +279,7 @@ func (box *DockerBoxClient) execBox(template *boxModel.BoxV1, info boxModel.BoxI
}
},
}

box.eventBus.Publish(newContainerExecDockerEvent(info.Id, info.Name, template.Shell))
return box.client.ContainerExec(execOpts)
}

Expand All @@ -301,26 +297,22 @@ func (box *DockerBoxClient) publishPortInfo(networkMap map[string]boxModel.BoxPo
func (box *DockerBoxClient) logsBox(info boxModel.BoxInfo, streamOpts *commonModel.StreamOptions, deleteOnExit bool) error {

if deleteOnExit {
// captures CTRL+C
signalChan := make(chan os.Signal, 1)
signal.Notify(signalChan, os.Interrupt, syscall.SIGTERM)
go func() {
<-signalChan
util.InterruptHandler(func() {
box.deleteBox(info)
}()
})
}

opts := &docker.ContainerLogsOpts{
ContainerId: info.Id,
OutStream: streamOpts.Out,
ErrStream: streamOpts.Err,
OnStreamCloseCallback: func() {
box.eventBus.Publish(newContainerExecExitDockerEvent(info.Id))
box.eventBus.Publish(newContainerLogsExitDockerEvent(info.Id))
},
OnStreamErrorCallback: func(err error) {
box.eventBus.Publish(newContainerExecErrorDockerEvent(info.Id, err))
box.eventBus.Publish(newContainerLogsErrorDockerEvent(info.Id, err))
},
}
box.eventBus.Publish(newContainerLogsDockerEvent(info.Id))
return box.client.ContainerLogs(opts)
}

Expand Down
12 changes: 12 additions & 0 deletions pkg/box/docker/events.go
Original file line number Diff line number Diff line change
Expand Up @@ -84,6 +84,18 @@ func newContainerExecErrorDockerEvent(containerId string, err error) *dockerBoxE
return &dockerBoxEvent{kind: event.LogError, value: fmt.Sprintf("container exec error: containerId=%s error=%v", containerId, err)}
}

func newContainerLogsDockerEvent(containerId string) *dockerBoxEvent {
return &dockerBoxEvent{kind: event.LogDebug, value: fmt.Sprintf("container logs: containerId=%s", containerId)}
}

func newContainerLogsExitDockerEvent(containerId string) *dockerBoxEvent {
return &dockerBoxEvent{kind: event.LogDebug, value: fmt.Sprintf("container logs exit: containerId=%s", containerId)}
}

func newContainerLogsErrorDockerEvent(containerId string, err error) *dockerBoxEvent {
return &dockerBoxEvent{kind: event.LogError, value: fmt.Sprintf("container logs error: containerId=%s error=%v", containerId, err)}
}

func newContainerListDockerEvent(index int, containerName string, containerId string, healthy bool) *dockerBoxEvent {
return &dockerBoxEvent{kind: event.LogDebug, value: fmt.Sprintf("container list: (%d) containerName=%s containerId=%s healthy=%v", index, containerName, containerId, healthy)}
}
Expand Down
16 changes: 16 additions & 0 deletions pkg/box/kubernetes/events.go
Original file line number Diff line number Diff line change
Expand Up @@ -92,6 +92,22 @@ func newPodExecKubeLoaderEvent() *kubeBoxEvent {
return &kubeBoxEvent{kind: event.LoaderStop, value: "waiting"}
}

func newPodLogsKubeEvent(namespace string, name string) *kubeBoxEvent {
return &kubeBoxEvent{kind: event.LogDebug, value: fmt.Sprintf("pod logs: namespace=%s name=%s", namespace, name)}
}

func newPodLogsExitKubeEvent(namespace string, name string) *kubeBoxEvent {
return &kubeBoxEvent{kind: event.LogDebug, value: fmt.Sprintf("pod logs exit: namespace=%s name=%s", namespace, name)}
}

func newPodLogsExitKubeConsoleEvent() *kubeBoxEvent {
return &kubeBoxEvent{kind: event.LoaderUpdate, value: "killing"}
}

func newPodLogsErrorKubeEvent(namespace string, name string, err error) *kubeBoxEvent {
return &kubeBoxEvent{kind: event.LogError, value: fmt.Sprintf("pod logs error: namespace=%s name=%s error=%v", namespace, name, err)}
}

func newPodPortForwardIgnoreKubeEvent(namespace string, name string) *kubeBoxEvent {
return &kubeBoxEvent{kind: event.LogWarning, value: fmt.Sprintf("pod port-forward ignored: namespace=%s name=%s", namespace, name)}
}
Expand Down
31 changes: 30 additions & 1 deletion pkg/box/kubernetes/kubernetes.go
Original file line number Diff line number Diff line change
Expand Up @@ -204,7 +204,15 @@ func (box *KubeBoxClient) searchBox(name string) (*boxModel.BoxInfo, error) {
}

func (box *KubeBoxClient) execBox(template *boxModel.BoxV1, info *boxModel.BoxInfo, streamOpts *commonModel.StreamOptions, deleteOnExit bool) error {
box.eventBus.Publish(newPodExecKubeEvent(template.Name, box.clientOpts.Namespace, info.Id, template.Shell))

// TODO if BoxInfo not Healthy attempt scale 1

if template.Shell == boxModel.BoxShellNone {
// stop loader
box.eventBus.Publish(newPodExecKubeLoaderEvent())

return box.logsBox(template, info, deleteOnExit)
}

// exec
opts := &kubernetes.PodExecOpts{
Expand All @@ -226,9 +234,30 @@ func (box *KubeBoxClient) execBox(template *boxModel.BoxV1, info *boxModel.BoxIn
defer box.deleteBox(info.Name)
}

box.eventBus.Publish(newPodExecKubeEvent(template.Name, box.clientOpts.Namespace, info.Id, template.Shell))
return box.client.PodExecShell(opts)
}

func (box *KubeBoxClient) logsBox(template *boxModel.BoxV1, info *boxModel.BoxInfo, deleteOnExit bool) error {
namespace := box.clientOpts.Namespace

if deleteOnExit {
util.InterruptHandler(func() {
box.eventBus.Publish(newPodLogsExitKubeEvent(namespace, info.Id))
box.eventBus.Publish(newPodLogsExitKubeConsoleEvent())
box.deleteBox(info.Name)
})
}

opts := &kubernetes.PodLogsOpts{
Namespace: namespace,
PodName: info.Id,
ContainerName: template.MainContainerName(),
}
box.eventBus.Publish(newPodLogsKubeEvent(namespace, info.Id))
return box.client.PodLogsStd(opts)
}

func (box *KubeBoxClient) podPortForward(template *boxModel.BoxV1, boxInfo *boxModel.BoxInfo, isWait bool) error {
namespace := box.clientOpts.Namespace

Expand Down
13 changes: 4 additions & 9 deletions pkg/client/docker/client.go
Original file line number Diff line number Diff line change
Expand Up @@ -5,11 +5,9 @@ import (
"fmt"
"io"
"os"
"os/signal"
"path/filepath"
"strings"
"sync"
"syscall"
"time"

"github.com/pkg/errors"
Expand Down Expand Up @@ -99,13 +97,9 @@ func (client *DockerClient) ContainerCreate(opts *ContainerCreateOpts) (string,
}

if opts.CaptureInterrupt {
// captures CTRL+C
signalChan := make(chan os.Signal, 1)
signal.Notify(signalChan, os.Interrupt, syscall.SIGTERM)
go func() {
<-signalChan
util.InterruptHandler(func() {
opts.OnContainerInterruptCallback(newContainer.ID)
}()
})
}

if err := opts.OnContainerCreateCallback(newContainer.ID); err != nil {
Expand Down Expand Up @@ -355,8 +349,9 @@ func (client *DockerClient) ContainerLogs(opts *ContainerLogsOpts) error {

var once sync.Once
go func() {
// do not copy ErrStream because it fails with "Unrecognized input header"
if _, err := io.Copy(opts.OutStream, outStream); err != nil {
opts.OnStreamErrorCallback(errors.Wrap(err, "error copy stdout and stderr docker->local"))
opts.OnStreamErrorCallback(errors.Wrap(err, "error copy stdout docker->local"))
}
once.Do(onStreamCloseCallback)
}()
Expand Down
1 change: 0 additions & 1 deletion pkg/client/docker/options.go
Original file line number Diff line number Diff line change
Expand Up @@ -75,7 +75,6 @@ type ContainerExecOpts struct {
type ContainerLogsOpts struct {
ContainerId string
OutStream io.Writer
ErrStream io.Writer
OnStreamCloseCallback func()
OnStreamErrorCallback func(error)
}
10 changes: 2 additions & 8 deletions pkg/client/kubernetes/client.go
Original file line number Diff line number Diff line change
Expand Up @@ -7,11 +7,9 @@ import (
"io"
"net/http"
"os"
"os/signal"
"path/filepath"
"strconv"
"strings"
"syscall"

"github.com/pkg/errors"

Expand Down Expand Up @@ -582,13 +580,9 @@ func (client *KubeClient) JobCreate(opts *JobCreateOpts) error {
}

if opts.CaptureInterrupt {
// captures CTRL+C
signalChan := make(chan os.Signal, 1)
signal.Notify(signalChan, os.Interrupt, syscall.SIGTERM)
go func() {
<-signalChan
util.InterruptHandler(func() {
opts.OnContainerInterruptCallback(job.Name)
}()
})
}

// blocks until the job is ready, then stop watching
Expand Down
4 changes: 4 additions & 0 deletions pkg/task/kubernetes/events.go
Original file line number Diff line number Diff line change
Expand Up @@ -40,6 +40,10 @@ func newJobDeleteKubeEvent(namespace string, name string) *kubeTaskEvent {
return &kubeTaskEvent{kind: event.LogInfo, value: fmt.Sprintf("job delete: namespace=%s name=%s", namespace, name)}
}

func newJobDeleteKubeConsoleEvent() *kubeTaskEvent {
return &kubeTaskEvent{kind: event.LoaderUpdate, value: "killing"}
}

func newPodNameKubeEvent(namespace string, name string, containerName string) *kubeTaskEvent {
return &kubeTaskEvent{kind: event.LogInfo, value: fmt.Sprintf("found unique pod: namespace=%s name=%s containerName=%s", namespace, name, containerName)}
}
Expand Down
2 changes: 1 addition & 1 deletion pkg/task/kubernetes/kubernetes.go
Original file line number Diff line number Diff line change
Expand Up @@ -91,7 +91,7 @@ func (task *KubeTaskClient) runTask(opts *taskModel.RunOptions) error {
CaptureInterrupt: true,
OnContainerInterruptCallback: func(name string) {
// ignore error when interrupted: it will attempt to delete the job twice
task.eventBus.Publish(newJobDeleteKubeEvent(namespace, jobName))
task.eventBus.Publish(newJobDeleteKubeConsoleEvent())
task.client.JobDelete(namespace, name)
},
OnStatusEventCallback: func(event string) {
Expand Down
17 changes: 17 additions & 0 deletions pkg/util/interrupt.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,17 @@
package util

import (
"os"
"os/signal"
"syscall"
)

func InterruptHandler(callback func()) {
// captures CTRL+C
signalChan := make(chan os.Signal, 1)
signal.Notify(signalChan, os.Interrupt, syscall.SIGTERM)
go func() {
<-signalChan
callback()
}()
}

0 comments on commit 40bd907

Please sign in to comment.