Skip to content

Commit

Permalink
Add graceful shutdown support for gvprovy on Win
Browse files Browse the repository at this point in the history 
Leverages WM_CLOSE and WM_QUIT events on win32 for SIGTERM like behavior
The improved/updated logic is now implemented in a shared containers/winquit module
Also, refactors win-sshproxy.exe to be consistent (had it's own precursor impl)

Signed-off-by: Jason T. Greene <[email protected]>
  • Loading branch information
@n1hility
n1hility committed Sep 14, 2023
1 parent 9ebd937 commit 2d105d3
Show file tree
Hide file tree
Showing 33 changed files with 2,413 additions and 155 deletions.
4 changes: 4 additions & 0 deletions cmd/gvproxy/main.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,7 @@ import (
"github.com/containers/gvisor-tap-vsock/pkg/transport"
"github.com/containers/gvisor-tap-vsock/pkg/types"
"github.com/containers/gvisor-tap-vsock/pkg/virtualnetwork"
"github.com/containers/winquit/pkg/winquit"
"github.com/dustin/go-humanize"
"github.com/pkg/errors"
log "github.com/sirupsen/logrus"
Expand Down Expand Up @@ -84,6 +85,9 @@ func main() {
log.SetLevel(log.DebugLevel)
}

// Intercept WM_QUIT/WM_CLOSE events if on Windows as SIGTERM (noop on other OSs)
winquit.SimulateSigTermOnQuit(sigChan)

// Make sure the qemu socket provided is valid syntax
if len(qemuSocket) > 0 {
uri, err := url.Parse(qemuSocket)
Expand Down
68 changes: 15 additions & 53 deletions cmd/win-sshproxy/main.go
View file
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
//go:build windows
// +build windows

package main
Expand All @@ -8,15 +9,14 @@ import (
"net/url"
"os"
"path/filepath"
"runtime"
"strings"
"syscall"
"unsafe"

"github.com/containers/gvisor-tap-vsock/pkg/sshclient"
"github.com/containers/winquit/pkg/winquit"
"github.com/sirupsen/logrus"
"golang.org/x/sync/errgroup"
"golang.org/x/sys/windows"
"golang.org/x/sys/windows/svc/eventlog"
)

Expand All @@ -25,15 +25,6 @@ const (
WM_QUIT = 0x12
)

type MSG struct {
hwnd uintptr
message uint32
wParam uintptr
lParam uintptr
time uint32
pt struct{ X, Y int32 }
}

var (
stateDir string
debug bool
Expand Down Expand Up @@ -73,11 +64,18 @@ func main() {
ctx, cancel := context.WithCancel(context.Background())
group, ctx := errgroup.WithContext(ctx)

quit := make(chan bool, 1)
// Wait for a WM_QUIT message to exit
group.Go(func() error {
logrus.Debug("Starting message loop")
return messageLoop(ctx, group, cancel)
})
winquit.NotifyOnQuit(quit)
go func() {
<-quit
cancel()
}()

// Save thread for legacy callers which use it to post a quit
if _, err := saveThreadId(); err != nil {
logrus.Errorf("Error saving thread id: " + err.Error())
}

logrus.Debug("Setting up proxies")
setupProxies(ctx, group, sources, dests, identities)
Expand Down Expand Up @@ -108,36 +106,6 @@ func setupLogging(name string) (*eventlog.Log, error) {
return log, nil
}

func messageLoop(ctx context.Context, group *errgroup.Group, cancel func()) error {
user32 := syscall.NewLazyDLL("user32.dll")
getMessage := user32.NewProc("GetMessageW")

runtime.LockOSThread() // GetMessageW relies on thread state
defer runtime.UnlockOSThread()
tid, err := saveThreadId()
if err != nil {
return err
}

// Abort the message loop thread on cancellation
group.Go(func() error {
<-ctx.Done()
terminate(tid)
return nil
})

for {
var msg = &MSG{}
ret, _, _ := getMessage.Call(uintptr(unsafe.Pointer(msg)), 0, 0, 0, 1)
if ret == 0 || int(ret) == -1 {
logrus.Info("Received QUIT notification")
cancel()
return nil
}
logrus.Infof("Unhandled message: %d", msg.message)
}
}

func setupProxies(ctx context.Context, g *errgroup.Group, sources []string, dests []string, identities []string) error {
for i := 0; i < len(sources); i++ {
var (
Expand Down Expand Up @@ -199,17 +167,11 @@ func saveThreadId() (uint32, error) {
return 0, err
}
defer file.Close()
tid := windows.GetCurrentThreadId()
tid := winquit.GetCurrentMessageLoopThreadId()
fmt.Fprintf(file, "%d:%d\n", os.Getpid(), tid)
return tid, nil
}

func terminate(tid uint32) {
user32 := syscall.NewLazyDLL("user32.dll")
postMessage := user32.NewProc("PostThreadMessageW")
postMessage.Call(uintptr(tid), WM_QUIT, 0, 0)
}

// Creates an "error" style pop-up window
func alert(caption string) int {
// Error box style
Expand All @@ -225,4 +187,4 @@ func alert(caption string) int {
uintptr(format))

return int(ret)
}
}
7 changes: 4 additions & 3 deletions go.mod
View file
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@ go 1.20
require (
github.com/Microsoft/go-winio v0.6.1
github.com/apparentlymart/go-cidr v1.1.0
github.com/containers/winquit v1.1.0
github.com/coreos/stream-metadata-go v0.4.3
github.com/dustin/go-humanize v1.0.1
github.com/google/gopacket v1.1.19
Expand Down Expand Up @@ -38,11 +39,11 @@ require (
github.com/pmezard/go-difflib v1.0.0 // indirect
github.com/u-root/uio v0.0.0-20210528114334-82958018845c // indirect
github.com/vishvananda/netns v0.0.0-20211101163701-50045581ed74 // indirect
golang.org/x/mod v0.10.0 // indirect
golang.org/x/net v0.12.0 // indirect
golang.org/x/mod v0.12.0 // indirect
golang.org/x/net v0.14.0 // indirect
golang.org/x/text v0.13.0 // indirect
golang.org/x/time v0.0.0-20220210224613-90d013bbcef8 // indirect
golang.org/x/tools v0.9.3 // indirect
golang.org/x/tools v0.12.0 // indirect
gopkg.in/tomb.v1 v1.0.0-20141024135613-dd632973f1e7 // indirect
gopkg.in/yaml.v3 v3.0.1 // indirect
)
16 changes: 9 additions & 7 deletions go.sum
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,8 @@ github.com/Microsoft/go-winio v0.6.1/go.mod h1:LRdKpFKfdobln8UmuiYcKPot9D2v6svN5
github.com/apparentlymart/go-cidr v1.1.0 h1:2mAhrMoF+nhXqxTzSZMUzDHkLjmIHC+Zzn4tdgBZjnU=
github.com/apparentlymart/go-cidr v1.1.0/go.mod h1:EBcsNrHc3zQeuaeCeCtQruQm+n9/YjEn/vI25Lg7Gwc=
github.com/armon/go-proxyproto v0.0.0-20210323213023-7e956b284f0a/go.mod h1:QmP9hvJ91BbJmGVGSbutW19IC0Q9phDCLGaomwTJbgU=
github.com/containers/winquit v1.1.0 h1:jArun04BNDQvt2W0Y78kh9TazN2EIEMG5Im6/JY7+pE=
github.com/containers/winquit v1.1.0/go.mod h1:PsPeZlnbkmGGIToMPHF1zhWjBUkd8aHjMOr/vFcPxw8=
github.com/coreos/stream-metadata-go v0.4.3 h1:5GykJ8dtZSx1rdlzEAiDVzA73cwmUF3ceTuIP293L6E=
github.com/coreos/stream-metadata-go v0.4.3/go.mod h1:fMObQqQm8Ku91G04btKzEH3AsdP1mrAb986z9aaK0tE=
github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
Expand Down Expand Up @@ -71,7 +73,7 @@ github.com/onsi/ginkgo v1.6.0/go.mod h1:lLunBs/Ym6LB5Z9jYTR76FiuTmxDTDusOGeTQH+W
github.com/onsi/ginkgo v1.12.1/go.mod h1:zj2OWP4+oCPe1qIXoGWkgMRwljMUYCdkwsT2108oapk=
github.com/onsi/ginkgo v1.16.5 h1:8xi0RTUf59SOSfEtZMvwTvXYMzG4gV23XVHOZiXNtnE=
github.com/onsi/ginkgo v1.16.5/go.mod h1:+E8gABHa3K6zRBolWtd+ROzc/U5bkGt0FwiG042wbpU=
github.com/onsi/ginkgo/v2 v2.11.0 h1:WgqUCUt/lT6yXoQ8Wef0fsNn5cAuMK7+KT9UFRz2tcU=
github.com/onsi/ginkgo/v2 v2.12.0 h1:UIVDowFPwpg6yMUpPjGkYvf06K3RAiJXUhCxEwQVHRI=
github.com/onsi/gomega v1.7.1/go.mod h1:XdKZgCCFLUoM/7CFJVPcG8C1xQ1AJ0vpAezJrB7JYyY=
github.com/onsi/gomega v1.10.1/go.mod h1:iN09h71vgCQne3DLsj+A5owkum+a2tYe+TOCB1ybHNo=
github.com/onsi/gomega v1.27.10 h1:naR28SdDFlqrG6kScpT8VWpu1xWY5nJRCF3XaYyBjhI=
Expand Down Expand Up @@ -112,8 +114,8 @@ golang.org/x/crypto v0.13.0/go.mod h1:y6Z2r+Rw4iayiXXAIxJIDAJ1zMW4yaTpebo8fPOliY
golang.org/x/lint v0.0.0-20200302205851-738671d3881b/go.mod h1:3xt1FjdF8hUf6vQPIChWIBhFzV8gjjsPE/fR3IyQdNY=
golang.org/x/mod v0.1.1-0.20191105210325-c90efee705ee/go.mod h1:QqPTAvyqsEbceGzBzNggFXnrqF1CaUcvgkdR5Ot7KZg=
golang.org/x/mod v0.3.0/go.mod h1:s0Qsj1ACt9ePp/hMypM3fl4fZqREWJwdYDEqhRiZZUA=
golang.org/x/mod v0.10.0 h1:lFO9qtOdlre5W1jxS3r/4szv2/6iXxScdzjoBMXNhYk=
golang.org/x/mod v0.10.0/go.mod h1:iBbtSCu2XBx23ZKBPSOrRkjjQPZFPuis4dIYUhu/chs=
golang.org/x/mod v0.12.0 h1:rmsUpXtvNzj340zd98LZ4KntptpfRHwpFOHG188oHXc=
golang.org/x/mod v0.12.0/go.mod h1:iBbtSCu2XBx23ZKBPSOrRkjjQPZFPuis4dIYUhu/chs=
golang.org/x/net v0.0.0-20180906233101-161cd47e91fd/go.mod h1:mL1N/T3taQHkDXs73rZJwtUhF3w3ftmwwsq0BUmARs4=
golang.org/x/net v0.0.0-20190311183353-d8887717615a/go.mod h1:t9HGtf8HONx5eT2rtn7q6eTqICYqUVnKs3thJo3Qplg=
golang.org/x/net v0.0.0-20190404232315-eb5bcb51f2a3/go.mod h1:t9HGtf8HONx5eT2rtn7q6eTqICYqUVnKs3thJo3Qplg=
Expand All @@ -127,8 +129,8 @@ golang.org/x/net v0.0.0-20200520004742-59133d7f0dd7/go.mod h1:qpuaurCH72eLCgpAm/
golang.org/x/net v0.0.0-20201010224723-4f7140c49acb/go.mod h1:sp8m0HH+o8qH0wwXwYZr8TS3Oi6o0r6Gce1SSxlDquU=
golang.org/x/net v0.0.0-20201021035429-f5854403a974/go.mod h1:sp8m0HH+o8qH0wwXwYZr8TS3Oi6o0r6Gce1SSxlDquU=
golang.org/x/net v0.0.0-20201110031124-69a78807bb2b/go.mod h1:sp8m0HH+o8qH0wwXwYZr8TS3Oi6o0r6Gce1SSxlDquU=
golang.org/x/net v0.12.0 h1:cfawfvKITfUsFCeJIHJrbSxpeu/E81khclypR0GVT50=
golang.org/x/net v0.12.0/go.mod h1:zEVYFnQC7m/vmpQFELhcD1EWkZlX69l4oqgmer6hfKA=
golang.org/x/net v0.14.0 h1:BONx9s002vGdD9umnlX1Po8vOZmrgH34qlHcD1MfK14=
golang.org/x/net v0.14.0/go.mod h1:PpSgVXXLK0OxS0F31C1/tv6XNguvCrnXIDrFMspZIUI=
golang.org/x/sync v0.0.0-20180314180146-1d60e4601c6f/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
golang.org/x/sync v0.0.0-20190423024810-112230192c58/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
golang.org/x/sync v0.0.0-20201020160332-67f06af15bc9/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
Expand Down Expand Up @@ -170,8 +172,8 @@ golang.org/x/tools v0.0.0-20190328211700-ab21143f2384/go.mod h1:LCzVGOaR6xXOjkQ3
golang.org/x/tools v0.0.0-20191119224855-298f0cb1881e/go.mod h1:b+2E5dAYhXwXZwtnZ6UAqBI28+e2cm9otk0dWdXHAEo=
golang.org/x/tools v0.0.0-20200130002326-2f3ba24bd6e7/go.mod h1:TB2adYChydJhpapKDTa4BR/hXlZSLoq2Wpct/0txZ28=
golang.org/x/tools v0.0.0-20201224043029-2b0845dc783e/go.mod h1:emZCQorbCU4vsT4fOWvOPXz4eW1wZW4PmDk9uLelYpA=
golang.org/x/tools v0.9.3 h1:Gn1I8+64MsuTb/HpH+LmQtNas23LhUVr3rYZ0eKuaMM=
golang.org/x/tools v0.9.3/go.mod h1:owI94Op576fPu3cIGQeHs3joujW/2Oc6MtlxbF5dfNc=
golang.org/x/tools v0.12.0 h1:YW6HUoUmYBpwSgyaGaZq1fHjrBjX1rlpZ54T6mu2kss=
golang.org/x/tools v0.12.0/go.mod h1:Sc0INKfu04TlqNoRA1hgpFZbhYXHPr4V5DzpSBTPqQM=
golang.org/x/xerrors v0.0.0-20190717185122-a985d3407aa7/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
golang.org/x/xerrors v0.0.0-20191011141410-1b5146add898/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
Expand Down
Loading

0 comments on commit 2d105d3

Please sign in to comment.