From c2522ffb59514443816a96386c16bb7527cbe57c Mon Sep 17 00:00:00 2001 From: Paul Buetow Date: Sat, 21 Aug 2021 14:54:24 +0300 Subject: read files bytewise for more control of whats happening - change transport protocol for more control over newlines --- internal/clients/handlers/basehandler.go | 6 +++--- internal/clients/handlers/maprhandler.go | 3 ++- 2 files changed, 5 insertions(+), 4 deletions(-) (limited to 'internal/clients') diff --git a/internal/clients/handlers/basehandler.go b/internal/clients/handlers/basehandler.go index f07fd90..acafe0e 100644 --- a/internal/clients/handlers/basehandler.go +++ b/internal/clients/handlers/basehandler.go @@ -9,7 +9,7 @@ import ( "github.com/mimecast/dtail/internal" "github.com/mimecast/dtail/internal/io/logger" - "github.com/mimecast/dtail/internal/version" + "github.com/mimecast/dtail/internal/protocol" ) type baseHandler struct { @@ -43,7 +43,7 @@ func (h *baseHandler) SendMessage(command string) error { logger.Debug("Sending command", h.server, command, encoded) select { - case h.commands <- fmt.Sprintf("protocol %s base64 %v;", version.ProtocolCompat, encoded): + case h.commands <- fmt.Sprintf("protocol %s base64 %v;", protocol.ProtocolCompat, encoded): case <-time.After(time.Second * 5): return fmt.Errorf("Timed out sending command '%s' (base64: '%s')", command, encoded) case <-h.Done(): @@ -57,7 +57,7 @@ func (h *baseHandler) SendMessage(command string) error { func (h *baseHandler) Write(p []byte) (n int, err error) { for _, b := range p { h.receiveBuf = append(h.receiveBuf, b) - if b == '\n' { + if b == protocol.MessageDelimiter { if len(h.receiveBuf) == 0 { continue } diff --git a/internal/clients/handlers/maprhandler.go b/internal/clients/handlers/maprhandler.go index fb71c8f..7ac5895 100644 --- a/internal/clients/handlers/maprhandler.go +++ b/internal/clients/handlers/maprhandler.go @@ -7,6 +7,7 @@ import ( "github.com/mimecast/dtail/internal/io/logger" "github.com/mimecast/dtail/internal/mapr" "github.com/mimecast/dtail/internal/mapr/client" + "github.com/mimecast/dtail/internal/protocol" ) // MaprHandler is the handler used on the client side for running mapreduce aggregations. @@ -58,7 +59,7 @@ func (h *MaprHandler) Write(p []byte) (n int, err error) { // related data. func (h *MaprHandler) handleAggregateMessage(message string) { h.count++ - parts := strings.Split(message, "➔") + parts := strings.Split(message, protocol.AggregateDelimiter) // Index 0 contains 'AGGREGATE', 1 contains server host. // Aggregation data begins from index 2. -- cgit v1.2.3 From 9883a190109623b64e6d311dc2b462a6eae68003 Mon Sep 17 00:00:00 2001 From: Paul Buetow Date: Sun, 22 Aug 2021 10:07:00 +0300 Subject: introduces the protocol package --- internal/clients/handlers/basehandler.go | 3 ++- internal/clients/handlers/healthhandler.go | 3 ++- internal/clients/handlers/maprhandler.go | 2 +- internal/clients/maprclient.go | 3 ++- 4 files changed, 7 insertions(+), 4 deletions(-) (limited to 'internal/clients') diff --git a/internal/clients/handlers/basehandler.go b/internal/clients/handlers/basehandler.go index acafe0e..fe83faa 100644 --- a/internal/clients/handlers/basehandler.go +++ b/internal/clients/handlers/basehandler.go @@ -56,13 +56,14 @@ func (h *baseHandler) SendMessage(command string) error { // Read data from the dtail server via Writer interface. func (h *baseHandler) Write(p []byte) (n int, err error) { for _, b := range p { - h.receiveBuf = append(h.receiveBuf, b) if b == protocol.MessageDelimiter { if len(h.receiveBuf) == 0 { continue } message := string(h.receiveBuf) h.handleMessageType(message) + } else { + h.receiveBuf = append(h.receiveBuf, b) } } diff --git a/internal/clients/handlers/healthhandler.go b/internal/clients/handlers/healthhandler.go index 0440706..213748c 100644 --- a/internal/clients/handlers/healthhandler.go +++ b/internal/clients/handlers/healthhandler.go @@ -6,6 +6,7 @@ import ( "time" "github.com/mimecast/dtail/internal" + "github.com/mimecast/dtail/internal/protocol" ) // HealthHandler implements the handler required for health checks. @@ -72,7 +73,7 @@ func (h *HealthHandler) SendMessage(command string) error { func (h *HealthHandler) Write(p []byte) (n int, err error) { for _, b := range p { h.receiveBuf = append(h.receiveBuf, b) - if b == '\n' { + if b == protocol.MessageDelimiter { // '\n' { h.receive <- string(h.receiveBuf) h.receiveBuf = h.receiveBuf[:0] } diff --git a/internal/clients/handlers/maprhandler.go b/internal/clients/handlers/maprhandler.go index 7ac5895..afad507 100644 --- a/internal/clients/handlers/maprhandler.go +++ b/internal/clients/handlers/maprhandler.go @@ -37,7 +37,7 @@ func NewMaprHandler(server string, query *mapr.Query, globalGroup *mapr.GlobalGr func (h *MaprHandler) Write(p []byte) (n int, err error) { for _, b := range p { h.baseHandler.receiveBuf = append(h.baseHandler.receiveBuf, b) - if b == '\n' { + if b == protocol.MessageDelimiter { // '\n' { if len(h.baseHandler.receiveBuf) == 0 { continue } diff --git a/internal/clients/maprclient.go b/internal/clients/maprclient.go index 1c0c2cc..77b674b 100644 --- a/internal/clients/maprclient.go +++ b/internal/clients/maprclient.go @@ -170,7 +170,8 @@ func (c *MaprClient) printResults() { return } - logger.Raw(fmt.Sprintf("%s\n", c.query.RawQuery)) + //logger.Raw(fmt.Sprintf("%s\n", c.query.RawQuery)) + logger.Raw(c.query.RawQuery) logger.Raw(result) } -- cgit v1.2.3 From 23982f331c2154a66b86d596226c24454fd06be5 Mon Sep 17 00:00:00 2001 From: Paul Buetow Date: Sat, 28 Aug 2021 20:26:32 +0100 Subject: 1. Major performance gain by not checking for file truncation aftter each bytes read. 2. Introduce field separator to the protocol package. --- internal/clients/healthclient.go | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) (limited to 'internal/clients') diff --git a/internal/clients/healthclient.go b/internal/clients/healthclient.go index e93f6be..692464c 100644 --- a/internal/clients/healthclient.go +++ b/internal/clients/healthclient.go @@ -11,6 +11,7 @@ import ( "github.com/mimecast/dtail/internal/clients/remote" "github.com/mimecast/dtail/internal/config" "github.com/mimecast/dtail/internal/omode" + "github.com/mimecast/dtail/internal/protocol" gossh "golang.org/x/crypto/ssh" ) @@ -57,7 +58,7 @@ func (c *HealthClient) Start(ctx context.Context) (status int) { select { case data := <-receive: // Parse recieved data. - s := strings.Split(data, "|") + s := strings.Split(data, protocol.FieldDelimiter) message := s[len(s)-1] if strings.HasPrefix(message, "done;") { return -- cgit v1.2.3 From 2c1c70313bb03cf2b2d7e7afadb07a48ff6bb690 Mon Sep 17 00:00:00 2001 From: Paul Buetow Date: Mon, 6 Sep 2021 09:22:21 +0300 Subject: REMOTE and CLIENT colors are brushed correctly too now --- internal/clients/stats.go | 14 ++++++++++---- 1 file changed, 10 insertions(+), 4 deletions(-) (limited to 'internal/clients') diff --git a/internal/clients/stats.go b/internal/clients/stats.go index d8163d4..3f76e0f 100644 --- a/internal/clients/stats.go +++ b/internal/clients/stats.go @@ -8,6 +8,7 @@ import ( "sync" "time" + "github.com/mimecast/dtail/internal/color" "github.com/mimecast/dtail/internal/config" "github.com/mimecast/dtail/internal/io/logger" ) @@ -54,7 +55,6 @@ func (s *stats) Start(ctx context.Context, throttleCh <-chan struct{}, statsCh < throttle := len(throttleCh) newConnections := connected - connectedLast - if (connected == connectedLast || quiet) && !force { continue } @@ -77,7 +77,15 @@ func (s *stats) Start(ctx context.Context, throttleCh <-chan struct{}, statsCh < func (s *stats) printStatsDueInterrupt(messages []string) { logger.Pause() - for _, message := range messages { + for i, message := range messages { + if i > 0 && config.Client.TermColorsEnable { + fmt.Println(color.PaintStrWithAttr(message, + config.Client.TermColors.Client.TextFg, + config.Client.TermColors.Client.TextBg, + config.Client.TermColors.Client.TextAttr, + )) + continue + } fmt.Println(fmt.Sprintf(" %s", message)) } time.Sleep(time.Second * time.Duration(config.InterruptTimeoutS)) @@ -99,7 +107,6 @@ func (s *stats) statsLine(connected, newConnections int, throttle int) string { func (s *stats) numConnected() int { s.mutex.Lock() defer s.mutex.Unlock() - return s.connected } @@ -107,6 +114,5 @@ func percentOf(total float64, value float64) float64 { if total == 0 || total == value { return 100 } - return value / (total / 100.0) } -- cgit v1.2.3 From cc89d3fb8be2465af276d7ef03ea2a8affd87b2e Mon Sep 17 00:00:00 2001 From: Paul Buetow Date: Mon, 6 Sep 2021 13:48:55 +0300 Subject: Print out client/server update notice even from dtail server 4 to dtail client 3. --- internal/clients/handlers/basehandler.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) (limited to 'internal/clients') diff --git a/internal/clients/handlers/basehandler.go b/internal/clients/handlers/basehandler.go index fe83faa..63ceaac 100644 --- a/internal/clients/handlers/basehandler.go +++ b/internal/clients/handlers/basehandler.go @@ -56,7 +56,7 @@ func (h *baseHandler) SendMessage(command string) error { // Read data from the dtail server via Writer interface. func (h *baseHandler) Write(p []byte) (n int, err error) { for _, b := range p { - if b == protocol.MessageDelimiter { + if b == protocol.MessageDelimiter || b == '\n' { if len(h.receiveBuf) == 0 { continue } -- cgit v1.2.3 From 6ae75e8f106d3eee18ea61e6c4d6925c6f514460 Mon Sep 17 00:00:00 2001 From: Paul Buetow Date: Mon, 6 Sep 2021 14:07:07 +0300 Subject: fine tweak color schema --- internal/clients/stats.go | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) (limited to 'internal/clients') diff --git a/internal/clients/stats.go b/internal/clients/stats.go index 3f76e0f..faeb9fb 100644 --- a/internal/clients/stats.go +++ b/internal/clients/stats.go @@ -80,9 +80,9 @@ func (s *stats) printStatsDueInterrupt(messages []string) { for i, message := range messages { if i > 0 && config.Client.TermColorsEnable { fmt.Println(color.PaintStrWithAttr(message, - config.Client.TermColors.Client.TextFg, - config.Client.TermColors.Client.TextBg, - config.Client.TermColors.Client.TextAttr, + config.Client.TermColors.Client.ClientFg, + config.Client.TermColors.Client.ClientBg, + config.Client.TermColors.Client.ClientAttr, )) continue } -- cgit v1.2.3 From f74a9e4b35feb8c07d8a70b5a581088a0a59889d Mon Sep 17 00:00:00 2001 From: Paul Buetow Date: Tue, 7 Sep 2021 10:01:32 +0300 Subject: Produce MAPREDUCE lines, can aggregate these via default log format --- internal/clients/stats.go | 49 ++++++++++++++++++++++++++++++++++------------- 1 file changed, 36 insertions(+), 13 deletions(-) (limited to 'internal/clients') diff --git a/internal/clients/stats.go b/internal/clients/stats.go index faeb9fb..6da443c 100644 --- a/internal/clients/stats.go +++ b/internal/clients/stats.go @@ -11,12 +11,13 @@ import ( "github.com/mimecast/dtail/internal/color" "github.com/mimecast/dtail/internal/config" "github.com/mimecast/dtail/internal/io/logger" + "github.com/mimecast/dtail/internal/protocol" ) // Used to collect and display various client stats. type stats struct { // Total amount servers to connect to. - connectionsTotal int + servers int // To keep track of what connected and disconnected connectionsEstCh chan struct{} // Amount of servers connections are established. @@ -25,10 +26,10 @@ type stats struct { mutex sync.Mutex } -func newTailStats(connectionsTotal int) *stats { +func newTailStats(servers int) *stats { return &stats{ - connectionsTotal: connectionsTotal, - connectionsEstCh: make(chan struct{}, connectionsTotal), + servers: servers, + connectionsEstCh: make(chan struct{}, servers), connected: 0, } } @@ -59,13 +60,14 @@ func (s *stats) Start(ctx context.Context, throttleCh <-chan struct{}, statsCh < continue } - stats := s.statsLine(connected, newConnections, throttle) switch force { case true: + stats := s.statsLine(connected, newConnections, throttle) messages = append(messages, fmt.Sprintf("Connection stats: %s", stats)) s.printStatsDueInterrupt(messages) default: - logger.Info(stats) + data := s.statsData(connected, newConnections, throttle) + logger.Mapreduce("STATS", data) } connectedLast = connected @@ -92,16 +94,37 @@ func (s *stats) printStatsDueInterrupt(messages []string) { logger.Resume() } +func (s *stats) statsData(connected, newConnections int, throttle int) map[string]interface{} { + percConnected := percentOf(float64(s.servers), float64(connected)) + + data := make(map[string]interface{}) + data["connected"] = connected + data["servers"] = s.servers + data["connected%"] = int(percConnected) + data["new"] = newConnections + data["throttle"] = throttle + data["goroutines"] = runtime.NumGoroutine() + data["cgocalls"] = runtime.NumCgoCall() + data["cpu"] = runtime.NumCPU() + + return data +} + func (s *stats) statsLine(connected, newConnections int, throttle int) string { - percConnected := percentOf(float64(s.connectionsTotal), float64(connected)) + sb := strings.Builder{} - 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())) + i := 0 + for k, v := range s.statsData(connected, newConnections, throttle) { + if i > 0 { + sb.WriteString(protocol.FieldDelimiter) + } + sb.WriteString(k) + sb.WriteByte('=') + sb.WriteString(fmt.Sprintf("%v", v)) + i++ + } - return strings.Join(stats, "|") + return sb.String() } func (s *stats) numConnected() int { -- cgit v1.2.3 From 16dc57e1e1c28e9d762424e596223a980770e059 Mon Sep 17 00:00:00 2001 From: Paul Buetow Date: Wed, 8 Sep 2021 19:10:50 +0300 Subject: mapreduce tables are in colors now too --- internal/clients/baseclient.go | 1 + internal/clients/handlers/basehandler.go | 29 +++++++++++++-------- internal/clients/handlers/healthhandler.go | 11 ++++---- internal/clients/handlers/maprhandler.go | 42 +++++++++++++++--------------- internal/clients/maprclient.go | 31 +++++++++++++++++----- 5 files changed, 70 insertions(+), 44 deletions(-) (limited to 'internal/clients') diff --git a/internal/clients/baseclient.go b/internal/clients/baseclient.go index f20156f..de0c101 100644 --- a/internal/clients/baseclient.go +++ b/internal/clients/baseclient.go @@ -71,6 +71,7 @@ func (c *baseClient) Start(ctx context.Context, statsCh <-chan string) (status i go c.hostKeyCallback.PromptAddHosts(ctx) // Print client stats every time something on statsCh is recieved. go c.stats.Start(ctx, c.throttleCh, statsCh, c.Args.Quiet) + // Keep count of active connections active := make(chan struct{}, len(c.connections)) diff --git a/internal/clients/handlers/basehandler.go b/internal/clients/handlers/basehandler.go index 63ceaac..51f33c1 100644 --- a/internal/clients/handlers/basehandler.go +++ b/internal/clients/handlers/basehandler.go @@ -1,6 +1,7 @@ package handlers import ( + "bytes" "encoding/base64" "fmt" "io" @@ -17,7 +18,7 @@ type baseHandler struct { server string shellStarted bool commands chan string - receiveBuf []byte + receiveBuf bytes.Buffer status int } @@ -56,14 +57,23 @@ func (h *baseHandler) SendMessage(command string) error { // Read data from the dtail server via Writer interface. func (h *baseHandler) Write(p []byte) (n int, err error) { for _, b := range p { - if b == protocol.MessageDelimiter || b == '\n' { - if len(h.receiveBuf) == 0 { + switch b { + /* + // TODO: Next DTail version make it so that '\n' gets ignored. For now + // leave it for compatibility with older DTail server + ability to display + // the protocol mismatch warn message. + case '\n' { + continue + */ + case '\n', protocol.MessageDelimiter: + message := h.receiveBuf.String() + if len(message) == 0 { continue } - message := string(h.receiveBuf) h.handleMessageType(message) - } else { - h.receiveBuf = append(h.receiveBuf, b) + h.receiveBuf.Reset() + default: + h.receiveBuf.WriteByte(b) } } @@ -78,25 +88,22 @@ func (h *baseHandler) Read(p []byte) (n int, err error) { case <-h.Done(): return 0, io.EOF } - return } // Handle various message types. func (h *baseHandler) handleMessageType(message string) { - if len(h.receiveBuf) == 0 { + if len(message) == 0 { return } // Hidden server commands starti with a dot "." - if h.receiveBuf[0] == '.' { + if message[0] == '.' { h.handleHiddenMessage(message) - h.receiveBuf = h.receiveBuf[:0] return } logger.Raw(message) - h.receiveBuf = h.receiveBuf[:0] } // Handle messages received from server which are not meant to be displayed diff --git a/internal/clients/handlers/healthhandler.go b/internal/clients/handlers/healthhandler.go index 213748c..eca0348 100644 --- a/internal/clients/handlers/healthhandler.go +++ b/internal/clients/handlers/healthhandler.go @@ -1,6 +1,7 @@ package handlers import ( + "bytes" "errors" "fmt" "time" @@ -13,7 +14,7 @@ import ( type HealthHandler struct { done *internal.Done // Buffer of incoming data from server. - receiveBuf []byte + receiveBuf bytes.Buffer // To send commands to the server. commands chan string // To receive messages from the server. @@ -72,10 +73,10 @@ func (h *HealthHandler) SendMessage(command string) error { // Server writes byte stream to client. func (h *HealthHandler) Write(p []byte) (n int, err error) { for _, b := range p { - h.receiveBuf = append(h.receiveBuf, b) - if b == protocol.MessageDelimiter { // '\n' { - h.receive <- string(h.receiveBuf) - h.receiveBuf = h.receiveBuf[:0] + h.receiveBuf.WriteByte(b) + if b == protocol.MessageDelimiter { + h.receive <- h.receiveBuf.String() + h.receiveBuf.Reset() } } diff --git a/internal/clients/handlers/maprhandler.go b/internal/clients/handlers/maprhandler.go index afad507..65b1454 100644 --- a/internal/clients/handlers/maprhandler.go +++ b/internal/clients/handlers/maprhandler.go @@ -15,7 +15,6 @@ type MaprHandler struct { baseHandler aggregate *client.Aggregate query *mapr.Query - count uint64 } // NewMaprHandler returns a new mapreduce client handler. @@ -36,19 +35,20 @@ func NewMaprHandler(server string, query *mapr.Query, globalGroup *mapr.GlobalGr // Read data from the dtail server via Writer interface. func (h *MaprHandler) Write(p []byte) (n int, err error) { for _, b := range p { - h.baseHandler.receiveBuf = append(h.baseHandler.receiveBuf, b) - if b == protocol.MessageDelimiter { // '\n' { - if len(h.baseHandler.receiveBuf) == 0 { - continue + switch b { + case '\n': + continue + case protocol.MessageDelimiter: + message := h.baseHandler.receiveBuf.String() + logger.Debug(message) + if message[0] == 'A' { + h.handleAggregateMessage(message) + } else { + h.baseHandler.handleMessageType(message) } - message := string(h.baseHandler.receiveBuf) - - if h.baseHandler.receiveBuf[0] == 'A' { - h.handleAggregateMessage(strings.TrimSpace(message)) - h.baseHandler.receiveBuf = h.baseHandler.receiveBuf[:0] - continue - } - h.baseHandler.handleMessageType(message) + h.baseHandler.receiveBuf.Reset() + default: + h.baseHandler.receiveBuf.WriteByte(b) } } @@ -58,12 +58,12 @@ func (h *MaprHandler) Write(p []byte) (n int, err error) { // Handle a message received from server including mapr aggregation // related data. func (h *MaprHandler) handleAggregateMessage(message string) { - h.count++ - parts := strings.Split(message, protocol.AggregateDelimiter) - - // Index 0 contains 'AGGREGATE', 1 contains server host. - // Aggregation data begins from index 2. - logger.Debug("Received aggregate data", h.server, h.count, parts) - h.aggregate.Aggregate(parts[2:]) - logger.Debug("Aggregated aggregate data", h.server, h.count) + parts := strings.SplitN(message, protocol.FieldDelimiter, 3) + if len(parts) != 3 { + logger.Error("Unable to aggregate data", h.server, message, parts, len(parts), "expected 3 parts") + return + } + if err := h.aggregate.Aggregate(parts[2]); err != nil { + logger.Error("Unable to aggregate data", h.server, message, err) + } } diff --git a/internal/clients/maprclient.go b/internal/clients/maprclient.go index 77b674b..cab9a6c 100644 --- a/internal/clients/maprclient.go +++ b/internal/clients/maprclient.go @@ -9,6 +9,8 @@ import ( "time" "github.com/mimecast/dtail/internal/clients/handlers" + "github.com/mimecast/dtail/internal/color" + "github.com/mimecast/dtail/internal/config" "github.com/mimecast/dtail/internal/io/logger" "github.com/mimecast/dtail/internal/mapr" "github.com/mimecast/dtail/internal/omode" @@ -37,6 +39,8 @@ type MaprClient struct { query *mapr.Query // Additative result or new result every interval run? cumulative bool + // The last result string received + lastResult string } // NewMaprClient returns a new mapreduce client. @@ -154,24 +158,37 @@ func (c *MaprClient) reportResults() { func (c *MaprClient) printResults() { var result string var err error - var numLines int + var numRows int if c.cumulative { - result, numLines, err = c.globalGroup.Result(c.query) + result, numRows, err = c.globalGroup.Result(c.query) } else { - result, numLines, err = c.globalGroup.SwapOut().Result(c.query) + result, numRows, err = c.globalGroup.SwapOut().Result(c.query) } + if err != nil { logger.FatalExit(err) } - if numLines == 0 { - logger.Warn("Empty result set this time...") + if result == c.lastResult { + logger.Debug("Result hasn't changed compared to last time...") return } + c.lastResult = result - //logger.Raw(fmt.Sprintf("%s\n", c.query.RawQuery)) - logger.Raw(c.query.RawQuery) + if numRows == 0 { + logger.Debug("Empty result set this time...") + return + } + + rawQuery := c.query.RawQuery + if config.Client.TermColorsEnable { + rawQuery = color.PaintStrWithAttr(rawQuery, + config.Client.TermColors.MaprTable.RawQueryFg, + config.Client.TermColors.MaprTable.RawQueryBg, + config.Client.TermColors.MaprTable.RawQueryAttr) + } + logger.Raw(rawQuery) logger.Raw(result) } -- cgit v1.2.3 From 842fd5800000bb68d6306a9ecad80a98ed762a2f Mon Sep 17 00:00:00 2001 From: Paul Buetow Date: Sun, 12 Sep 2021 19:01:34 +0300 Subject: limit mapreduce table output to 10 rows by default --- internal/clients/maprclient.go | 17 ++++++++++++++--- 1 file changed, 14 insertions(+), 3 deletions(-) (limited to 'internal/clients') diff --git a/internal/clients/maprclient.go b/internal/clients/maprclient.go index cab9a6c..2cad15d 100644 --- a/internal/clients/maprclient.go +++ b/internal/clients/maprclient.go @@ -158,12 +158,18 @@ func (c *MaprClient) reportResults() { func (c *MaprClient) printResults() { var result string var err error - var numRows int + var numRows, rowsLimit int + + if c.query.Limit == -1 { + // Limit output to 10 rows when the result is printed to stdout. + // This can be overriden with the limit clause though. + rowsLimit = 10 + } if c.cumulative { - result, numRows, err = c.globalGroup.Result(c.query) + result, numRows, err = c.globalGroup.Result(c.query, rowsLimit) } else { - result, numRows, err = c.globalGroup.SwapOut().Result(c.query) + result, numRows, err = c.globalGroup.SwapOut().Result(c.query, rowsLimit) } if err != nil { @@ -189,6 +195,11 @@ func (c *MaprClient) printResults() { config.Client.TermColors.MaprTable.RawQueryAttr) } logger.Raw(rawQuery) + + if rowsLimit > 0 && numRows > rowsLimit { + logger.Warn(fmt.Sprintf("Got %d results but limited output to %d rows! Use 'limit' clause to override!", + numRows, rowsLimit)) + } logger.Raw(result) } -- cgit v1.2.3 From 2ebe7e9d63ba62c6f19749c39fe0a577d86ca775 Mon Sep 17 00:00:00 2001 From: Paul Buetow Date: Sun, 12 Sep 2021 19:04:42 +0300 Subject: bugfix: dmap skipped the last couple of mapreduce lines --- internal/clients/maprclient.go | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) (limited to 'internal/clients') diff --git a/internal/clients/maprclient.go b/internal/clients/maprclient.go index 2cad15d..68352ea 100644 --- a/internal/clients/maprclient.go +++ b/internal/clients/maprclient.go @@ -158,7 +158,8 @@ func (c *MaprClient) reportResults() { func (c *MaprClient) printResults() { var result string var err error - var numRows, rowsLimit int + var numRows int + rowsLimit := -1 if c.query.Limit == -1 { // Limit output to 10 rows when the result is printed to stdout. -- cgit v1.2.3 From 5e717af91e8012c72ec7dc0204420dea46f187db Mon Sep 17 00:00:00 2001 From: Paul Buetow Date: Sat, 18 Sep 2021 11:57:09 +0300 Subject: new docker test cases - also change default FATAL bg color to magenta --- internal/clients/args.go | 46 +++++++++++++++++++++++++++++++- internal/clients/baseclient.go | 3 +++ internal/clients/handlers/basehandler.go | 7 ----- 3 files changed, 48 insertions(+), 8 deletions(-) (limited to 'internal/clients') diff --git a/internal/clients/args.go b/internal/clients/args.go index 7f782f1..7ce1634 100644 --- a/internal/clients/args.go +++ b/internal/clients/args.go @@ -1,6 +1,13 @@ package clients import ( + "fmt" + "os" + "path/filepath" + "strings" + + "github.com/mimecast/dtail/internal/config" + "github.com/mimecast/dtail/internal/io/logger" "github.com/mimecast/dtail/internal/omode" gossh "golang.org/x/crypto/ssh" @@ -22,5 +29,42 @@ type Args struct { SSHAuthMethods []gossh.AuthMethod SSHHostKeyCallback gossh.HostKeyCallback PrivateKeyPathFile string - Quiet bool + Quiet bool +} + +// When no servers are given, connect to localhost! +func (a *Args) handleEmptyServer() { + if a.Discovery != "" || a.ServersStr != "" { + return + } + + fqdn, err := os.Hostname() + if err != nil { + logger.FatalExit(err) + } + a.ServersStr = fmt.Sprintf("%s:%d", fqdn, config.Common.SSHPort) + // I am trusting my own hostname. + a.TrustAllHosts = true + logger.Debug("Will connect to local server", a.ServersStr) + + cleanPath := func(dirtyPath string) string { + cleanPath, err := filepath.EvalSymlinks(dirtyPath) + if err != nil { + logger.FatalExit("Unable to evaluate symlinks", dirtyPath, err) + } + cleanPath, err = filepath.Abs(cleanPath) + if err != nil { + logger.FatalExit("Unable to make file path absolute", dirtyPath, cleanPath, err) + } + return cleanPath + } + + logger.Debug("Dirty file paths", a.What) + var filePaths []string + for _, dirtyPath := range strings.Split(a.What, ",") { + filePaths = append(filePaths, cleanPath(dirtyPath)) + } + + a.What = strings.Join(filePaths, ",") + logger.Debug("Clean file paths", a.What) } diff --git a/internal/clients/baseclient.go b/internal/clients/baseclient.go index de0c101..455174e 100644 --- a/internal/clients/baseclient.go +++ b/internal/clients/baseclient.go @@ -41,6 +41,9 @@ type baseClient struct { func (c *baseClient) init() { logger.Debug("Initiating base client") + // When empty server list provided, connect to localhost by default. + c.Args.handleEmptyServer() + flag := regex.Default if c.Args.RegexInvert { flag = regex.Invert diff --git a/internal/clients/handlers/basehandler.go b/internal/clients/handlers/basehandler.go index 51f33c1..74559e9 100644 --- a/internal/clients/handlers/basehandler.go +++ b/internal/clients/handlers/basehandler.go @@ -112,12 +112,5 @@ func (h *baseHandler) handleHiddenMessage(message string) { switch { case strings.HasPrefix(message, ".syn close connection"): h.SendMessage(".ack close connection") - select { - case <-time.After(time.Second * 5): - logger.Debug("Shutting down client after timeout and sending ack to server") - h.Shutdown() - case <-h.Done(): - return - } } } -- cgit v1.2.3 From 6506e20f6c80f4acb7434eb9dd14f784a67189cd Mon Sep 17 00:00:00 2001 From: Paul Buetow Date: Sat, 18 Sep 2021 14:41:25 +0300 Subject: add spartan mode --- internal/clients/args.go | 34 ++++++++++++++++++++++++++++---- internal/clients/baseclient.go | 3 --- internal/clients/catclient.go | 6 ++++-- internal/clients/grepclient.go | 7 +++++-- internal/clients/handlers/basehandler.go | 15 +++++++------- internal/clients/maprclient.go | 7 +++++-- internal/clients/tailclient.go | 7 +++++-- 7 files changed, 56 insertions(+), 23 deletions(-) (limited to 'internal/clients') diff --git a/internal/clients/args.go b/internal/clients/args.go index 7ce1634..081b911 100644 --- a/internal/clients/args.go +++ b/internal/clients/args.go @@ -1,6 +1,7 @@ package clients import ( + "flag" "fmt" "os" "path/filepath" @@ -30,14 +31,35 @@ type Args struct { SSHHostKeyCallback gossh.HostKeyCallback PrivateKeyPathFile string Quiet bool + Spartan bool + NoColor bool } -// When no servers are given, connect to localhost! -func (a *Args) handleEmptyServer() { - if a.Discovery != "" || a.ServersStr != "" { - return +// Transform the arguments based on certain conditions. +func (a *Args) Transform(args []string) { + // Interpret additional args as file list. + if a.What == "" { + var files []string + for _, file := range flag.Args() { + files = append(files, file) + } + a.What = strings.Join(files, ",") + } + + if a.Spartan { + a.Quiet = true + a.NoColor = true } +} +// TransformAfterConfigFile same as Transform, but after the config file has been read. +func (a *Args) TransformAfterConfigFile() { + if a.Discovery == "" && a.ServersStr == "" { + a.handleEmptyServer() + } +} + +func (a *Args) handleEmptyServer() { fqdn, err := os.Hostname() if err != nil { logger.FatalExit(err) @@ -68,3 +90,7 @@ func (a *Args) handleEmptyServer() { a.What = strings.Join(filePaths, ",") logger.Debug("Clean file paths", a.What) } + +func (a *Args) SerializeOptions() string { + return fmt.Sprintf("quiet=%v:spartan=%v", a.Quiet, a.Spartan) +} diff --git a/internal/clients/baseclient.go b/internal/clients/baseclient.go index 455174e..de0c101 100644 --- a/internal/clients/baseclient.go +++ b/internal/clients/baseclient.go @@ -41,9 +41,6 @@ type baseClient struct { func (c *baseClient) init() { logger.Debug("Initiating base client") - // When empty server list provided, connect to localhost by default. - c.Args.handleEmptyServer() - flag := regex.Default if c.Args.RegexInvert { flag = regex.Invert diff --git a/internal/clients/catclient.go b/internal/clients/catclient.go index b7b6131..d14cdcc 100644 --- a/internal/clients/catclient.go +++ b/internal/clients/catclient.go @@ -42,9 +42,11 @@ func (c CatClient) makeHandler(server string) handlers.Handler { } func (c CatClient) makeCommands() (commands []string) { - options := fmt.Sprintf("quiet=%v", c.Args.Quiet) for _, file := range strings.Split(c.What, ",") { - commands = append(commands, fmt.Sprintf("%s:%s %s %s", c.Mode.String(), options, file, c.Regex.Serialize())) + commands = append(commands, fmt.Sprintf("%s:%s %s %s", + c.Mode.String(), + c.Args.SerializeOptions(), + file, c.Regex.Serialize())) } return } diff --git a/internal/clients/grepclient.go b/internal/clients/grepclient.go index 652c31b..ea5022b 100644 --- a/internal/clients/grepclient.go +++ b/internal/clients/grepclient.go @@ -41,9 +41,12 @@ func (c GrepClient) makeHandler(server string) handlers.Handler { } func (c GrepClient) makeCommands() (commands []string) { - options := fmt.Sprintf("quiet=%v", c.Args.Quiet) for _, file := range strings.Split(c.What, ",") { - commands = append(commands, fmt.Sprintf("%s:%s %s %s", c.Mode.String(), options, file, c.Regex.Serialize())) + commands = append(commands, fmt.Sprintf("%s:%s %s %s", + c.Mode.String(), + c.Args.SerializeOptions(), + file, + c.Regex.Serialize())) } return diff --git a/internal/clients/handlers/basehandler.go b/internal/clients/handlers/basehandler.go index 74559e9..0f2d1b5 100644 --- a/internal/clients/handlers/basehandler.go +++ b/internal/clients/handlers/basehandler.go @@ -67,9 +67,12 @@ func (h *baseHandler) Write(p []byte) (n int, err error) { */ case '\n', protocol.MessageDelimiter: message := h.receiveBuf.String() - if len(message) == 0 { - continue - } + /* + // dcat/grep should actually display empty lines. + if len(message) == 0 { + continue + } + */ h.handleMessageType(message) h.receiveBuf.Reset() default: @@ -93,12 +96,8 @@ func (h *baseHandler) Read(p []byte) (n int, err error) { // Handle various message types. func (h *baseHandler) handleMessageType(message string) { - if len(message) == 0 { - return - } - // Hidden server commands starti with a dot "." - if message[0] == '.' { + if len(message) > 0 && message[0] == '.' { h.handleHiddenMessage(message) return } diff --git a/internal/clients/maprclient.go b/internal/clients/maprclient.go index 68352ea..e6ab96b 100644 --- a/internal/clients/maprclient.go +++ b/internal/clients/maprclient.go @@ -116,7 +116,6 @@ func (c MaprClient) makeHandler(server string) handlers.Handler { func (c MaprClient) makeCommands() (commands []string) { commands = append(commands, fmt.Sprintf("map %s", c.query.RawQuery)) - options := fmt.Sprintf("quiet=%v", c.Args.Quiet) modeStr := "cat" if c.Mode == omode.TailClient { @@ -128,7 +127,11 @@ 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 %s %s", modeStr, options, file, c.Regex.Serialize())) + commands = append(commands, fmt.Sprintf("%s:%s %s %s", + modeStr, + c.Args.SerializeOptions(), + file, + c.Regex.Serialize())) } return diff --git a/internal/clients/tailclient.go b/internal/clients/tailclient.go index cefbaa7..360354b 100644 --- a/internal/clients/tailclient.go +++ b/internal/clients/tailclient.go @@ -38,9 +38,12 @@ func (c TailClient) makeHandler(server string) handlers.Handler { } func (c TailClient) makeCommands() (commands []string) { - options := fmt.Sprintf("quiet=%v", c.Args.Quiet) for _, file := range strings.Split(c.What, ",") { - commands = append(commands, fmt.Sprintf("%s:%s %s %s", c.Mode.String(), options, file, c.Regex.Serialize())) + commands = append(commands, fmt.Sprintf("%s:%s %s %s", + c.Mode.String(), + c.Args.SerializeOptions(), + file, + c.Regex.Serialize())) } logger.Debug(commands) -- cgit v1.2.3 From 69b88a1cae0a61bd22530c384f40166b37b9f1ea Mon Sep 17 00:00:00 2001 From: Paul Buetow Date: Sat, 18 Sep 2021 18:43:19 +0300 Subject: remote connector is now an interface --- internal/clients/baseclient.go | 24 +-- internal/clients/connectors/connector.go | 17 ++ internal/clients/connectors/serverconnection.go | 225 ++++++++++++++++++++++++ internal/clients/healthclient.go | 12 +- internal/clients/remote/connection.go | 212 ---------------------- 5 files changed, 259 insertions(+), 231 deletions(-) create mode 100644 internal/clients/connectors/connector.go create mode 100644 internal/clients/connectors/serverconnection.go delete mode 100644 internal/clients/remote/connection.go (limited to 'internal/clients') diff --git a/internal/clients/baseclient.go b/internal/clients/baseclient.go index de0c101..d0631fc 100644 --- a/internal/clients/baseclient.go +++ b/internal/clients/baseclient.go @@ -5,7 +5,7 @@ import ( "sync" "time" - "github.com/mimecast/dtail/internal/clients/remote" + "github.com/mimecast/dtail/internal/clients/connectors" "github.com/mimecast/dtail/internal/discovery" "github.com/mimecast/dtail/internal/io/logger" "github.com/mimecast/dtail/internal/omode" @@ -23,7 +23,7 @@ type baseClient struct { // List of remote servers to connect to. servers []string // We have one connection per remote server. - connections []*remote.Connection + connections []connectors.Connector // SSH auth methods to use to connect to the remote servers. sshAuthMethods []gossh.AuthMethod // To deal with SSH host keys @@ -77,7 +77,7 @@ func (c *baseClient) Start(ctx context.Context, statsCh <-chan string) (status i var mutex sync.Mutex for i, conn := range c.connections { - go func(i int, conn *remote.Connection) { + go func(i int, conn connectors.Connector) { connStatus := c.start(ctx, active, i, conn) // Update global status. @@ -93,7 +93,7 @@ func (c *baseClient) Start(ctx context.Context, statsCh <-chan string) (status i return } -func (c *baseClient) start(ctx context.Context, active chan struct{}, i int, conn *remote.Connection) (status int) { +func (c *baseClient) start(ctx context.Context, active chan struct{}, i int, conn connectors.Connector) (status int) { // Increment connection count active <- struct{}{} // Derement connection count @@ -105,26 +105,20 @@ func (c *baseClient) start(ctx context.Context, active chan struct{}, i int, con conn.Start(connCtx, cancel, c.throttleCh, c.stats.connectionsEstCh) // Retrieve status code from handler (dtail client will exit with that status) - status = conn.Handler.Status() + status = conn.Handler().Status() if !c.retry { return } time.Sleep(time.Second * 2) - logger.Debug(conn.Server, "Reconnecting") - - conn = c.makeConnection(conn.Server, c.sshAuthMethods, c.hostKeyCallback) - c.connections[i] = conn + logger.Debug(conn.Server(), "Reconnecting") + c.connections[i] = c.makeConnection(conn.Server(), c.sshAuthMethods, c.hostKeyCallback) } } -func (c *baseClient) makeConnection(server string, sshAuthMethods []gossh.AuthMethod, hostKeyCallback client.HostKeyCallback) *remote.Connection { - conn := remote.NewConnection(server, c.UserName, sshAuthMethods, hostKeyCallback) - conn.Handler = c.maker.makeHandler(server) - conn.Commands = c.maker.makeCommands() - - return conn +func (c *baseClient) makeConnection(server string, sshAuthMethods []gossh.AuthMethod, hostKeyCallback client.HostKeyCallback) connectors.Connector { + return connectors.NewServerConnection(server, c.UserName, sshAuthMethods, hostKeyCallback, c.maker.makeHandler(server), c.maker.makeCommands()) } func (c *baseClient) waitUntilDone(ctx context.Context, active chan struct{}) { diff --git a/internal/clients/connectors/connector.go b/internal/clients/connectors/connector.go new file mode 100644 index 0000000..3ab6a08 --- /dev/null +++ b/internal/clients/connectors/connector.go @@ -0,0 +1,17 @@ +package connectors + +import ( + "context" + + "github.com/mimecast/dtail/internal/clients/handlers" +) + +// Connector interface. +type Connector interface { + // Start the connection. + Start(ctx context.Context, cancel context.CancelFunc, throttleCh, statsCh chan struct{}) + // Server hostname. + Server() string + // Handler for the connection. + Handler() handlers.Handler +} diff --git a/internal/clients/connectors/serverconnection.go b/internal/clients/connectors/serverconnection.go new file mode 100644 index 0000000..fab2f87 --- /dev/null +++ b/internal/clients/connectors/serverconnection.go @@ -0,0 +1,225 @@ +package connectors + +import ( + "context" + "fmt" + "io" + "strconv" + "strings" + "time" + + "github.com/mimecast/dtail/internal/clients/handlers" + "github.com/mimecast/dtail/internal/config" + "github.com/mimecast/dtail/internal/io/logger" + "github.com/mimecast/dtail/internal/ssh/client" + + "golang.org/x/crypto/ssh" +) + +// ServerConnection represents a client connection connection to a single server. +type ServerConnection struct { + // The remote server's hostname connected to. + server string + // The remote server's port connected to. + port int + // The SSH client configuration used. + config *ssh.ClientConfig + // The SSH client handler to use. + handler handlers.Handler + // DTail commands sent from client to server. When client loses + // connection to the server it re-connects automatically and sends the + // same commands again. + commands []string + // Is it a persistent connection or a one-off? + isOneOff bool + // To deal with SSH server host keys + hostKeyCallback client.HostKeyCallback + // To determine if connection throttling has finished or not + throttlingDone bool +} + +// NewServerConnection returns a new connection. +func NewServerConnection(server string, userName string, authMethods []ssh.AuthMethod, hostKeyCallback client.HostKeyCallback, handler handlers.Handler, commands []string) *ServerConnection { + logger.Debug(server, "Creating new connection") + + c := ServerConnection{ + hostKeyCallback: hostKeyCallback, + server: server, + handler: handler, + commands: commands, + config: &ssh.ClientConfig{ + User: userName, + Auth: authMethods, + HostKeyCallback: hostKeyCallback.Wrap(), + Timeout: time.Second * 2, + }, + } + + c.initServerPort() + return &c +} + +// NewOneOffServerConnection creates new one-off connection (only for sending a series of commands and then quit). +func NewOneOffServerConnection(server string, userName string, authMethods []ssh.AuthMethod, handler handlers.Handler, commands []string) *ServerConnection { + c := ServerConnection{ + server: server, + handler: handler, + commands: commands, + config: &ssh.ClientConfig{ + User: userName, + Auth: authMethods, + HostKeyCallback: ssh.InsecureIgnoreHostKey(), + }, + isOneOff: true, + } + + c.initServerPort() + return &c +} + +// Server hostname +func (c *ServerConnection) Server() string { + return c.server +} + +// Handler for the connection +func (c *ServerConnection) Handler() handlers.Handler { + return c.handler +} + +// Attempt to parse the server port address from the provided server FQDN. +func (c *ServerConnection) initServerPort() { + c.port = config.Common.SSHPort + parts := strings.Split(c.server, ":") + + if len(parts) == 2 { + logger.Debug("Parsing port from hostname", parts) + port, err := strconv.Atoi(parts[1]) + if err != nil { + logger.FatalExit("Unable to parse client port", c.server, parts, err) + } + c.server = parts[0] + c.port = port + } +} + +// Start the server connection. Build up SSH session and send some DTail commands. +func (c *ServerConnection) Start(ctx context.Context, cancel context.CancelFunc, throttleCh, statsCh chan struct{}) { + // Throttle how many connections can be established concurrently (based on ch length) + logger.Debug(c.server, "Throttling connection", len(throttleCh), cap(throttleCh)) + + select { + case throttleCh <- struct{}{}: + case <-ctx.Done(): + logger.Debug(c.server, "Not establishing connection as context is done", len(throttleCh), cap(throttleCh)) + return + } + + logger.Debug(c.server, "Throttling says that the connection can be established", len(throttleCh), cap(throttleCh)) + + go func() { + defer func() { + if !c.throttlingDone { + logger.Debug(c.server, "Unthrottling connection (1)", len(throttleCh), cap(throttleCh)) + c.throttlingDone = true + <-throttleCh + } + cancel() + }() + + if err := c.dial(ctx, cancel, throttleCh, statsCh); err != nil { + logger.Warn(c.server, c.port, err) + if c.hostKeyCallback.Untrusted(fmt.Sprintf("%s:%d", c.server, c.port)) { + logger.Debug(c.server, "Not trusting host") + } + } + }() + + <-ctx.Done() +} + +// Dail into a new SSH connection. Close connection in case of an error. +func (c *ServerConnection) dial(ctx context.Context, cancel context.CancelFunc, throttleCh, statsCh chan struct{}) error { + logger.Debug(c.server, "Incrementing connection stats") + statsCh <- struct{}{} + defer func() { + logger.Debug(c.server, "Decrementing connection stats") + <-statsCh + }() + + logger.Debug(c.server, "Dialing into the connection") + address := fmt.Sprintf("%s:%d", c.server, c.port) + + client, err := ssh.Dial("tcp", address, c.config) + if err != nil { + return err + } + defer client.Close() + + return c.session(ctx, cancel, client, throttleCh) +} + +// Create the SSH session. Close the session in case of an error. +func (c *ServerConnection) session(ctx context.Context, cancel context.CancelFunc, client *ssh.Client, throttleCh chan struct{}) error { + logger.Debug(c.server, "session") + + session, err := client.NewSession() + if err != nil { + return err + } + defer session.Close() + + return c.handle(ctx, cancel, session, throttleCh) +} + +func (c *ServerConnection) handle(ctx context.Context, cancel context.CancelFunc, session *ssh.Session, throttleCh chan struct{}) error { + logger.Debug(c.server, "handle") + + stdinPipe, err := session.StdinPipe() + if err != nil { + return err + } + + stdoutPipe, err := session.StdoutPipe() + if err != nil { + return err + } + + if err := session.Shell(); err != nil { + return err + } + + go func() { + io.Copy(stdinPipe, c.handler) + cancel() + }() + + go func() { + io.Copy(c.handler, stdoutPipe) + cancel() + }() + + go func() { + select { + case <-c.handler.Done(): + case <-ctx.Done(): + } + cancel() + }() + + // Send all commands to client. + for _, command := range c.commands { + logger.Debug(command) + c.handler.SendMessage(command) + } + + if !c.throttlingDone { + logger.Debug(c.server, "Unthrottling connection (2)", len(throttleCh), cap(throttleCh)) + c.throttlingDone = true + <-throttleCh + } + + <-ctx.Done() + c.handler.Shutdown() + return nil +} diff --git a/internal/clients/healthclient.go b/internal/clients/healthclient.go index 692464c..47007b6 100644 --- a/internal/clients/healthclient.go +++ b/internal/clients/healthclient.go @@ -7,8 +7,8 @@ import ( "strings" "time" + "github.com/mimecast/dtail/internal/clients/connectors" "github.com/mimecast/dtail/internal/clients/handlers" - "github.com/mimecast/dtail/internal/clients/remote" "github.com/mimecast/dtail/internal/config" "github.com/mimecast/dtail/internal/omode" "github.com/mimecast/dtail/internal/protocol" @@ -47,9 +47,13 @@ func (c *HealthClient) Start(ctx context.Context) (status int) { throttleCh := make(chan struct{}, runtime.NumCPU()) statsCh := make(chan struct{}, 1) - conn := remote.NewOneOffConnection(c.server, c.userName, c.sshAuthMethods) - conn.Handler = handlers.NewHealthHandler(c.server, receive) - conn.Commands = []string{c.mode.String()} + conn := connectors.NewOneOffServerConnection( + c.server, + c.userName, + c.sshAuthMethods, + handlers.NewHealthHandler(c.server, receive), + []string{c.mode.String()}, + ) connCtx, cancel := context.WithCancel(ctx) go conn.Start(connCtx, cancel, throttleCh, statsCh) diff --git a/internal/clients/remote/connection.go b/internal/clients/remote/connection.go deleted file mode 100644 index b29ffed..0000000 --- a/internal/clients/remote/connection.go +++ /dev/null @@ -1,212 +0,0 @@ -package remote - -import ( - "context" - "fmt" - "io" - "strconv" - "strings" - "time" - - "github.com/mimecast/dtail/internal/clients/handlers" - "github.com/mimecast/dtail/internal/config" - "github.com/mimecast/dtail/internal/io/logger" - "github.com/mimecast/dtail/internal/ssh/client" - - "golang.org/x/crypto/ssh" -) - -// Connection represents a client connection connection to a single server. -type Connection struct { - // The remote server's hostname connected to. - Server string - // The remote server's port connected to. - port int - // The SSH client configuration used. - config *ssh.ClientConfig - // The SSH client handler to use. - Handler handlers.Handler - // DTail commands sent from client to server. When client loses - // connection to the server it re-connects automatically and sends the - // same commands again. - Commands []string - // Is it a persistent connection or a one-off? - isOneOff bool - // To deal with SSH server host keys - hostKeyCallback client.HostKeyCallback - // To determine if connection throttling has finished or not - throttlingDone bool -} - -// NewConnection returns a new connection. -func NewConnection(server string, userName string, authMethods []ssh.AuthMethod, hostKeyCallback client.HostKeyCallback) *Connection { - logger.Debug(server, "Creating new connection") - - c := Connection{ - hostKeyCallback: hostKeyCallback, - config: &ssh.ClientConfig{ - User: userName, - Auth: authMethods, - HostKeyCallback: hostKeyCallback.Wrap(), - Timeout: time.Second * 3, - }, - } - - c.initServerPort(server) - - return &c -} - -// NewOneOffConnection creates new one-off connection (only for sending a series of commands and then quit). -func NewOneOffConnection(server string, userName string, authMethods []ssh.AuthMethod) *Connection { - c := Connection{ - config: &ssh.ClientConfig{ - User: userName, - Auth: authMethods, - HostKeyCallback: ssh.InsecureIgnoreHostKey(), - }, - isOneOff: true, - } - - c.initServerPort(server) - - return &c -} - -// Attempt to parse the server port address from the provided server FQDN. -func (c *Connection) initServerPort(server string) { - c.Server = server - c.port = config.Common.SSHPort - parts := strings.Split(server, ":") - - if len(parts) == 2 { - logger.Debug("Parsing port from hostname", parts) - port, err := strconv.Atoi(parts[1]) - if err != nil { - logger.FatalExit("Unable to parse client port", server, parts, err) - } - c.Server = parts[0] - c.port = port - } -} - -// Start the server connection. Build up SSH session and send some DTail commands. -func (c *Connection) Start(ctx context.Context, cancel context.CancelFunc, throttleCh, statsCh chan struct{}) { - // Throttle how many connections can be established concurrently (based on ch length) - logger.Debug(c.Server, "Throttling connection", len(throttleCh), cap(throttleCh)) - - select { - case throttleCh <- struct{}{}: - case <-ctx.Done(): - logger.Debug(c.Server, "Not establishing connection as context is done", len(throttleCh), cap(throttleCh)) - return - } - - logger.Debug(c.Server, "Throttling says that the connection can be established", len(throttleCh), cap(throttleCh)) - - go func() { - defer func() { - if !c.throttlingDone { - logger.Debug(c.Server, "Unthrottling connection (1)", len(throttleCh), cap(throttleCh)) - c.throttlingDone = true - <-throttleCh - } - cancel() - }() - - if err := c.dial(ctx, cancel, throttleCh, statsCh); err != nil { - logger.Warn(c.Server, c.port, err) - if c.hostKeyCallback.Untrusted(fmt.Sprintf("%s:%d", c.Server, c.port)) { - logger.Debug(c.Server, "Not trusting host") - } - } - }() - - <-ctx.Done() -} - -// Dail into a new SSH connection. Close connection in case of an error. -func (c *Connection) dial(ctx context.Context, cancel context.CancelFunc, throttleCh, statsCh chan struct{}) error { - logger.Debug(c.Server, "Incrementing connection stats") - statsCh <- struct{}{} - defer func() { - logger.Debug(c.Server, "Decrementing connection stats") - <-statsCh - }() - - logger.Debug(c.Server, "Dialing into the connection") - address := fmt.Sprintf("%s:%d", c.Server, c.port) - - client, err := ssh.Dial("tcp", address, c.config) - if err != nil { - return err - } - defer client.Close() - - return c.session(ctx, cancel, client, throttleCh) -} - -// Create the SSH session. Close the session in case of an error. -func (c *Connection) session(ctx context.Context, cancel context.CancelFunc, client *ssh.Client, throttleCh chan struct{}) error { - logger.Debug(c.Server, "session") - - session, err := client.NewSession() - if err != nil { - return err - } - defer session.Close() - - return c.handle(ctx, cancel, session, throttleCh) -} - -func (c *Connection) handle(ctx context.Context, cancel context.CancelFunc, session *ssh.Session, throttleCh chan struct{}) error { - logger.Debug(c.Server, "handle") - - stdinPipe, err := session.StdinPipe() - if err != nil { - return err - } - - stdoutPipe, err := session.StdoutPipe() - if err != nil { - return err - } - - if err := session.Shell(); err != nil { - return err - } - - go func() { - io.Copy(stdinPipe, c.Handler) - cancel() - }() - - go func() { - io.Copy(c.Handler, stdoutPipe) - cancel() - }() - - go func() { - select { - case <-c.Handler.Done(): - case <-ctx.Done(): - } - cancel() - }() - - // Send all commands to client. - for _, command := range c.Commands { - logger.Debug(command) - c.Handler.SendMessage(command) - } - - if !c.throttlingDone { - logger.Debug(c.Server, "Unthrottling connection (2)", len(throttleCh), cap(throttleCh)) - c.throttlingDone = true - <-throttleCh - } - - <-ctx.Done() - c.Handler.Shutdown() - return nil -} -- cgit v1.2.3 From abeac87aec44249bf67f1b0eca471a31086265ca Mon Sep 17 00:00:00 2001 From: Paul Buetow Date: Sat, 18 Sep 2021 19:27:50 +0300 Subject: fix auto reconnect --- internal/clients/args.go | 43 ++---------- internal/clients/baseclient.go | 9 ++- internal/clients/connectors/serverconnection.go | 38 ++++------- internal/clients/connectors/serverless.go | 90 +++++++++++++++++++++++++ internal/clients/handlers/basehandler.go | 10 +++ 5 files changed, 125 insertions(+), 65 deletions(-) create mode 100644 internal/clients/connectors/serverless.go (limited to 'internal/clients') diff --git a/internal/clients/args.go b/internal/clients/args.go index 081b911..dc71e83 100644 --- a/internal/clients/args.go +++ b/internal/clients/args.go @@ -3,12 +3,8 @@ package clients import ( "flag" "fmt" - "os" - "path/filepath" "strings" - "github.com/mimecast/dtail/internal/config" - "github.com/mimecast/dtail/internal/io/logger" "github.com/mimecast/dtail/internal/omode" gossh "golang.org/x/crypto/ssh" @@ -33,6 +29,7 @@ type Args struct { Quiet bool Spartan bool NoColor bool + Serverless bool } // Transform the arguments based on certain conditions. @@ -50,45 +47,15 @@ func (a *Args) Transform(args []string) { a.Quiet = true a.NoColor = true } -} -// TransformAfterConfigFile same as Transform, but after the config file has been read. -func (a *Args) TransformAfterConfigFile() { if a.Discovery == "" && a.ServersStr == "" { - a.handleEmptyServer() + a.Serverless = true } } -func (a *Args) handleEmptyServer() { - fqdn, err := os.Hostname() - if err != nil { - logger.FatalExit(err) - } - a.ServersStr = fmt.Sprintf("%s:%d", fqdn, config.Common.SSHPort) - // I am trusting my own hostname. - a.TrustAllHosts = true - logger.Debug("Will connect to local server", a.ServersStr) - - cleanPath := func(dirtyPath string) string { - cleanPath, err := filepath.EvalSymlinks(dirtyPath) - if err != nil { - logger.FatalExit("Unable to evaluate symlinks", dirtyPath, err) - } - cleanPath, err = filepath.Abs(cleanPath) - if err != nil { - logger.FatalExit("Unable to make file path absolute", dirtyPath, cleanPath, err) - } - return cleanPath - } - - logger.Debug("Dirty file paths", a.What) - var filePaths []string - for _, dirtyPath := range strings.Split(a.What, ",") { - filePaths = append(filePaths, cleanPath(dirtyPath)) - } - - a.What = strings.Join(filePaths, ",") - logger.Debug("Clean file paths", a.What) +// TransformAfterConfigFile same as Transform, but after the config file has been read. +func (a *Args) TransformAfterConfigFile() { + // TODO: Remove this method. It's not used. } func (a *Args) SerializeOptions() string { diff --git a/internal/clients/baseclient.go b/internal/clients/baseclient.go index d0631fc..5523052 100644 --- a/internal/clients/baseclient.go +++ b/internal/clients/baseclient.go @@ -113,12 +113,17 @@ func (c *baseClient) start(ctx context.Context, active chan struct{}, i int, con time.Sleep(time.Second * 2) logger.Debug(conn.Server(), "Reconnecting") - c.connections[i] = c.makeConnection(conn.Server(), c.sshAuthMethods, c.hostKeyCallback) + conn = c.makeConnection(conn.Server(), c.sshAuthMethods, c.hostKeyCallback) + c.connections[i] = conn } } func (c *baseClient) makeConnection(server string, sshAuthMethods []gossh.AuthMethod, hostKeyCallback client.HostKeyCallback) connectors.Connector { - return connectors.NewServerConnection(server, c.UserName, sshAuthMethods, hostKeyCallback, c.maker.makeHandler(server), c.maker.makeCommands()) + if c.Args.Serverless { + return connectors.NewServerless(c.UserName, c.maker.makeHandler(server), c.maker.makeCommands()) + } + return connectors.NewServerConnection(server, c.UserName, sshAuthMethods, + hostKeyCallback, c.maker.makeHandler(server), c.maker.makeCommands()) } func (c *baseClient) waitUntilDone(ctx context.Context, active chan struct{}) { diff --git a/internal/clients/connectors/serverconnection.go b/internal/clients/connectors/serverconnection.go index fab2f87..0904ba1 100644 --- a/internal/clients/connectors/serverconnection.go +++ b/internal/clients/connectors/serverconnection.go @@ -16,31 +16,21 @@ import ( "golang.org/x/crypto/ssh" ) -// ServerConnection represents a client connection connection to a single server. +// ServerConnection represents a connection to a single remote dtail server via SSH protocol. type ServerConnection struct { - // The remote server's hostname connected to. - server string - // The remote server's port connected to. - port int - // The SSH client configuration used. - config *ssh.ClientConfig - // The SSH client handler to use. - handler handlers.Handler - // DTail commands sent from client to server. When client loses - // connection to the server it re-connects automatically and sends the - // same commands again. - commands []string - // Is it a persistent connection or a one-off? - isOneOff bool - // To deal with SSH server host keys + server string + port int + config *ssh.ClientConfig + handler handlers.Handler + commands []string + isOneOff bool hostKeyCallback client.HostKeyCallback - // To determine if connection throttling has finished or not - throttlingDone bool + throttlingDone bool } -// NewServerConnection returns a new connection. +// NewServerConnection returns a new DTail SSH server connection. func NewServerConnection(server string, userName string, authMethods []ssh.AuthMethod, hostKeyCallback client.HostKeyCallback, handler handlers.Handler, commands []string) *ServerConnection { - logger.Debug(server, "Creating new connection") + logger.Debug(server, "Creating new connection", server, handler, commands) c := ServerConnection{ hostKeyCallback: hostKeyCallback, @@ -77,12 +67,10 @@ func NewOneOffServerConnection(server string, userName string, authMethods []ssh return &c } -// Server hostname func (c *ServerConnection) Server() string { return c.server } -// Handler for the connection func (c *ServerConnection) Handler() handlers.Handler { return c.handler } @@ -103,7 +91,6 @@ func (c *ServerConnection) initServerPort() { } } -// Start the server connection. Build up SSH session and send some DTail commands. func (c *ServerConnection) Start(ctx context.Context, cancel context.CancelFunc, throttleCh, statsCh chan struct{}) { // Throttle how many connections can be established concurrently (based on ch length) logger.Debug(c.server, "Throttling connection", len(throttleCh), cap(throttleCh)) @@ -161,7 +148,7 @@ func (c *ServerConnection) dial(ctx context.Context, cancel context.CancelFunc, // Create the SSH session. Close the session in case of an error. func (c *ServerConnection) session(ctx context.Context, cancel context.CancelFunc, client *ssh.Client, throttleCh chan struct{}) error { - logger.Debug(c.server, "session") + logger.Debug(c.server, "Creating SSH session") session, err := client.NewSession() if err != nil { @@ -173,7 +160,7 @@ func (c *ServerConnection) session(ctx context.Context, cancel context.CancelFun } func (c *ServerConnection) handle(ctx context.Context, cancel context.CancelFunc, session *ssh.Session, throttleCh chan struct{}) error { - logger.Debug(c.server, "handle") + logger.Debug(c.server, "Creating handler for SSH session") stdinPipe, err := session.StdinPipe() if err != nil { @@ -221,5 +208,6 @@ func (c *ServerConnection) handle(ctx context.Context, cancel context.CancelFunc <-ctx.Done() c.handler.Shutdown() + return nil } diff --git a/internal/clients/connectors/serverless.go b/internal/clients/connectors/serverless.go new file mode 100644 index 0000000..0500645 --- /dev/null +++ b/internal/clients/connectors/serverless.go @@ -0,0 +1,90 @@ +package connectors + +import ( + "context" + "io" + + "github.com/mimecast/dtail/internal/clients/handlers" + "github.com/mimecast/dtail/internal/config" + "github.com/mimecast/dtail/internal/io/logger" + serverHandlers "github.com/mimecast/dtail/internal/server/handlers" + user "github.com/mimecast/dtail/internal/user/server" +) + +// Serverless creates a server object directly without TCP. +type Serverless struct { + handler handlers.Handler + commands []string + userName string +} + +// NewServerConnection returns a new connection. +func NewServerless(userName string, handler handlers.Handler, commands []string) *Serverless { + s := Serverless{ + userName: userName, + handler: handler, + commands: commands, + } + + logger.Debug("Creating new serverless connector", handler, commands) + return &s +} + +func (s *Serverless) Server() string { + return "local(serverless)" +} + +func (s *Serverless) Handler() handlers.Handler { + return s.handler +} + +func (s *Serverless) Start(ctx context.Context, cancel context.CancelFunc, throttleCh, statsCh chan struct{}) { + go func() { + defer cancel() + + if err := s.handle(ctx, cancel); err != nil { + logger.Warn(err) + } + }() + + <-ctx.Done() +} + +func (s *Serverless) handle(ctx context.Context, cancel context.CancelFunc) error { + logger.Debug("Creating server handler for a serverless session") + + serverHandler := serverHandlers.NewServerHandler( + user.New(s.userName, s.Server()), + make(chan struct{}, config.Server.MaxConcurrentCats), + make(chan struct{}, config.Server.MaxConcurrentTails), + ) + + go func() { + io.Copy(serverHandler, s.handler) + cancel() + }() + + go func() { + io.Copy(s.handler, serverHandler) + cancel() + }() + + go func() { + select { + case <-s.handler.Done(): + case <-ctx.Done(): + } + cancel() + }() + + // Send all commands to client. + for _, command := range s.commands { + logger.Debug(command) + s.handler.SendMessage(command) + } + + <-ctx.Done() + s.handler.Shutdown() + + return nil +} diff --git a/internal/clients/handlers/basehandler.go b/internal/clients/handlers/basehandler.go index 0f2d1b5..af1ad62 100644 --- a/internal/clients/handlers/basehandler.go +++ b/internal/clients/handlers/basehandler.go @@ -22,6 +22,16 @@ type baseHandler struct { status int } +func (h *baseHandler) String() string { + return fmt.Sprintf("baseHandler(%s,server:%s,shellStarted:%v,status:%d)@%p", + h.done, + h.server, + h.shellStarted, + h.status, + h, + ) +} + func (h *baseHandler) Server() string { return h.server } -- cgit v1.2.3 From fe3e68afd99d8ea246be52893730f987e138ec24 Mon Sep 17 00:00:00 2001 From: Paul Buetow Date: Sun, 19 Sep 2021 13:22:59 +0300 Subject: move args to config package logger package rewrite as dlog --- internal/clients/args.go | 63 ------------------------- internal/clients/baseclient.go | 28 +++++++---- internal/clients/catclient.go | 12 +++-- internal/clients/connectors/serverconnection.go | 34 ++++++------- internal/clients/connectors/serverless.go | 21 +++++---- internal/clients/grepclient.go | 13 +++-- internal/clients/handlers/basehandler.go | 6 +-- internal/clients/handlers/clienthandler.go | 4 +- internal/clients/handlers/maprhandler.go | 8 ++-- internal/clients/maprclient.go | 40 ++++++++-------- internal/clients/stats.go | 8 ++-- internal/clients/tailclient.go | 16 ++++--- 12 files changed, 107 insertions(+), 146 deletions(-) delete mode 100644 internal/clients/args.go (limited to 'internal/clients')