From d402cf103b00b8a55c6cab665ab36d4fd282f3a5 Mon Sep 17 00:00:00 2001 From: Paul Buetow Date: Fri, 11 Sep 2020 11:34:12 +0100 Subject: no SIGINFO for now as not supported on Linux --- internal/io/signal/signal.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) (limited to 'internal') diff --git a/internal/io/signal/signal.go b/internal/io/signal/signal.go index bca7e6e..0d49485 100644 --- a/internal/io/signal/signal.go +++ b/internal/io/signal/signal.go @@ -10,7 +10,7 @@ import ( // StatsCh returns a channel for "please print stats" signalling. func StatsCh(ctx context.Context) <-chan struct{} { sigCh := make(chan os.Signal) - gosignal.Notify(sigCh, syscall.SIGINFO, syscall.SIGUSR1) + gosignal.Notify(sigCh, syscall.SIGUSR1) statsCh := make(chan struct{}) -- cgit v1.2.3 From ec67d9833095dfbe620dd3c99ea0caba391c4b87 Mon Sep 17 00:00:00 2001 From: Paul Buetow Date: Fri, 11 Sep 2020 12:30:29 +0100 Subject: hit ctrl+c twice to exit --- internal/clients/stats.go | 17 ++++++++++------- internal/io/signal/signal.go | 14 +++++++++++--- 2 files changed, 21 insertions(+), 10 deletions(-) (limited to 'internal') diff --git a/internal/clients/stats.go b/internal/clients/stats.go index a6ac0c5..e7eabd8 100644 --- a/internal/clients/stats.go +++ b/internal/clients/stats.go @@ -4,6 +4,7 @@ import ( "context" "fmt" "runtime" + "strings" "sync" "time" @@ -54,7 +55,8 @@ func (s *stats) Start(ctx context.Context, throttleCh, statsCh <-chan struct{}) if connected == connectedLast && !force { continue } - s.log(connected, newConnections, throttle) + + logger.Info(s.statsLine(connected, newConnections, throttle)) connectedLast = connected s.mutex.Lock() @@ -63,15 +65,16 @@ func (s *stats) Start(ctx context.Context, throttleCh, statsCh <-chan struct{}) } } -func (s *stats) log(connected, newConnections int, throttle int) { +func (s *stats) statsLine(connected, newConnections int, throttle int) string { percConnected := percentOf(float64(s.connectionsTotal), float64(connected)) - connectedStr := fmt.Sprintf("connected=%d/%d(%d%%)", connected, s.connectionsTotal, int(percConnected)) - newConnStr := fmt.Sprintf("new=%d", newConnections) - throttleStr := fmt.Sprintf("throttle=%d", throttle) - cpusGoroutinesStr := fmt.Sprintf("cpus/goroutines=%d/%d", runtime.NumCPU(), runtime.NumGoroutine()) + var stats []string + stats = append(stats, fmt.Sprintf("connected=%d/%d(%d%%)", connected, s.connectionsTotal, int(percConnected))) + stats = append(stats, fmt.Sprintf("new=%d", newConnections)) + stats = append(stats, fmt.Sprintf("throttle=%d", throttle)) + stats = append(stats, fmt.Sprintf("cpus/goroutines=%d/%d", runtime.NumCPU(), runtime.NumGoroutine())) - logger.Info("stats", connectedStr, newConnStr, throttleStr, cpusGoroutinesStr) + return strings.Join(stats, "|") } func (s *stats) numConnected() int { diff --git a/internal/io/signal/signal.go b/internal/io/signal/signal.go index 0d49485..3785d71 100644 --- a/internal/io/signal/signal.go +++ b/internal/io/signal/signal.go @@ -4,13 +4,15 @@ import ( "context" "os" gosignal "os/signal" - "syscall" + "time" + + "github.com/mimecast/dtail/internal/io/logger" ) // StatsCh returns a channel for "please print stats" signalling. -func StatsCh(ctx context.Context) <-chan struct{} { +func InterruptCh(ctx context.Context) <-chan struct{} { sigCh := make(chan os.Signal) - gosignal.Notify(sigCh, syscall.SIGUSR1) + gosignal.Notify(sigCh, os.Interrupt) statsCh := make(chan struct{}) @@ -20,6 +22,12 @@ func StatsCh(ctx context.Context) <-chan struct{} { case <-sigCh: select { case statsCh <- struct{}{}: + logger.Info("Hit Ctrl+C twice to exit") + select { + case <-sigCh: + os.Exit(0) + case <-time.After(time.Second): + } default: // Stats currently already printed. } -- cgit v1.2.3 From cdbcef72a1b9f93d08455b5a77f7f7afa2b9f53d Mon Sep 17 00:00:00 2001 From: Paul Buetow Date: Fri, 11 Sep 2020 14:55:50 +0100 Subject: fix nil pointer panic --- internal/clients/maprclient.go | 2 +- internal/regex/regex.go | 4 ++++ 2 files changed, 5 insertions(+), 1 deletion(-) (limited to 'internal') diff --git a/internal/clients/maprclient.go b/internal/clients/maprclient.go index 581db44..149129d 100644 --- a/internal/clients/maprclient.go +++ b/internal/clients/maprclient.go @@ -123,7 +123,7 @@ func (c MaprClient) makeCommands() (commands []string) { commands = append(commands, fmt.Sprintf("timeout %d %s %s %s", c.Timeout, modeStr, file, c.Regex.Serialize())) continue } - commands = append(commands, fmt.Sprintf("%s %s regex %s", modeStr, file, c.Regex.Serialize())) + commands = append(commands, fmt.Sprintf("%s %s %s", modeStr, file, c.Regex.Serialize())) } return diff --git a/internal/regex/regex.go b/internal/regex/regex.go index 707cb48..f668c38 100644 --- a/internal/regex/regex.go +++ b/internal/regex/regex.go @@ -39,6 +39,10 @@ func New(regexStr string, flag Flag) (Regex, error) { } func new(regexStr string, flags []Flag) (Regex, error) { + if len(flags) == 0 { + flags = append(flags, Default) + } + r := Regex{ regexStr: regexStr, flags: flags, -- cgit v1.2.3 From 813d2d00ec581c801d64091c7774988b559c3e93 Mon Sep 17 00:00:00 2001 From: Paul Buetow Date: Sat, 19 Sep 2020 17:52:45 +0100 Subject: refactor to have no context.Context in client handler structs --- internal/clients/baseclient.go | 2 +- internal/clients/handlers/basehandler.go | 22 +++++++++++++----- internal/clients/handlers/clienthandler.go | 4 +--- internal/clients/handlers/done.go | 37 ++++++++++++++++++++++++++++++ internal/clients/handlers/handler.go | 3 +-- internal/clients/handlers/healthhandler.go | 15 ++++++++---- internal/clients/handlers/maprhandler.go | 4 +--- internal/clients/handlers/withcancel.go | 24 ------------------- internal/clients/healthclient.go | 2 +- internal/clients/remote/connection.go | 7 +++--- 10 files changed, 73 insertions(+), 47 deletions(-) create mode 100644 internal/clients/handlers/done.go delete mode 100644 internal/clients/handlers/withcancel.go (limited to 'internal') diff --git a/internal/clients/baseclient.go b/internal/clients/baseclient.go index 008a01e..d8d4fde 100644 --- a/internal/clients/baseclient.go +++ b/internal/clients/baseclient.go @@ -99,7 +99,7 @@ func (c *baseClient) start(ctx context.Context, active chan struct{}, i int, con defer func() { <-active }() for { - connCtx, cancel := conn.Handler.WithCancel(ctx) + connCtx, cancel := context.WithCancel(ctx) defer cancel() conn.Start(connCtx, cancel, c.throttleCh, c.stats.connectionsEstCh) diff --git a/internal/clients/handlers/basehandler.go b/internal/clients/handlers/basehandler.go index 65bbfd7..54b80ae 100644 --- a/internal/clients/handlers/basehandler.go +++ b/internal/clients/handlers/basehandler.go @@ -13,7 +13,7 @@ import ( ) type baseHandler struct { - withCancel + done *Done server string shellStarted bool commands chan string @@ -29,6 +29,14 @@ func (h *baseHandler) Status() int { return h.status } +func (h *baseHandler) Done() <-chan struct{} { + return h.done.Done() +} + +func (h *baseHandler) Shutdown() { + h.done.Shutdown() +} + // SendMessage to the server. func (h *baseHandler) SendMessage(command string) error { encoded := base64.StdEncoding.EncodeToString([]byte(command)) @@ -38,7 +46,8 @@ func (h *baseHandler) SendMessage(command string) error { case h.commands <- fmt.Sprintf("protocol %s base64 %v;", version.ProtocolCompat, encoded): case <-time.After(time.Second * 5): return fmt.Errorf("Timed out sending command '%s' (base64: '%s')", command, encoded) - case <-h.ctx.Done(): + case <-h.Done(): + return nil } return nil @@ -65,7 +74,7 @@ func (h *baseHandler) Read(p []byte) (n int, err error) { select { case command := <-h.commands: n = copy(p, []byte(command)) - case <-h.ctx.Done(): + case <-h.Done(): return 0, io.EOF } return @@ -95,10 +104,11 @@ func (h *baseHandler) handleHiddenMessage(message string) { case strings.HasPrefix(message, ".syn close connection"): h.SendMessage(".ack close connection") select { - case <-time.After(time.Second * 1): + case <-time.After(time.Second * 5): logger.Debug("Shutting down client after timeout and sending ack to server") - h.withCancel.shutdown() - case <-h.ctx.Done(): + h.Shutdown() + case <-h.Done(): + return } case strings.HasPrefix(message, ".run exitstatus"): diff --git a/internal/clients/handlers/clienthandler.go b/internal/clients/handlers/clienthandler.go index fcd8052..2908f7c 100644 --- a/internal/clients/handlers/clienthandler.go +++ b/internal/clients/handlers/clienthandler.go @@ -19,9 +19,7 @@ func NewClientHandler(server string) *ClientHandler { shellStarted: false, commands: make(chan string), status: -1, - withCancel: withCancel{ - done: make(chan struct{}), - }, + done: NewDone(), }, } } diff --git a/internal/clients/handlers/done.go b/internal/clients/handlers/done.go new file mode 100644 index 0000000..5b1335e --- /dev/null +++ b/internal/clients/handlers/done.go @@ -0,0 +1,37 @@ +package handlers + +import ( + "sync" + + "github.com/mimecast/dtail/internal/io/logger" +) + +type Done struct { + ch chan struct{} + mutex sync.Mutex +} + +func NewDone() *Done { + return &Done{ + ch: make(chan struct{}), + } +} + +func (d *Done) Done() <-chan struct{} { + return d.ch +} + +func (d *Done) Shutdown() { + d.mutex.Lock() + defer d.mutex.Unlock() + + logger.Debug("Done.Shutdown()") + + select { + case <-d.ch: + return + default: + logger.Debug("Done.Shutdown() -> close") + close(d.ch) + } +} diff --git a/internal/clients/handlers/handler.go b/internal/clients/handlers/handler.go index c53ca34..afa87e2 100644 --- a/internal/clients/handlers/handler.go +++ b/internal/clients/handlers/handler.go @@ -1,7 +1,6 @@ package handlers import ( - "context" "io" ) @@ -11,6 +10,6 @@ type Handler interface { SendMessage(command string) error Server() string Status() int - WithCancel(ctx context.Context) (context.Context, context.CancelFunc) + Shutdown() Done() <-chan struct{} } diff --git a/internal/clients/handlers/healthhandler.go b/internal/clients/handlers/healthhandler.go index 9051015..9fc2671 100644 --- a/internal/clients/handlers/healthhandler.go +++ b/internal/clients/handlers/healthhandler.go @@ -8,7 +8,7 @@ import ( // HealthHandler implements the handler required for health checks. type HealthHandler struct { - withCancel + done *Done // Buffer of incoming data from server. receiveBuf []byte // To send commands to the server. @@ -27,9 +27,7 @@ func NewHealthHandler(server string, receive chan<- string) *HealthHandler { receive: receive, commands: make(chan string), status: -1, - withCancel: withCancel{ - done: make(chan struct{}), - }, + done: NewDone(), } return &h @@ -45,12 +43,21 @@ func (h *HealthHandler) Status() int { return h.status } +func (h *HealthHandler) Done() <-chan struct{} { + return h.done.Done() +} + +func (h *HealthHandler) Shutdown() { + h.done.Shutdown() +} + // SendMessage sends a DTail command to the server. func (h *HealthHandler) SendMessage(command string) error { select { case h.commands <- fmt.Sprintf("%s;", command): case <-time.NewTimer(time.Second * 10).C: return errors.New("Timed out sending command " + command) + case <-h.Done(): } return nil diff --git a/internal/clients/handlers/maprhandler.go b/internal/clients/handlers/maprhandler.go index b908f3b..5d98690 100644 --- a/internal/clients/handlers/maprhandler.go +++ b/internal/clients/handlers/maprhandler.go @@ -24,9 +24,7 @@ func NewMaprHandler(server string, query *mapr.Query, globalGroup *mapr.GlobalGr shellStarted: false, commands: make(chan string), status: -1, - withCancel: withCancel{ - done: make(chan struct{}), - }, + done: NewDone(), }, query: query, aggregate: client.NewAggregate(server, query, globalGroup), diff --git a/internal/clients/handlers/withcancel.go b/internal/clients/handlers/withcancel.go deleted file mode 100644 index 7c9cf4e..0000000 --- a/internal/clients/handlers/withcancel.go +++ /dev/null @@ -1,24 +0,0 @@ -package handlers - -import "context" - -type withCancel struct { - ctx context.Context - done chan struct{} -} - -// WithCancel sets and returns the context used. -func (w *withCancel) WithCancel(ctx context.Context) (context.Context, context.CancelFunc) { - cancelCtx, cancel := context.WithCancel(ctx) - w.ctx = cancelCtx - - return cancelCtx, cancel -} - -func (w *withCancel) Done() <-chan struct{} { - return w.done -} - -func (w *withCancel) shutdown() { - close(w.done) -} diff --git a/internal/clients/healthclient.go b/internal/clients/healthclient.go index 7313583..e93f6be 100644 --- a/internal/clients/healthclient.go +++ b/internal/clients/healthclient.go @@ -50,7 +50,7 @@ func (c *HealthClient) Start(ctx context.Context) (status int) { conn.Handler = handlers.NewHealthHandler(c.server, receive) conn.Commands = []string{c.mode.String()} - connCtx, cancel := conn.Handler.WithCancel(ctx) + connCtx, cancel := context.WithCancel(ctx) go conn.Start(connCtx, cancel, throttleCh, statsCh) for { diff --git a/internal/clients/remote/connection.go b/internal/clients/remote/connection.go index 2d97d14..b29ffed 100644 --- a/internal/clients/remote/connection.go +++ b/internal/clients/remote/connection.go @@ -177,21 +177,21 @@ func (c *Connection) handle(ctx context.Context, cancel context.CancelFunc, sess } go func() { - defer cancel() io.Copy(stdinPipe, c.Handler) + cancel() }() go func() { - defer cancel() io.Copy(c.Handler, stdoutPipe) + cancel() }() go func() { - defer cancel() select { case <-c.Handler.Done(): case <-ctx.Done(): } + cancel() }() // Send all commands to client. @@ -207,5 +207,6 @@ func (c *Connection) handle(ctx context.Context, cancel context.CancelFunc, sess } <-ctx.Done() + c.Handler.Shutdown() return nil } -- cgit v1.2.3 From bad8e04bb4410b94b8e875ccde287f74ab94121a Mon Sep 17 00:00:00 2001 From: Paul Buetow Date: Sat, 19 Sep 2020 18:38:36 +0100 Subject: server handler context refactoring --- internal/clients/handlers/basehandler.go | 3 +- internal/clients/handlers/clienthandler.go | 3 +- internal/clients/handlers/done.go | 37 --------------- internal/clients/handlers/healthhandler.go | 6 ++- internal/clients/handlers/maprhandler.go | 3 +- internal/done.go | 32 +++++++++++++ internal/server/handlers/controlhandler.go | 26 +++++++---- internal/server/handlers/handler.go | 2 + internal/server/handlers/serverhandler.go | 73 ++++++++++++++++-------------- internal/server/server.go | 39 +++++++--------- 10 files changed, 116 insertions(+), 108 deletions(-) delete mode 100644 internal/clients/handlers/done.go create mode 100644 internal/done.go (limited to 'internal') diff --git a/internal/clients/handlers/basehandler.go b/internal/clients/handlers/basehandler.go index 54b80ae..b5045e2 100644 --- a/internal/clients/handlers/basehandler.go +++ b/internal/clients/handlers/basehandler.go @@ -8,12 +8,13 @@ import ( "strings" "time" + "github.com/mimecast/dtail/internal" "github.com/mimecast/dtail/internal/io/logger" "github.com/mimecast/dtail/internal/version" ) type baseHandler struct { - done *Done + done *internal.Done server string shellStarted bool commands chan string diff --git a/internal/clients/handlers/clienthandler.go b/internal/clients/handlers/clienthandler.go index 2908f7c..2bcb038 100644 --- a/internal/clients/handlers/clienthandler.go +++ b/internal/clients/handlers/clienthandler.go @@ -1,6 +1,7 @@ package handlers import ( + "github.com/mimecast/dtail/internal" "github.com/mimecast/dtail/internal/io/logger" ) @@ -19,7 +20,7 @@ func NewClientHandler(server string) *ClientHandler { shellStarted: false, commands: make(chan string), status: -1, - done: NewDone(), + done: internal.NewDone(), }, } } diff --git a/internal/clients/handlers/done.go b/internal/clients/handlers/done.go deleted file mode 100644 index 5b1335e..0000000 --- a/internal/clients/handlers/done.go +++ /dev/null @@ -1,37 +0,0 @@ -package handlers - -import ( - "sync" - - "github.com/mimecast/dtail/internal/io/logger" -) - -type Done struct { - ch chan struct{} - mutex sync.Mutex -} - -func NewDone() *Done { - return &Done{ - ch: make(chan struct{}), - } -} - -func (d *Done) Done() <-chan struct{} { - return d.ch -} - -func (d *Done) Shutdown() { - d.mutex.Lock() - defer d.mutex.Unlock() - - logger.Debug("Done.Shutdown()") - - select { - case <-d.ch: - return - default: - logger.Debug("Done.Shutdown() -> close") - close(d.ch) - } -} diff --git a/internal/clients/handlers/healthhandler.go b/internal/clients/handlers/healthhandler.go index 9fc2671..95693ab 100644 --- a/internal/clients/handlers/healthhandler.go +++ b/internal/clients/handlers/healthhandler.go @@ -4,11 +4,13 @@ import ( "errors" "fmt" "time" + + "github.com/mimecast/dtail/internal" ) // HealthHandler implements the handler required for health checks. type HealthHandler struct { - done *Done + done *internal.Done // Buffer of incoming data from server. receiveBuf []byte // To send commands to the server. @@ -27,7 +29,7 @@ func NewHealthHandler(server string, receive chan<- string) *HealthHandler { receive: receive, commands: make(chan string), status: -1, - done: NewDone(), + done: internal.NewDone(), } return &h diff --git a/internal/clients/handlers/maprhandler.go b/internal/clients/handlers/maprhandler.go index 5d98690..fb71c8f 100644 --- a/internal/clients/handlers/maprhandler.go +++ b/internal/clients/handlers/maprhandler.go @@ -3,6 +3,7 @@ package handlers import ( "strings" + "github.com/mimecast/dtail/internal" "github.com/mimecast/dtail/internal/io/logger" "github.com/mimecast/dtail/internal/mapr" "github.com/mimecast/dtail/internal/mapr/client" @@ -24,7 +25,7 @@ func NewMaprHandler(server string, query *mapr.Query, globalGroup *mapr.GlobalGr shellStarted: false, commands: make(chan string), status: -1, - done: NewDone(), + done: internal.NewDone(), }, query: query, aggregate: client.NewAggregate(server, query, globalGroup), diff --git a/internal/done.go b/internal/done.go new file mode 100644 index 0000000..2326eee --- /dev/null +++ b/internal/done.go @@ -0,0 +1,32 @@ +package internal + +import ( + "sync" +) + +type Done struct { + ch chan struct{} + mutex sync.Mutex +} + +func NewDone() *Done { + return &Done{ + ch: make(chan struct{}), + } +} + +func (d *Done) Done() <-chan struct{} { + return d.ch +} + +func (d *Done) Shutdown() { + d.mutex.Lock() + defer d.mutex.Unlock() + + select { + case <-d.ch: + return + default: + close(d.ch) + } +} diff --git a/internal/server/handlers/controlhandler.go b/internal/server/handlers/controlhandler.go index daa9835..9a8eb75 100644 --- a/internal/server/handlers/controlhandler.go +++ b/internal/server/handlers/controlhandler.go @@ -1,20 +1,19 @@ package handlers import ( - "context" "fmt" "io" "os" "strings" + "github.com/mimecast/dtail/internal" "github.com/mimecast/dtail/internal/io/logger" user "github.com/mimecast/dtail/internal/user/server" ) // ControlHandler is used for control functions and health monitoring. type ControlHandler struct { - ctx context.Context - done chan struct{} + done *internal.Done hostname string payload []byte serverMessages chan string @@ -22,12 +21,11 @@ type ControlHandler struct { } // NewControlHandler returns a new control handler. -func NewControlHandler(ctx context.Context, user *user.User) (*ControlHandler, <-chan struct{}) { +func NewControlHandler(user *user.User) *ControlHandler { logger.Debug(user, "Creating control handler") h := ControlHandler{ - ctx: ctx, - done: make(chan struct{}), + done: internal.NewDone(), serverMessages: make(chan string, 10), user: user, } @@ -40,7 +38,15 @@ func NewControlHandler(ctx context.Context, user *user.User) (*ControlHandler, < s := strings.Split(fqdn, ".") h.hostname = s[0] - return &h, h.done + return &h +} + +func (h *ControlHandler) Shutdown() { + h.done.Shutdown() +} + +func (h *ControlHandler) Done() <-chan struct{} { + return h.done.Done() } // Read is to send data to the client via the Reader interface. @@ -51,7 +57,7 @@ func (h *ControlHandler) Read(p []byte) (n int, err error) { wholePayload := []byte(fmt.Sprintf("SERVER|%s|%s\n", h.hostname, message)) n = copy(p, wholePayload) return - case <-h.ctx.Done(): + case <-h.done.Done(): return 0, io.EOF } } @@ -63,7 +69,7 @@ func (h *ControlHandler) Write(p []byte) (n int, err error) { switch c { case ';': wholePayload := strings.TrimSpace(string(h.payload)) - h.handleCommand(h.ctx, wholePayload) + h.handleCommand(wholePayload) h.payload = nil default: @@ -75,7 +81,7 @@ func (h *ControlHandler) Write(p []byte) (n int, err error) { return } -func (h *ControlHandler) handleCommand(ctx context.Context, command string) { +func (h *ControlHandler) handleCommand(command string) { logger.Info(h.user, command) s := strings.Split(command, " ") logger.Debug(h.user, "Receiving command", command, s) diff --git a/internal/server/handlers/handler.go b/internal/server/handlers/handler.go index c42ceb9..b04e854 100644 --- a/internal/server/handlers/handler.go +++ b/internal/server/handlers/handler.go @@ -5,4 +5,6 @@ import "io" // Handler interface for server side functionality. type Handler interface { io.ReadWriter + Shutdown() + Done() <-chan struct{} } diff --git a/internal/server/handlers/serverhandler.go b/internal/server/handlers/serverhandler.go index 7017f3e..5af3ebb 100644 --- a/internal/server/handlers/serverhandler.go +++ b/internal/server/handlers/serverhandler.go @@ -13,6 +13,7 @@ import ( "sync/atomic" "time" + "github.com/mimecast/dtail/internal" "github.com/mimecast/dtail/internal/config" "github.com/mimecast/dtail/internal/io/line" "github.com/mimecast/dtail/internal/io/logger" @@ -31,33 +32,28 @@ const ( // the Bi-directional communication between SSH client and server. // This handler implements the handler of the SSH server. type ServerHandler struct { - lines chan line.Line - regex string - aggregate *server.Aggregate - aggregatedMessages chan string - serverMessages chan string - payload []byte - hostname string - user *user.User - // TODO: Move all these channels into a separate struct for readability! + done *internal.Done + lines chan line.Line + regex string + aggregate *server.Aggregate + aggregatedMessages chan string + serverMessages chan string + payload []byte + hostname string + user *user.User catLimiter chan struct{} tailLimiter chan struct{} globalServerWaitFor chan struct{} ackCloseReceived chan struct{} - serverCtx context.Context - handlerCtx context.Context - done chan struct{} activeCommands int32 activeReaders int32 background background.Background } // NewServerHandler returns the server handler. -func NewServerHandler(handlerCtx, serverCtx context.Context, user *user.User, catLimiter, tailLimiter, globalServerWaitFor chan struct{}, background background.Background) (*ServerHandler, <-chan struct{}) { +func NewServerHandler(user *user.User, catLimiter, tailLimiter, globalServerWaitFor chan struct{}, background background.Background) *ServerHandler { h := ServerHandler{ - serverCtx: serverCtx, - handlerCtx: handlerCtx, - done: make(chan struct{}), + done: internal.NewDone(), lines: make(chan line.Line, 100), serverMessages: make(chan string, 10), aggregatedMessages: make(chan string, 10), @@ -78,7 +74,15 @@ func NewServerHandler(handlerCtx, serverCtx context.Context, user *user.User, ca s := strings.Split(fqdn, ".") h.hostname = s[0] - return &h, h.done + return &h +} + +func (h *ServerHandler) Shutdown() { + h.done.Shutdown() +} + +func (h *ServerHandler) Done() <-chan struct{} { + return h.done.Done() } // Read is to send data to the dtail client via Reader interface. @@ -120,7 +124,7 @@ func (h *ServerHandler) Read(p []byte) (n int, err error) { case <-time.After(time.Second): // Once in a while check whether we are done. select { - case <-h.handlerCtx.Done(): + case <-h.done.Done(): return 0, io.EOF default: } @@ -134,7 +138,7 @@ func (h *ServerHandler) Write(p []byte) (n int, err error) { switch c { case ';': commandStr := strings.TrimSpace(string(h.payload)) - h.handleCommand(h.handlerCtx, commandStr) + h.handleCommand(commandStr) h.payload = nil default: h.payload = append(h.payload, c) @@ -145,9 +149,10 @@ func (h *ServerHandler) Write(p []byte) (n int, err error) { return } -func (h *ServerHandler) handleCommand(ctx context.Context, commandStr string) { +func (h *ServerHandler) handleCommand(commandStr string) { logger.Debug(h.user, commandStr) var timeout time.Duration + ctx := context.Background() args, argc, err := h.handleProtocolVersion(strings.Split(commandStr, " ")) if err != nil { @@ -172,15 +177,21 @@ func (h *ServerHandler) handleCommand(ctx context.Context, commandStr string) { return } + ctx, cancel := context.WithCancel(ctx) + go func() { + <-h.done.Done() + cancel() + }() + if timeout > 0 { logger.Info(h.user, "Command with timeout context", argc, args, timeout) - commandCtx, cancel := context.WithTimeout(ctx, timeout) + ctx, cancel := context.WithTimeout(ctx, timeout) go func() { - <-commandCtx.Done() + <-ctx.Done() logger.Info(h.user, "Command timed out, canceling it", args, args, timeout) cancel() }() - h.handleUserCommand(commandCtx, argc, args, timeout) + h.handleUserCommand(ctx, argc, args, timeout) return } @@ -348,9 +359,8 @@ func (h *ServerHandler) handleUserCommand(ctx context.Context, argc int, args [] // Set default background timeout. timeout = time.Hour * 1 } - // Use a new context based on the server context, so that background job does not get - // terminated when handler/SSH connection terminates. - commandCtx, cancel := context.WithTimeout(h.serverCtx, timeout) + + commandCtx, cancel := context.WithTimeout(ctx, timeout) if err := h.background.Add(h.user.Name, jobName, cancel, &wg); err != nil { h.sendServerMessage(logger.Error(h.user, err, jobName, args)) @@ -406,7 +416,7 @@ func (h *ServerHandler) handleAckCommand(argc int, args []string) { func (h *ServerHandler) send(ch chan<- string, message string) { select { case ch <- message: - case <-h.handlerCtx.Done(): + case <-h.done.Done(): } } @@ -447,7 +457,7 @@ func (h *ServerHandler) shutdown() { go func() { select { case h.serverMessageC() <- ".syn close connection": - case <-h.handlerCtx.Done(): + case <-h.done.Done(): } }() @@ -455,13 +465,10 @@ func (h *ServerHandler) shutdown() { case <-h.ackCloseReceived: case <-time.After(time.Second * 5): logger.Debug(h.user, "Shutdown timeout reached, enforcing shutdown") - case <-h.handlerCtx.Done(): + case <-h.done.Done(): } - select { - case h.done <- struct{}{}: - default: - } + h.done.Shutdown() } func (h *ServerHandler) incrementActiveCommands() { diff --git a/internal/server/server.go b/internal/server/server.go index a446738..5e2a521 100644 --- a/internal/server/server.go +++ b/internal/server/server.go @@ -178,53 +178,46 @@ func (s *Server) handleRequests(ctx context.Context, sshConn gossh.Conn, in <-ch switch req.Type { case "shell": - handlerCtx, cancel := context.WithCancel(ctx) - var handler handlers.Handler - var done <-chan struct{} - switch user.Name { case config.ControlUser: - handler, done = handlers.NewControlHandler(handlerCtx, user) + handler = handlers.NewControlHandler(user) default: - handler, done = handlers.NewServerHandler(handlerCtx, ctx, user, s.catLimiter, s.tailLimiter, s.shutdownWaitFor, s.background) + handler = handlers.NewServerHandler(user, s.catLimiter, s.tailLimiter, s.shutdownWaitFor, s.background) } - go func() { - // Handler finished work, cancel all remaining routines - defer cancel() - - <-done - }() + terminate := func() { + handler.Shutdown() + sshConn.Close() + } go func() { // Broken pipe, cancel - defer cancel() - io.Copy(channel, handler) + terminate() }() go func() { // Broken pipe, cancel - defer cancel() - io.Copy(handler, channel) + terminate() }() go func() { - defer cancel() + select { + case <-ctx.Done(): + case <-handler.Done(): + } + terminate() + }() + go func() { if err := sshConn.Wait(); err != nil && err != io.EOF { logger.Error(user, err) } s.stats.decrementConnections() logger.Info(user, "Good bye Mister!") - }() - - go func() { - <-handlerCtx.Done() - sshConn.Close() - logger.Info(user, "Closed SSH connection") + terminate() }() // Only serving shell type -- cgit v1.2.3 From df2ff83897cde61d04b12958c6f6d458c69502f4 Mon Sep 17 00:00:00 2001 From: Paul Buetow Date: Sat, 19 Sep 2020 19:15:08 +0100 Subject: refactor to have no context.Context in server mapreduce aggregate structs --- internal/mapr/server/aggregate.go | 80 +++++++++++++++++-------------- internal/server/handlers/serverhandler.go | 2 +- 2 files changed, 44 insertions(+), 38 deletions(-) (limited to 'internal') diff --git a/internal/mapr/server/aggregate.go b/internal/mapr/server/aggregate.go index 1028943..cd59b63 100644 --- a/internal/mapr/server/aggregate.go +++ b/internal/mapr/server/aggregate.go @@ -6,6 +6,7 @@ import ( "strings" "time" + "github.com/mimecast/dtail/internal" "github.com/mimecast/dtail/internal/config" "github.com/mimecast/dtail/internal/io/line" "github.com/mimecast/dtail/internal/io/logger" @@ -15,6 +16,7 @@ import ( // Aggregate is for aggregating mapreduce data on the DTail server side. type Aggregate struct { + done *internal.Done // Log lines to process (parsing MAPREDUCE lines). Lines chan line.Line // Hostname of the current server (used to populate $hostname field). @@ -23,12 +25,12 @@ type Aggregate struct { serialize chan struct{} // Signals to flush data. flush chan struct{} + // Signals that data has been flushed + flushed chan struct{} // The mapr query query *mapr.Query // The mapr log format parser parser *logformat.Parser - cancel context.CancelFunc - ctx context.Context } // NewAggregate return a new server side aggregator. @@ -64,56 +66,63 @@ func NewAggregate(queryStr string) (*Aggregate, error) { } } - ctx, cancel := context.WithCancel(context.Background()) - a := Aggregate{ + done: internal.NewDone(), Lines: make(chan line.Line, 100), serialize: make(chan struct{}), flush: make(chan struct{}), + flushed: make(chan struct{}), hostname: s[0], query: query, parser: logParser, - ctx: ctx, - cancel: cancel, } return &a, nil } +func (a *Aggregate) Shutdown() { + a.Flush() + a.done.Shutdown() +} + // Start an aggregation. func (a *Aggregate) Start(ctx context.Context, maprLines chan<- string) { - defer a.cancel() - fieldsCh := a.linesToFields(ctx) + myCtx, cancel := context.WithCancel(ctx) + defer cancel() + + go func() { + select { + case <-myCtx.Done(): + a.done.Shutdown() + case <-a.done.Done(): + cancel() + } + }() + + fieldsCh := a.makeFields(myCtx) // Add fields (e.g. via 'set' clause) if len(a.query.Set) > 0 { - fieldsCh = a.addMoreFields(ctx, fieldsCh) + fieldsCh = a.addFields(myCtx, fieldsCh) } - go a.fieldsToMaprLines(ctx, fieldsCh, maprLines) - a.periodicAggregateTimer(ctx) + go a.aggregateTimer(myCtx) + a.makeMaprLines(myCtx, fieldsCh, maprLines) } -// Cancel the aggregation. -func (a *Aggregate) Cancel() { - a.cancel() -} - -func (a *Aggregate) periodicAggregateTimer(ctx context.Context) { +func (a *Aggregate) aggregateTimer(ctx context.Context) { for { select { case <-time.After(a.query.Interval): a.Serialize(ctx) case <-ctx.Done(): return - case <-a.ctx.Done(): - return } } } -func (a *Aggregate) linesToFields(ctx context.Context) <-chan map[string]string { +func (a *Aggregate) makeFields(ctx context.Context) <-chan map[string]string { ch := make(chan map[string]string) go func() { @@ -144,8 +153,6 @@ func (a *Aggregate) linesToFields(ctx context.Context) <-chan map[string]string } case <-ctx.Done(): return - case <-a.ctx.Done(): - return } } }() @@ -153,14 +160,14 @@ func (a *Aggregate) linesToFields(ctx context.Context) <-chan map[string]string return ch } -func (a *Aggregate) addMoreFields(ctx context.Context, fieldsCh <-chan map[string]string) <-chan map[string]string { +func (a *Aggregate) addFields(ctx context.Context, fieldsCh <-chan map[string]string) <-chan map[string]string { ch := make(chan map[string]string) go func() { defer close(ch) for { - // fieldsCh will be closed via 'linesToFields' if ctx is done + // fieldsCh will be closed via 'makeFields' if ctx is done fields, ok := <-fieldsCh if !ok { return @@ -179,7 +186,7 @@ func (a *Aggregate) addMoreFields(ctx context.Context, fieldsCh <-chan map[strin return ch } -func (a *Aggregate) fieldsToMaprLines(ctx context.Context, fieldsCh <-chan map[string]string, maprLines chan<- string) { +func (a *Aggregate) makeMaprLines(ctx context.Context, fieldsCh <-chan map[string]string, maprLines chan<- string) { group := mapr.NewGroupSet() serialize := func() { @@ -200,18 +207,10 @@ func (a *Aggregate) fieldsToMaprLines(ctx context.Context, fieldsCh <-chan map[s case <-a.serialize: serialize() case <-a.flush: - logger.Info("Flushing mapreduce result") serialize() - a.flush <- struct{}{} - logger.Info("Done flushing mapreduce result") + a.flushed <- struct{}{} case <-ctx.Done(): return - case <-a.ctx.Done(): - logger.Info("Flushing mapreduce result") - serialize() - a.flush <- struct{}{} - logger.Info("Done flushing mapreduce result") - return } } } @@ -254,6 +253,8 @@ func (a *Aggregate) aggregate(group *mapr.GroupSet, fields map[string]string) { func (a *Aggregate) Serialize(ctx context.Context) { select { case a.serialize <- struct{}{}: + case <-time.After(time.Minute): + logger.Warn("Starting to serialize mapredice data takes over a minute") case <-ctx.Done(): } } @@ -261,15 +262,20 @@ func (a *Aggregate) Serialize(ctx context.Context) { // Flush all data. func (a *Aggregate) Flush() { select { - case <-a.ctx.Done(): - return case a.flush <- struct{}{}: + logger.Info("Flushing mapreduce data") case <-time.After(time.Minute): + logger.Warn("Starting to flush mapreduce data takes over a minute") + return + case <-a.done.Done(): return } select { - case <-a.flush: + case <-a.flushed: + logger.Info("Done flushing") case <-time.After(time.Minute): + logger.Warn("Waiting for data to be flushed takes over a minute") + case <-a.done.Done(): } } diff --git a/internal/server/handlers/serverhandler.go b/internal/server/handlers/serverhandler.go index 5af3ebb..164a280 100644 --- a/internal/server/handlers/serverhandler.go +++ b/internal/server/handlers/serverhandler.go @@ -266,7 +266,7 @@ func (h *ServerHandler) handleUserCommand(ctx context.Context, argc int, args [] if h.aggregate == nil { return } - h.aggregate.Cancel() + h.aggregate.Shutdown() } } -- cgit v1.2.3 From 98ed1491676d1e69dfbd6126ce1761a473951e9f Mon Sep 17 00:00:00 2001 From: Paul Buetow Date: Sat, 19 Sep 2020 20:16:45 +0100 Subject: INFO|INFO -> INFO --- internal/io/logger/logger.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) (limited to 'internal') diff --git a/internal/io/logger/logger.go b/internal/io/logger/logger.go index d059cbb..b7db0a7 100644 --- a/internal/io/logger/logger.go +++ b/internal/io/logger/logger.go @@ -224,7 +224,7 @@ func log(what string, severity string, args []interface{}) string { return "" } - messages := []string{severity} + messages := []string{} for _, arg := range args { switch v := arg.(type) { -- cgit v1.2.3 From 7df612f527bd5dc2e785bf766d7d61124c260b94 Mon Sep 17 00:00:00 2001 From: Paul Buetow Date: Mon, 16 Nov 2020 22:11:44 +0000 Subject: remove drun command for simplicity. only focus on interactive commands dealing with log streams --- internal/clients/runclient.go | 87 ----------------- internal/io/run/run.go | 150 ------------------------------ internal/server/background/background.go | 126 ------------------------- internal/server/handlers/runcommand.go | 111 ---------------------- internal/server/handlers/serverhandler.go | 85 +---------------- internal/server/server.go | 6 +- 6 files changed, 2 insertions(+), 563 deletions(-) delete mode 100644 internal/clients/runclient.go delete mode 100644 internal/io/run/run.go delete mode 100644 internal/server/background/background.go delete mode 100644 internal/server/handlers/runcommand.go (limited to 'internal') diff --git a/internal/clients/runclient.go b/internal/clients/runclient.go deleted file mode 100644 index 5464d54..0000000 --- a/internal/clients/runclient.go +++ /dev/null @@ -1,87 +0,0 @@ -package clients - -import ( - "crypto/sha256" - "encoding/base64" - "encoding/hex" - "fmt" - "runtime" - "strings" - - "github.com/mimecast/dtail/internal/clients/handlers" - "github.com/mimecast/dtail/internal/io/logger" - "github.com/mimecast/dtail/internal/omode" -) - -// RunClient is a client to run various commands on the server. -type RunClient struct { - baseClient - jobName string - background string -} - -// NewRunClient returns a new run client to execute commands on the remote server. -func NewRunClient(args Args, background, jobName string) (*RunClient, error) { - args.Mode = omode.RunClient - - if jobName == "" { - jobName = hash(strings.Join(args.Arguments, " ")) - } - - c := RunClient{ - baseClient: baseClient{ - Args: args, - throttleCh: make(chan struct{}, args.ConnectionsPerCPU*runtime.NumCPU()), - retry: false, - }, - jobName: jobName, - background: background, - } - - c.init() - c.makeConnections(c) - - return &c, nil -} - -func (c RunClient) makeHandler(server string) handlers.Handler { - return handlers.NewClientHandler(server) -} - -func (c RunClient) makeCommands() (commands []string) { - if c.Timeout > 0 { - commands = append(commands, fmt.Sprintf("timeout %d run%s %s", c.Timeout, c.options(), c.What)) - return - } - - commands = append(commands, fmt.Sprintf("run%s %s", c.options(), c.What)) - return -} - -func (c RunClient) options() string { - var sb strings.Builder - - logger.Debug("options", fmt.Sprintf(":background=%s", c.background)) - sb.WriteString(fmt.Sprintf(":background=%s", c.background)) - - logger.Debug("options", fmt.Sprintf(":jobName=%s", c.jobName)) - sb.WriteString(fmt.Sprintf(":jobName=%s", c.jobName)) - - if len(c.Arguments) > 0 { - logger.Debug("options", fmt.Sprintf(":outerArgs=base64%%%s", strings.Join(c.Arguments, " "))) - sb.WriteString(fmt.Sprintf(":outerArgs=base64%%%s", encode64(strings.Join(c.Arguments, " ")))) - } - - return sb.String() -} - -func encode64(str string) string { - return base64.StdEncoding.EncodeToString([]byte(str)) -} - -func hash(str string) string { - h := sha256.New() - h.Write([]byte(str)) - - return hex.EncodeToString(h.Sum(nil)) -} diff --git a/internal/io/run/run.go b/internal/io/run/run.go deleted file mode 100644 index 2bb3756..0000000 --- a/internal/io/run/run.go +++ /dev/null @@ -1,150 +0,0 @@ -package run - -import ( - "bufio" - "context" - "io" - "os/exec" - "sync" - "syscall" - "time" - - "github.com/mimecast/dtail/internal/io/line" - "github.com/mimecast/dtail/internal/io/logger" -) - -// Run is for execute a command. -type Run struct { - command string - args []string - cmd *exec.Cmd - pgroupKilled chan struct{} -} - -// New returns a new command runner. -func New(command string, args []string) Run { - return Run{ - command: command, - args: args, - pgroupKilled: make(chan struct{}), - } -} - -// StartBackground starts running the command in background. -func (r Run) StartBackground(ctx context.Context, wg *sync.WaitGroup, ec chan<- int, lines chan<- line.Line) (pid int, err error) { - pid = -1 - - if len(r.args) > 0 { - logger.Debug(r.command, r.args, " ") - r.cmd = exec.CommandContext(ctx, r.command, r.args...) - } else { - logger.Debug(r.command) - r.cmd = exec.CommandContext(ctx, r.command) - } - - // Create a new process group, so that kill() will only kill this command + pgroup. - r.cmd.SysProcAttr = &syscall.SysProcAttr{Setpgid: true} - - stdoutPipe, myErr := r.cmd.StdoutPipe() - if err != nil { - wg.Done() - err = myErr - return - } - - stderrPipe, myErr := r.cmd.StderrPipe() - if myErr != nil { - wg.Done() - err = myErr - return - } - - if myErr := r.cmd.Start(); err != nil { - wg.Done() - err = myErr - return - } - - if r.cmd.Process != nil { - pid = r.cmd.Process.Pid - } - - commandExited := make(chan struct{}) - - var pipeWg sync.WaitGroup - pipeWg.Add(2) - - go r.killPgroup(ctx, commandExited, pid) - go r.pipeToLines(commandExited, &pipeWg, pid, stdoutPipe, "STDOUT", lines) - go r.pipeToLines(commandExited, &pipeWg, pid, stderrPipe, "STDERR", lines) - - go func() { - exitCode := 255 - if waitErr := r.cmd.Wait(); waitErr != nil { - if exitError, ok := waitErr.(*exec.ExitError); ok { - exitCode = exitError.ExitCode() - } - } - ec <- exitCode - - // Tell pipes we are done - close(commandExited) - // Wait for process group to be killed - <-r.pgroupKilled - // Wait for the pipes to flush the contents - pipeWg.Wait() - // Now the job is truly done - wg.Done() - }() - - return -} - -func (r Run) pipeToLines(commandExited chan struct{}, wg *sync.WaitGroup, pid int, reader io.Reader, what string, lines chan<- line.Line) { - defer wg.Done() - bufReader := bufio.NewReader(reader) - - for { - time.Sleep(time.Millisecond * 10) - lineStr, err := bufReader.ReadString('\n') - - if err != nil { - select { - case <-commandExited: - return - default: - } - continue - } - - newLine := line.Line{ - Content: []byte(lineStr), - Count: uint64(pid), - TransmittedPerc: 100, - SourceID: what, - } - - select { - case lines <- newLine: - case <-commandExited: - return - } - } -} - -func (r Run) killPgroup(ctx context.Context, commandExited chan struct{}, pid int) { - if pid == -1 { - close(r.pgroupKilled) - return - } - - if pgid, err := syscall.Getpgid(pid); err == nil { - // Kill process group when done - select { - case <-ctx.Done(): - case <-commandExited: - } - syscall.Kill(-pgid, syscall.SIGKILL) - close(r.pgroupKilled) - } -} diff --git a/internal/server/background/background.go b/internal/server/background/background.go deleted file mode 100644 index 3907448..0000000 --- a/internal/server/background/background.go +++ /dev/null @@ -1,126 +0,0 @@ -package background - -import ( - "context" - "errors" - "fmt" - "strings" - "sync" - - "github.com/mimecast/dtail/internal/io/logger" -) - -type job struct { - cancel context.CancelFunc - wg *sync.WaitGroup -} - -// Background specifies a job or command run in background on server side. -// This does not require an active DTail client SSH connection/session. -type Background struct { - mutex *sync.Mutex - jobs map[string]job -} - -// New returns a new background manager. -func New() Background { - return Background{ - jobs: make(map[string]job), - mutex: &sync.Mutex{}, - } -} - -// Add a background job. -func (b Background) Add(userName, jobName string, cancel context.CancelFunc, wg *sync.WaitGroup) error { - key := b.key(userName, jobName) - logger.Debug("background", "Add", key) - - b.mutex.Lock() - defer b.mutex.Unlock() - - if _, ok := b.jobs[key]; ok { - return errors.New("job already exists") - } - - b.jobs[key] = job{cancel, wg} - - // Clean up background job database. - go func() { - wg.Wait() - b.cancel(key) - }() - - return nil -} - -// Cancel a background job. -func (b Background) Cancel(userName, jobName string) error { - key := b.key(userName, jobName) - logger.Debug("background", "Cancel", key) - - return b.cancel(key) -} - -func (b Background) cancel(key string) error { - job, ok := b.get(key) - logger.Debug("background", "cancel", key, job, ok) - - if !ok { - return errors.New("no job to cancel") - } - - logger.Debug("background", "cancel", "run job.cancel()") - job.cancel() - logger.Debug("background", "cancel", "run job.wg.Wait()") - job.wg.Wait() - logger.Debug("background", "cancel", "run b.delete(key)") - b.delete(key) - - return nil -} - -// ListJobsC returns a channel listing all jobs of the given user. -func (b Background) ListJobsC(userName string) <-chan string { - logger.Debug("background", "ListJobC", userName) - - ch := make(chan string) - - go func() { - defer close(ch) - - b.mutex.Lock() - defer b.mutex.Unlock() - - for k := range b.jobs { - logger.Debug("ListJobsC", k, userName) - if strings.HasPrefix(k, fmt.Sprintf("%s.", userName)) { - ch <- k - } - } - }() - - return ch -} - -func (b Background) get(key string) (job, bool) { - logger.Debug("background", "get", key) - - b.mutex.Lock() - defer b.mutex.Unlock() - - job, ok := b.jobs[key] - return job, ok -} - -func (b Background) delete(key string) { - logger.Debug("background", "delete", key) - - b.mutex.Lock() - defer b.mutex.Unlock() - - delete(b.jobs, key) -} - -func (Background) key(userName, jobName string) string { - return fmt.Sprintf("%s.%s", userName, jobName) -} diff --git a/internal/server/handlers/runcommand.go b/internal/server/handlers/runcommand.go deleted file mode 100644 index 8e5895b..0000000 --- a/internal/server/handlers/runcommand.go +++ /dev/null @@ -1,111 +0,0 @@ -package handlers - -import ( - "context" - "errors" - "fmt" - "io/ioutil" - "os" - "os/exec" - "strings" - "sync" - "time" - - "github.com/mimecast/dtail/internal/config" - "github.com/mimecast/dtail/internal/io/logger" - "github.com/mimecast/dtail/internal/io/run" -) - -type runCommand struct { - server *ServerHandler - run run.Run -} - -func newRunCommand(server *ServerHandler) runCommand { - return runCommand{ - server: server, - } -} - -func (r runCommand) StartBackground(ctx context.Context, wg *sync.WaitGroup, argc int, args, outerArgs []string) error { - if argc < 2 { - return fmt.Errorf("%s: args:%v argc:%d", commandParseWarning, args, argc) - } - - ec := make(chan int, 1) - var pid int - var err error - - command := strings.Join(args[1:], " ") - if strings.Contains(command, ";") || strings.Contains(command, "\n") { - if pid, err = r.startScript(ctx, wg, ec, command, outerArgs); err != nil { - r.server.sendServerMessage(".run exitstatus 255") - return err - } - return nil - } - - if pid, err = r.start(ctx, wg, ec, strings.TrimSpace(command), outerArgs); err != nil { - r.server.sendServerMessage(".run exitstatus 255") - return err - } - - exitCode := <-ec - r.server.sendServerMessage(fmt.Sprintf(".run exitstatus %d", exitCode)) - r.server.sendServerMessage(logger.Info(fmt.Sprintf("Process %d exited with status %d", pid, exitCode))) - - return nil -} - -func (r runCommand) startScript(ctx context.Context, wg *sync.WaitGroup, ec chan<- int, script string, outerArgs []string) (int, error) { - if _, err := os.Stat(config.Common.TmpDir); os.IsNotExist(err) { - return -1, err - } - - timestamp := time.Now().UnixNano() - scriptPath := fmt.Sprintf("%s/%s_%v.sh", config.Common.TmpDir, r.server.user.Name, timestamp) - - // TODO: On dserver startup delete all previously written scripts (there might be left overs due to a crash or so) - logger.Debug(r.server.user, "Writing temp script", scriptPath) - - script = fmt.Sprintf("#!/bin/sh\n%s", script) - if err := ioutil.WriteFile(scriptPath, []byte(script), 0700); err != nil { - return -1, err - } - - pid, err := r.start(ctx, wg, ec, scriptPath, outerArgs) - go func() { - wg.Wait() - logger.Debug("Deleting script", scriptPath) - os.Remove(scriptPath) - }() - - return pid, err -} - -func (r runCommand) start(ctx context.Context, wg *sync.WaitGroup, ec chan<- int, command string, outerArgs []string) (int, error) { - if len(command) == 0 { - return -1, errors.New("Empty command provided") - } - - splitted := strings.Split(command, " ") - path := splitted[0] - args := splitted[1:] - args = append(args, outerArgs...) - - qualifiedPath, err := exec.LookPath(path) - if err != nil { - return -1, err - } - - if !r.server.user.HasFilePermission(qualifiedPath, "runcommands") { - return -1, fmt.Errorf("No permission to execute path: %s", qualifiedPath) - } - - r.run = run.New(qualifiedPath, args) - pid, err := r.run.StartBackground(ctx, wg, ec, r.server.lines) - if err != nil { - return pid, err - } - return pid, nil -} diff --git a/internal/server/handlers/serverhandler.go b/internal/server/handlers/serverhandler.go index 164a280..7ad1224 100644 --- a/internal/server/handlers/serverhandler.go +++ b/internal/server/handlers/serverhandler.go @@ -9,7 +9,6 @@ import ( "os" "strconv" "strings" - "sync" "sync/atomic" "time" @@ -19,7 +18,6 @@ import ( "github.com/mimecast/dtail/internal/io/logger" "github.com/mimecast/dtail/internal/mapr/server" "github.com/mimecast/dtail/internal/omode" - "github.com/mimecast/dtail/internal/server/background" user "github.com/mimecast/dtail/internal/user/server" "github.com/mimecast/dtail/internal/version" ) @@ -47,11 +45,10 @@ type ServerHandler struct { ackCloseReceived chan struct{} activeCommands int32 activeReaders int32 - background background.Background } // NewServerHandler returns the server handler. -func NewServerHandler(user *user.User, catLimiter, tailLimiter, globalServerWaitFor chan struct{}, background background.Background) *ServerHandler { +func NewServerHandler(user *user.User, catLimiter, tailLimiter, globalServerWaitFor chan struct{}) *ServerHandler { h := ServerHandler{ done: internal.NewDone(), lines: make(chan line.Line, 100), @@ -63,7 +60,6 @@ func NewServerHandler(user *user.User, catLimiter, tailLimiter, globalServerWait globalServerWaitFor: globalServerWaitFor, regex: ".", user: user, - background: background, } fqdn, err := os.Hostname() @@ -314,85 +310,6 @@ func (h *ServerHandler) handleUserCommand(ctx context.Context, argc int, args [] commandFinished() }() - case "run": - // TODO: Refactor this "run" case, move code to runcommand.go - command := newRunCommand(h) - - jobName, _ := options["jobName"] - logger.Debug(h.user, "run", options) - - if val, ok := options["background"]; ok && (val == "cancel" || val == "stop") { - if err := h.background.Cancel(h.user.Name, jobName); err != nil { - h.sendServerMessage(logger.Error(h.user, err, jobName, args)) - } else { - h.sendServerMessage(logger.Info(h.user, "job cancelled", jobName)) - } - commandFinished() - return - } - - if val, ok := options["background"]; ok && val == "list" { - h.sendServerMessage("Listing jobs") - count := 0 - for jobName := range h.background.ListJobsC(h.user.Name) { - h.sendServerMessage(jobName) - count++ - } - h.sendServerMessage(fmt.Sprintf("Found %d jobs", count)) - commandFinished() - return - } - - str, _ := options["outerArgs"] - outerArgs := strings.Split(str, " ") - - var background bool - if val, ok := options["background"]; ok && val == "start" { - background = true - } - - var wg sync.WaitGroup - wg.Add(1) - - if background { - if timeout == 0 { - // Set default background timeout. - timeout = time.Hour * 1 - } - - commandCtx, cancel := context.WithTimeout(ctx, timeout) - - if err := h.background.Add(h.user.Name, jobName, cancel, &wg); err != nil { - h.sendServerMessage(logger.Error(h.user, err, jobName, args)) - commandFinished() - return - } - ctx = commandCtx - } - - if err := command.StartBackground(ctx, &wg, argc, args, outerArgs); err != nil { - h.sendServerMessage(logger.Error(h.user, "Unable to execute command", argc, args, err)) - commandFinished() - return - } - - // Make sure that server waits for all sub-processes to finish on shutdown - go func() { h.globalServerWaitFor <- struct{}{} }() - go func() { - wg.Wait() - <-h.globalServerWaitFor - }() - - if background { - h.sendServerMessage(logger.Info(h.user, jobName, "job started in background")) - commandFinished() - return - } - - // Command run in foreground, wait for it to complete before finishing the connection. - wg.Wait() - commandFinished() - case "ack", ".ack": h.handleAckCommand(argc, args) commandFinished() diff --git a/internal/server/server.go b/internal/server/server.go index 5e2a521..d4255a3 100644 --- a/internal/server/server.go +++ b/internal/server/server.go @@ -11,7 +11,6 @@ import ( "github.com/mimecast/dtail/internal/config" "github.com/mimecast/dtail/internal/io/logger" - "github.com/mimecast/dtail/internal/server/background" "github.com/mimecast/dtail/internal/server/handlers" "github.com/mimecast/dtail/internal/ssh/server" user "github.com/mimecast/dtail/internal/user/server" @@ -36,8 +35,6 @@ type Server struct { cont *continuous // Wait counter, e.g. there might be still subprocesses (forked by drun) to be killed. shutdownWaitFor chan struct{} - // Background jobs - background background.Background } // New returns a new server. @@ -51,7 +48,6 @@ func New() *Server { shutdownWaitFor: make(chan struct{}, 1000), sched: newScheduler(), cont: newContinuous(), - background: background.New(), } s.sshServerConfig.PasswordCallback = s.Callback @@ -183,7 +179,7 @@ func (s *Server) handleRequests(ctx context.Context, sshConn gossh.Conn, in <-ch case config.ControlUser: handler = handlers.NewControlHandler(user) default: - handler = handlers.NewServerHandler(user, s.catLimiter, s.tailLimiter, s.shutdownWaitFor, s.background) + handler = handlers.NewServerHandler(user, s.catLimiter, s.tailLimiter, s.shutdownWaitFor) } terminate := func() { -- cgit v1.2.3 From 1626b6237d87a6ae6238b0d297d7bfb5cf7ddd38 Mon Sep 17 00:00:00 2001 From: Paul Buetow Date: Mon, 16 Nov 2020 22:16:10 +0000 Subject: strip out timeout which was used for background command scheduling --- internal/server/handlers/serverhandler.go | 34 ++----------------------------- 1 file changed, 2 insertions(+), 32 deletions(-) (limited to 'internal') diff --git a/internal/server/handlers/serverhandler.go b/internal/server/handlers/serverhandler.go index 7ad1224..843eabc 100644 --- a/internal/server/handlers/serverhandler.go +++ b/internal/server/handlers/serverhandler.go @@ -7,7 +7,6 @@ import ( "fmt" "io" "os" - "strconv" "strings" "sync/atomic" "time" @@ -147,7 +146,6 @@ func (h *ServerHandler) Write(p []byte) (n int, err error) { func (h *ServerHandler) handleCommand(commandStr string) { logger.Debug(h.user, commandStr) - var timeout time.Duration ctx := context.Background() args, argc, err := h.handleProtocolVersion(strings.Split(commandStr, " ")) @@ -162,12 +160,6 @@ func (h *ServerHandler) handleCommand(commandStr string) { return } - args, argc, timeout, err = h.handleTimeout(args, argc) - if err != nil { - h.send(h.serverMessages, logger.Error(h.user, err)) - return - } - if h.user.Name == config.ControlUser { h.handleControlCommand(argc, args) return @@ -179,19 +171,7 @@ func (h *ServerHandler) handleCommand(commandStr string) { cancel() }() - if timeout > 0 { - logger.Info(h.user, "Command with timeout context", argc, args, timeout) - ctx, cancel := context.WithTimeout(ctx, timeout) - go func() { - <-ctx.Done() - logger.Info(h.user, "Command timed out, canceling it", args, args, timeout) - cancel() - }() - h.handleUserCommand(ctx, argc, args, timeout) - return - } - - h.handleUserCommand(ctx, argc, args, timeout) + h.handleUserCommand(ctx, argc, args) } func (h *ServerHandler) handleProtocolVersion(args []string) ([]string, int, error) { @@ -229,16 +209,6 @@ func (h *ServerHandler) handleBase64(args []string, argc int) ([]string, int, er return args, argc, nil } -func (h *ServerHandler) handleTimeout(args []string, argc int) ([]string, int, time.Duration, error) { - if argc <= 2 || args[0] != "timeout" { - // No timeout specified - return args, argc, time.Duration(0) * time.Second, nil - } - - timeout, err := strconv.Atoi(args[1]) - return args[2:], argc - 2, time.Duration(timeout) * time.Second, err -} - func (h *ServerHandler) handleControlCommand(argc int, args []string) { switch args[0] { case "debug": @@ -248,7 +218,7 @@ func (h *ServerHandler) handleControlCommand(argc int, args []string) { } } -func (h *ServerHandler) handleUserCommand(ctx context.Context, argc int, args []string, timeout time.Duration) { +func (h *ServerHandler) handleUserCommand(ctx context.Context, argc int, args []string) { logger.Debug(h.user, "handleUserCommand", argc, args) h.incrementActiveCommands() -- cgit v1.2.3 From 1bc1725bac7cda1f9bc3fcbcf02bf4e6413cf058 Mon Sep 17 00:00:00 2001 From: Paul Buetow Date: Sun, 22 Nov 2020 09:21:44 +0000 Subject: add dockerfile for building a dtail server dev/test container --- internal/server/server.go | 1 + 1 file changed, 1 insertion(+) (limited to 'internal') diff --git a/internal/server/server.go b/internal/server/server.go index d4255a3..31fa85d 100644 --- a/internal/server/server.go +++ b/internal/server/server.go @@ -34,6 +34,7 @@ type Server struct { // Mointor log files for pattern (if configured) cont *continuous // Wait counter, e.g. there might be still subprocesses (forked by drun) to be killed. + // TODO: Remove this counter. shutdownWaitFor chan struct{} } -- cgit v1.2.3 From 97827c2247498eefe0263f043a0cd6eca962077d Mon Sep 17 00:00:00 2001 From: Paul Buetow Date: Sun, 22 Nov 2020 10:29:40 +0000 Subject: add comments --- internal/done.go | 4 ++++ 1 file changed, 4 insertions(+) (limited to 'internal') diff --git a/internal/done.go b/internal/done.go index 2326eee..54e5e8e 100644 --- a/internal/done.go +++ b/internal/done.go @@ -4,21 +4,25 @@ import ( "sync" ) +// Done is a cleanup/shutdown helper. type Done struct { ch chan struct{} mutex sync.Mutex } +// NewDone returns a new cleanup/shutdown helper. func NewDone() *Done { return &Done{ ch: make(chan struct{}), } } +// Done returns the done channel (closed when done) func (d *Done) Done() <-chan struct{} { return d.ch } +// Shutdown closes the done channel. It can be called multiple times. func (d *Done) Shutdown() { d.mutex.Lock() defer d.mutex.Unlock() -- cgit v1.2.3 From ae8ffc84331ca72f0d33fff69edd85d6e03d29ae Mon Sep 17 00:00:00 2001 From: Paul Buetow Date: Mon, 28 Dec 2020 09:40:38 +0000 Subject: refactor --- internal/server/handlers/serverhandler.go | 50 +++++++++++++++---------------- internal/server/server.go | 20 +------------ 2 files changed, 25 insertions(+), 45 deletions(-) (limited to 'internal') diff --git a/internal/server/handlers/serverhandler.go b/internal/server/handlers/serverhandler.go index 843eabc..db917bd 100644 --- a/internal/server/handlers/serverhandler.go +++ b/internal/server/handlers/serverhandler.go @@ -29,36 +29,34 @@ const ( // the Bi-directional communication between SSH client and server. // This handler implements the handler of the SSH server. type ServerHandler struct { - done *internal.Done - lines chan line.Line - regex string - aggregate *server.Aggregate - aggregatedMessages chan string - serverMessages chan string - payload []byte - hostname string - user *user.User - catLimiter chan struct{} - tailLimiter chan struct{} - globalServerWaitFor chan struct{} - ackCloseReceived chan struct{} - activeCommands int32 - activeReaders int32 + done *internal.Done + lines chan line.Line + regex string + aggregate *server.Aggregate + aggregatedMessages chan string + serverMessages chan string + payload []byte + hostname string + user *user.User + catLimiter chan struct{} + tailLimiter chan struct{} + ackCloseReceived chan struct{} + activeCommands int32 + activeReaders int32 } // NewServerHandler returns the server handler. -func NewServerHandler(user *user.User, catLimiter, tailLimiter, globalServerWaitFor chan struct{}) *ServerHandler { +func NewServerHandler(user *user.User, catLimiter, tailLimiter chan struct{}) *ServerHandler { h := ServerHandler{ - done: internal.NewDone(), - lines: make(chan line.Line, 100), - serverMessages: make(chan string, 10), - aggregatedMessages: make(chan string, 10), - ackCloseReceived: make(chan struct{}), - catLimiter: catLimiter, - tailLimiter: tailLimiter, - globalServerWaitFor: globalServerWaitFor, - regex: ".", - user: user, + done: internal.NewDone(), + lines: make(chan line.Line, 100), + serverMessages: make(chan string, 10), + aggregatedMessages: make(chan string, 10), + ackCloseReceived: make(chan struct{}), + catLimiter: catLimiter, + tailLimiter: tailLimiter, + regex: ".", + user: user, } fqdn, err := os.Hostname() diff --git a/internal/server/server.go b/internal/server/server.go index 31fa85d..d8d43c9 100644 --- a/internal/server/server.go +++ b/internal/server/server.go @@ -7,7 +7,6 @@ import ( "io" "net" "strings" - "time" "github.com/mimecast/dtail/internal/config" "github.com/mimecast/dtail/internal/io/logger" @@ -33,9 +32,6 @@ type Server struct { sched *scheduler // Mointor log files for pattern (if configured) cont *continuous - // Wait counter, e.g. there might be still subprocesses (forked by drun) to be killed. - // TODO: Remove this counter. - shutdownWaitFor chan struct{} } // New returns a new server. @@ -46,7 +42,6 @@ func New() *Server { sshServerConfig: &gossh.ServerConfig{}, catLimiter: make(chan struct{}, config.Server.MaxConcurrentCats), tailLimiter: make(chan struct{}, config.Server.MaxConcurrentTails), - shutdownWaitFor: make(chan struct{}, 1000), sched: newScheduler(), cont: newContinuous(), } @@ -82,25 +77,12 @@ func (s *Server) Start(ctx context.Context) int { select { case <-ctx.Done(): - // Wait until all commands/jobs/children are no more! - s.wait() } // For future use. return 0 } -func (s *Server) wait() { - for { - num := len(s.shutdownWaitFor) - logger.Debug("Waiting for stuff to finish", num) - if num <= 0 { - return - } - time.Sleep(time.Second) - } -} - func (s *Server) listenerLoop(ctx context.Context, listener net.Listener) { logger.Debug("Starting listener loop") @@ -180,7 +162,7 @@ func (s *Server) handleRequests(ctx context.Context, sshConn gossh.Conn, in <-ch case config.ControlUser: handler = handlers.NewControlHandler(user) default: - handler = handlers.NewServerHandler(user, s.catLimiter, s.tailLimiter, s.shutdownWaitFor) + handler = handlers.NewServerHandler(user, s.catLimiter, s.tailLimiter) } terminate := func() { -- cgit v1.2.3 From 0099a7ab9e1d28300c69c3b50b4ebe1cde9a8cbc Mon Sep 17 00:00:00 2001 From: Paul Buetow Date: Tue, 29 Dec 2020 08:34:04 +0000 Subject: Make Linux ACL support optional, as it requires CGo and makes the binary