Skip to content

Commit

Permalink
Supports closing of sessions when killed
Browse files Browse the repository at this point in the history
  • Loading branch information
Hentioe committed Jul 26, 2024
1 parent c35640e commit 9f07a3e
Show file tree
Hide file tree
Showing 2 changed files with 38 additions and 0 deletions.
20 changes: 20 additions & 0 deletions tsshd/main.go
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,8 @@ import (
"io"
"os"
"os/exec"
"os/signal"
"syscall"
"time"

"github.com/trzsz/go-arg"
Expand Down Expand Up @@ -99,6 +101,9 @@ func TsshdMain() int {
// cleanup on exit
defer cleanupOnExit()

// handle exit signals
handleExitSignals()

kcpListener, quicListener, err := initServer(&args)
if err != nil {
fmt.Println(err)
Expand Down Expand Up @@ -127,3 +132,18 @@ func TsshdMain() int {

return <-exitChan
}

func handleExitSignals() {
sigChan := make(chan os.Signal, 1)
signal.Notify(sigChan,
syscall.SIGTERM, // Default signal for the kill command
syscall.SIGINT, // Ctrl+C signal
syscall.SIGHUP, // Terminal closed (System reboot/shutdown)
)

go func() {
<-sigChan
trySendErrorMessage("tsshd has been terminated")
closeAllSessions()
}()
}
18 changes: 18 additions & 0 deletions tsshd/session.go
Original file line number Diff line number Diff line change
Expand Up @@ -48,6 +48,7 @@ type sessionContext struct {
stdout io.ReadCloser
stderr io.ReadCloser
started bool
closed bool
}

type stderrStream struct {
Expand Down Expand Up @@ -130,6 +131,10 @@ func (c *sessionContext) Wait() {
}

func (c *sessionContext) Close() {
if c.closed {
return
}
c.closed = true
if err := sendBusMessage("exit", ExitMessage{
ID: c.id,
ExitCode: c.cmd.ProcessState.ExitCode(),
Expand Down Expand Up @@ -447,3 +452,16 @@ func handleChannelAccept(listener net.Listener, channelType string) {
}(conn)
}
}

func closeAllSessions() {
sessionMutex.Lock()
var sessions []*sessionContext
for _, session := range sessionMap {
sessions = append(sessions, session)
}
sessionMutex.Unlock()

for _, session := range sessions {
session.Close()
}
}

0 comments on commit 9f07a3e

Please sign in to comment.