Fix WS error when switching from settings to traffic viewer (#985)

* Added log message
* Remove Reset function from OasGenerator interface, use Stop+Start instead
* SetEntriesQuery returns a bool stating whether the query changed

agent/main.go

@@ -199,7 +199,7 @@ func runInHarReaderMode() {
 func enableExpFeatureIfNeeded() {
 	if config.Config.OAS {
 		oasGenerator := dependency.GetInstance(dependency.OasGeneratorDependency).(oas.OasGenerator)
-		oasGenerator.Start()
+		oasGenerator.Start(nil)
 	}
 	if config.Config.ServiceMap {
 		serviceMapGenerator := dependency.GetInstance(dependency.ServiceMapGeneratorDependency).(servicemap.ServiceMap)
@@ -371,7 +371,7 @@ func handleIncomingMessageAsTapper(socketConnection *websocket.Conn) {
 
 func initializeDependencies() {
 	dependency.RegisterGenerator(dependency.ServiceMapGeneratorDependency, func() interface{} { return servicemap.GetDefaultServiceMapInstance() })
-	dependency.RegisterGenerator(dependency.OasGeneratorDependency, func() interface{} { return oas.GetDefaultOasGeneratorInstance(nil) })
+	dependency.RegisterGenerator(dependency.OasGeneratorDependency, func() interface{} { return oas.GetDefaultOasGeneratorInstance() })
 	dependency.RegisterGenerator(dependency.EntriesProvider, func() interface{} { return &entries.BasenineEntriesProvider{} })
 	dependency.RegisterGenerator(dependency.EntriesSocketStreamer, func() interface{} { return &api.BasenineEntryStreamer{} })
 	dependency.RegisterGenerator(dependency.EntryStreamerSocketConnector, func() interface{} { return &api.DefaultEntryStreamerSocketConnector{} })

agent/pkg/controllers/oas_controller_test.go

@@ -58,12 +58,12 @@ func getRecorderAndContext() (*httptest.ResponseRecorder, *gin.Context) {
 		receiveBuffer: bytes.NewBufferString("\n"),
 	}
 	dependency.RegisterGenerator(dependency.OasGeneratorDependency, func() interface{} {
-		return oas.GetDefaultOasGeneratorInstance(dummyConn)
+		return oas.GetDefaultOasGeneratorInstance()
 	})
 
 	recorder := httptest.NewRecorder()
 	c, _ := gin.CreateTestContext(recorder)
-	oas.GetDefaultOasGeneratorInstance(dummyConn).Start()
-	oas.GetDefaultOasGeneratorInstance(dummyConn).GetServiceSpecs().Store("some", oas.NewGen("some"))
+	oas.GetDefaultOasGeneratorInstance().Start(dummyConn)
+	oas.GetDefaultOasGeneratorInstance().GetServiceSpecs().Store("some", oas.NewGen("some"))
 	return recorder, c
 }

agent/pkg/oas/oas_generator.go

@@ -19,46 +19,57 @@ var (
 )
 
 type OasGenerator interface {
-	Start()
+	Start(conn *basenine.Connection)
 	Stop()
 	IsStarted() bool
-	Reset()
 	GetServiceSpecs() *sync.Map
+	SetEntriesQuery(query string) bool
 }
 
 type defaultOasGenerator struct {
-	started      bool
-	ctx          context.Context
-	cancel       context.CancelFunc
-	serviceSpecs *sync.Map
-	dbConn       *basenine.Connection
+	started       bool
+	ctx           context.Context
+	cancel        context.CancelFunc
+	serviceSpecs  *sync.Map
+	dbConn        *basenine.Connection
+	entriesQuery string
 }
 
-func GetDefaultOasGeneratorInstance(conn *basenine.Connection) *defaultOasGenerator {
+func GetDefaultOasGeneratorInstance() *defaultOasGenerator {
 	syncOnce.Do(func() {
-		if conn == nil {
-			c, err := basenine.NewConnection(shared.BasenineHost, shared.BaseninePort)
-			if err != nil {
-				panic(err)
-			}
-			conn = c
-		}
-
-		instance = NewDefaultOasGenerator(conn)
+		instance = NewDefaultOasGenerator()
 		logger.Log.Debug("OAS Generator Initialized")
 	})
 	return instance
 }
 
-func (g *defaultOasGenerator) Start() {
+func (g *defaultOasGenerator) Start(conn *basenine.Connection) {
 	if g.started {
 		return
 	}
+
+	if g.dbConn == nil {
+		if conn == nil {
+			logger.Log.Infof("Creating new DB connection for OAS generator to address %s:%s", shared.BasenineHost, shared.BaseninePort)
+			newConn, err := basenine.NewConnection(shared.BasenineHost, shared.BaseninePort)
+			if err != nil {
+				logger.Log.Error("Error connecting to DB for OAS generator, err: %v", err)
+				return
+			}
+
+			conn = newConn
+		}
+
+		g.dbConn = conn
+	}
+
 	ctx, cancel := context.WithCancel(context.Background())
 	g.cancel = cancel
 	g.ctx = ctx
 	g.serviceSpecs = &sync.Map{}
+
 	g.started = true
+
 	go g.runGenerator()
 }
 
@@ -66,8 +77,15 @@ func (g *defaultOasGenerator) Stop() {
 	if !g.started {
 		return
 	}
+
+	if g.dbConn != nil {
+		g.dbConn.Close()
+		g.dbConn = nil
+	}
+
 	g.cancel()
-	g.Reset()
+	g.reset()
+
 	g.started = false
 }
 
@@ -76,16 +94,19 @@ func (g *defaultOasGenerator) IsStarted() bool {
 }
 
 func (g *defaultOasGenerator) runGenerator() {
-	// Make []byte channels to recieve the data and the meta
+	// Make []byte channels to receive the data and the meta
 	dataChan := make(chan []byte)
 	metaChan := make(chan []byte)
 
-	g.dbConn.Query("", dataChan, metaChan)
+	logger.Log.Infof("Querying DB for OAS generator with query '%s'", g.entriesQuery)
+	g.dbConn.Query(g.entriesQuery, dataChan, metaChan)
 
 	for {
 		select {
 		case <-g.ctx.Done():
 			logger.Log.Infof("OAS Generator was canceled")
+			close(dataChan)
+			close(metaChan)
 			return
 
 		case metaBytes, ok := <-metaChan:
@@ -173,7 +194,7 @@ func (g *defaultOasGenerator) getGen(dest string, urlStr string) *SpecGen {
 	return gen
 }
 
-func (g *defaultOasGenerator) Reset() {
+func (g *defaultOasGenerator) reset() {
 	g.serviceSpecs = &sync.Map{}
 }
 
@@ -181,12 +202,18 @@ func (g *defaultOasGenerator) GetServiceSpecs() *sync.Map {
 	return g.serviceSpecs
 }
 
-func NewDefaultOasGenerator(c *basenine.Connection) *defaultOasGenerator {
+func (g *defaultOasGenerator) SetEntriesQuery(query string) bool {
+	changed := g.entriesQuery != query
+	g.entriesQuery = query
+	return changed
+}
+
+func NewDefaultOasGenerator() *defaultOasGenerator {
 	return &defaultOasGenerator{
 		started:      false,
 		ctx:          nil,
 		cancel:       nil,
 		serviceSpecs: nil,
-		dbConn:       c,
+		dbConn:       nil,
 	}
 }

agent/pkg/oas/oas_generator_test.go

@@ -3,13 +3,11 @@ package oas
 import (
 	"encoding/json"
 	"github.com/up9inc/mizu/agent/pkg/har"
-	"sync"
 	"testing"
 )
 
 func TestOASGen(t *testing.T) {
 	gen := new(defaultOasGenerator)
-	gen.serviceSpecs = &sync.Map{}
 
 	e := new(har.Entry)
 	err := json.Unmarshal([]byte(`{"startedDateTime": "20000101","request": {"url": "https://host/path", "method": "GET"}, "response": {"status": 200}}`), e)
@@ -21,6 +19,9 @@ func TestOASGen(t *testing.T) {
 		Destination: "some",
 		Entry:       *e,
 	}
+
+	dummyConn := GetFakeDBConn(`{"startedDateTime": "20000101","request": {"url": "https://host/path", "method": "GET"}, "response": {"status": 200}}`)
+	gen.Start(dummyConn)
 	gen.handleHARWithSource(ews)
 	g, ok := gen.serviceSpecs.Load("some")
 	if !ok {
@@ -33,4 +34,9 @@ func TestOASGen(t *testing.T) {
 	}
 	specText, _ := json.Marshal(spec)
 	t.Log(string(specText))
+
+	if !gen.IsStarted() {
+		t.Errorf("Should be started")
+	}
+	gen.Stop()
 }

agent/pkg/oas/specgen_test.go

@@ -1,8 +1,10 @@
 package oas
 
 import (
+	"bytes"
 	"encoding/json"
 	"io/ioutil"
+	"net"
 	"os"
 	"regexp"
 	"strings"
@@ -11,13 +13,22 @@ import (
 	"time"
 
 	"github.com/chanced/openapi"
-	"github.com/op/go-logging"
 	"github.com/up9inc/mizu/shared/logger"
 	"github.com/wI2L/jsondiff"
 
+	basenine "github.com/up9inc/basenine/client/go"
 	"github.com/up9inc/mizu/agent/pkg/har"
 )
 
+func GetFakeDBConn(send string) *basenine.Connection {
+	dummyConn := new(basenine.Connection)
+	dummyConn.Conn = FakeConn{
+		sendBuffer:    bytes.NewBufferString(send),
+		receiveBuffer: bytes.NewBufferString(""),
+	}
+	return dummyConn
+}
+
 // if started via env, write file into subdir
 func outputSpec(label string, spec *openapi.OpenAPI, t *testing.T) string {
 	content, err := json.MarshalIndent(spec, "", "  ")
@@ -43,14 +54,14 @@ func outputSpec(label string, spec *openapi.OpenAPI, t *testing.T) string {
 }
 
 func TestEntries(t *testing.T) {
-	logger.InitLoggerStd(logging.INFO)
+	//logger.InitLoggerStd(logging.INFO) causes race condition
 	files, err := getFiles("./test_artifacts/")
 	if err != nil {
 		t.Log(err)
 		t.FailNow()
 	}
 
-	gen := NewDefaultOasGenerator(nil)
+	gen := NewDefaultOasGenerator()
 	gen.serviceSpecs = new(sync.Map)
 	loadStartingOAS("test_artifacts/catalogue.json", "catalogue", gen.serviceSpecs)
 	loadStartingOAS("test_artifacts/trcc.json", "trcc-api-service", gen.serviceSpecs)
@@ -124,7 +135,7 @@ func TestEntries(t *testing.T) {
 }
 
 func TestFileSingle(t *testing.T) {
-	gen := NewDefaultOasGenerator(nil)
+	gen := NewDefaultOasGenerator()
 	gen.serviceSpecs = new(sync.Map)
 	// loadStartingOAS()
 	file := "test_artifacts/params.har"
@@ -214,7 +225,7 @@ func loadStartingOAS(file string, label string, specs *sync.Map) {
 }
 
 func TestEntriesNegative(t *testing.T) {
-	gen := NewDefaultOasGenerator(nil)
+	gen := NewDefaultOasGenerator()
 	gen.serviceSpecs = new(sync.Map)
 	files := []string{"invalid"}
 	_, err := feedEntries(files, false, gen)
@@ -225,7 +236,7 @@ func TestEntriesNegative(t *testing.T) {
 }
 
 func TestEntriesPositive(t *testing.T) {
-	gen := NewDefaultOasGenerator(nil)
+	gen := NewDefaultOasGenerator()
 	gen.serviceSpecs = new(sync.Map)
 	files := []string{"test_artifacts/params.har"}
 	_, err := feedEntries(files, false, gen)
@@ -267,3 +278,17 @@ func TestLoadValid3_1(t *testing.T) {
 		t.FailNow()
 	}
 }
+
+type FakeConn struct {
+	sendBuffer    *bytes.Buffer
+	receiveBuffer *bytes.Buffer
+}
+
+func (f FakeConn) Read(p []byte) (int, error)       { return f.sendBuffer.Read(p) }
+func (f FakeConn) Write(p []byte) (int, error)      { return f.receiveBuffer.Write(p) }
+func (FakeConn) Close() error                       { return nil }
+func (FakeConn) LocalAddr() net.Addr                { return nil }
+func (FakeConn) RemoteAddr() net.Addr               { return nil }
+func (FakeConn) SetDeadline(t time.Time) error      { return nil }
+func (FakeConn) SetReadDeadline(t time.Time) error  { return nil }
+func (FakeConn) SetWriteDeadline(t time.Time) error { return nil }

agent/pkg/servicemap/servicemap.go

@@ -224,7 +224,8 @@ func (s *defaultServiceMap) GetStatus() ServiceMapStatus {
 }
 
 func (s *defaultServiceMap) GetNodes() []ServiceMapNode {
-	var nodes []ServiceMapNode
+	nodes := []ServiceMapNode{}
+
 	for i, n := range s.graph.Nodes {
 		nodes = append(nodes, ServiceMapNode{
 			Id:       n.id,
@@ -234,11 +235,13 @@ func (s *defaultServiceMap) GetNodes() []ServiceMapNode {
 			Count:    n.count,
 		})
 	}
+
 	return nodes
 }
 
 func (s *defaultServiceMap) GetEdges() []ServiceMapEdge {
-	var edges []ServiceMapEdge
+	edges := []ServiceMapEdge{}
+
 	for u, m := range s.graph.Edges {
 		for v := range m {
 			for _, p := range s.graph.Edges[u][v].data {
@@ -263,6 +266,7 @@ func (s *defaultServiceMap) GetEdges() []ServiceMapEdge {
 			}
 		}
 	}
+
 	return edges
 }
 

agent/pkg/servicemap/servicemap_test.go

@@ -403,10 +403,10 @@ func (s *ServiceMapEnabledSuite) TestServiceMap() {
 	assert.Equal(0, status.EdgeCount)
 
 	// Nodes after reset
-	assert.Equal([]ServiceMapNode(nil), nodes)
+	assert.Equal([]ServiceMapNode{}, nodes)
 
 	// Edges after reset
-	assert.Equal([]ServiceMapEdge(nil), edges)
+	assert.Equal([]ServiceMapEdge{}, edges)
 }
 
 func TestServiceMapSuite(t *testing.T) {

tap/api/api.go

@@ -161,7 +161,7 @@ type Entry struct {
 	Capture                Capture                `json:"capture"`
 	Source                 *TCP                   `json:"src"`
 	Destination            *TCP                   `json:"dst"`
-	Namespace              string                 `json:"namespace,omitempty"`
+	Namespace              string                 `json:"namespace"`
 	Outgoing               bool                   `json:"outgoing"`
 	Timestamp              int64                  `json:"timestamp"`
 	StartTime              time.Time              `json:"startTime"`

tap/extensions/amqp/Makefile

@@ -13,4 +13,4 @@ test-pull-bin:
 
 test-pull-expect:
 	@mkdir -p expect
-	@[ "${skipexpect}" ] && echo "Skipping downloading expected JSONs" || gsutil  -o 'GSUtil:parallel_process_count=5' -o 'GSUtil:parallel_thread_count=5'  -m cp -r gs://static.up9.io/mizu/test-pcap/expect5/amqp/\* expect
+	@[ "${skipexpect}" ] && echo "Skipping downloading expected JSONs" || gsutil  -o 'GSUtil:parallel_process_count=5' -o 'GSUtil:parallel_thread_count=5'  -m cp -r gs://static.up9.io/mizu/test-pcap/expect6/amqp/\* expect

tap/extensions/http/Makefile

@@ -13,4 +13,4 @@ test-pull-bin:
 
 test-pull-expect:
 	@mkdir -p expect
-	@[ "${skipexpect}" ] && echo "Skipping downloading expected JSONs" || gsutil  -o 'GSUtil:parallel_process_count=5' -o 'GSUtil:parallel_thread_count=5'  -m cp -r gs://static.up9.io/mizu/test-pcap/expect5/http/\* expect
+	@[ "${skipexpect}" ] && echo "Skipping downloading expected JSONs" || gsutil  -o 'GSUtil:parallel_process_count=5' -o 'GSUtil:parallel_thread_count=5'  -m cp -r gs://static.up9.io/mizu/test-pcap/expect6/http/\* expect

tap/extensions/http/http2_assembler.go

@@ -28,26 +28,6 @@ const protoMinorHTTP2 = 0
 
 var maxHTTP2DataLen = 1 * 1024 * 1024 // 1MB
 
-var grpcStatusCodes = []string{
-	"OK",
-	"CANCELLED",
-	"UNKNOWN",
-	"INVALID_ARGUMENT",
-	"DEADLINE_EXCEEDED",
-	"NOT_FOUND",
-	"ALREADY_EXISTS",
-	"PERMISSION_DENIED",
-	"RESOURCE_EXHAUSTED",
-	"FAILED_PRECONDITION",
-	"ABORTED",
-	"OUT_OF_RANGE",
-	"UNIMPLEMENTED",
-	"INTERNAL",
-	"UNAVAILABLE",
-	"DATA_LOSS",
-	"UNAUTHENTICATED",
-}
-
 type messageFragment struct {
 	headers []hpack.HeaderField
 	data    []byte
@@ -142,18 +122,8 @@ func (ga *Http2Assembler) readMessage() (streamID uint32, messageHTTP1 interface
 
 	// gRPC detection
 	grpcStatus := headersHTTP1.Get("Grpc-Status")
-	if grpcStatus != "" {
+	if grpcStatus != "" || strings.Contains(headersHTTP1.Get("Content-Type"), "application/grpc") {
 		isGrpc = true
-		status = grpcStatus
-	}
-
-	if strings.Contains(headersHTTP1.Get("Content-Type"), "application/grpc") {
-		isGrpc = true
-		grpcPath := headersHTTP1.Get(":path")
-		pathSegments := strings.Split(grpcPath, "/")
-		if len(pathSegments) > 0 {
-			method = pathSegments[len(pathSegments)-1]
-		}
 	}
 
 	if method != "" {

tap/extensions/http/main.go

@@ -248,11 +248,6 @@ func (d dissecting) Analyze(item *api.OutputChannelItem, resolvedSource string,
 	reqDetails["_queryStringMerged"] = mapSliceMergeRepeatedKeys(reqDetails["_queryString"].([]interface{}))
 	reqDetails["queryString"] = mapSliceRebuildAsMap(reqDetails["_queryStringMerged"].([]interface{}))
 
-	statusCode := int(resDetails["status"].(float64))
-	if item.Protocol.Abbreviation == "gRPC" {
-		resDetails["statusText"] = grpcStatusCodes[statusCode]
-	}
-
 	elapsedTime := item.Pair.Response.CaptureTime.Sub(item.Pair.Request.CaptureTime).Round(time.Millisecond).Milliseconds()
 	if elapsedTime < 0 {
 		elapsedTime = 0

tap/extensions/kafka/Makefile

@@ -13,4 +13,4 @@ test-pull-bin:
 
 test-pull-expect:
 	@mkdir -p expect
-	@[ "${skipexpect}" ] && echo "Skipping downloading expected JSONs" || gsutil  -o 'GSUtil:parallel_process_count=5' -o 'GSUtil:parallel_thread_count=5'  -m cp -r gs://static.up9.io/mizu/test-pcap/expect5/kafka/\* expect
+	@[ "${skipexpect}" ] && echo "Skipping downloading expected JSONs" || gsutil  -o 'GSUtil:parallel_process_count=5' -o 'GSUtil:parallel_thread_count=5'  -m cp -r gs://static.up9.io/mizu/test-pcap/expect6/kafka/\* expect

tap/extensions/redis/Makefile

@@ -13,4 +13,4 @@ test-pull-bin:
 
 test-pull-expect:
 	@mkdir -p expect
-	@[ "${skipexpect}" ] && echo "Skipping downloading expected JSONs" || gsutil  -o 'GSUtil:parallel_process_count=5' -o 'GSUtil:parallel_thread_count=5'  -m cp -r gs://static.up9.io/mizu/test-pcap/expect5/redis/\* expect
+	@[ "${skipexpect}" ] && echo "Skipping downloading expected JSONs" || gsutil  -o 'GSUtil:parallel_process_count=5' -o 'GSUtil:parallel_thread_count=5'  -m cp -r gs://static.up9.io/mizu/test-pcap/expect6/redis/\* expect

ui-common/src/components/TrafficViewer/EntryDetailed.tsx

@@ -89,12 +89,13 @@ const EntryTitle: React.FC<any> = ({ protocol, data, elapsedTime }) => {
     </div>;
 };
 
-const EntrySummary: React.FC<any> = ({ entry }) => {
+const EntrySummary: React.FC<any> = ({ entry, namespace }) => {
     return <EntryItem
         key={`entry-${entry.id}`}
         entry={entry}
         style={{}}
         headingMode={true}
+        namespace={namespace}
     />;
 };
 
@@ -140,7 +141,7 @@ export const EntryDetailed = () => {
             data={entryData.data}
             elapsedTime={entryData.data.elapsedTime}
         />}
-        {!isLoading && entryData && <EntrySummary entry={entryData.base} />}
+        {!isLoading && entryData && <EntrySummary entry={entryData.base} namespace={entryData.data.namespace} />}
         <React.Fragment>
             {!isLoading && entryData && <EntryViewer
                 representation={entryData.representation}

ui-common/src/components/TrafficViewer/EntryListItem/EntryListItem.tsx

@@ -52,6 +52,7 @@ interface EntryProps {
     entry: Entry;
     style: object;
     headingMode: boolean;
+    namespace?: string;
 }
 
 enum CaptureTypes {
@@ -62,7 +63,7 @@ enum CaptureTypes {
     Ebpf = "ebpf",
 }
 
-export const EntryItem: React.FC<EntryProps> = ({entry, style, headingMode}) => {
+export const EntryItem: React.FC<EntryProps> = ({entry, style, headingMode, namespace}) => {
 
     const [focusedEntryId, setFocusedEntryId] = useRecoilState(focusedEntryIdAtom);
     const [queryState, setQuery] = useRecoilState(queryAtom);
@@ -224,6 +225,19 @@ export const EntryItem: React.FC<EntryProps> = ({entry, style, headingMode}) =>
                 : ""
             }
             <div className={styles.separatorRight}>
+                {headingMode ? <Queryable
+                        query={`namespace == "${namespace}"`}
+                        displayIconOnMouseOver={true}
+                        flipped={true}
+                        iconStyle={{marginRight: "16px"}}
+                >
+                    <span
+                        className={`${styles.tcpInfo} ${styles.ip}`}
+                        title="Namespace"
+                    >
+                        {namespace}
+                    </span>
+                </Queryable> : null}
                 <Queryable
                         query={`src.ip == "${entry.src.ip}"`}
                         displayIconOnMouseOver={true}

ui-common/src/components/TrafficViewer/TrafficViewer.tsx

@@ -201,7 +201,9 @@ export const TrafficViewer: React.FC<TrafficViewerProps> = ({
 
   useEffect(() => {
     return () => {
-      ws.current.close();
+      if (ws?.current?.readyState === WebSocket.OPEN) {
+        ws.current.close();
+      }
     };
   }, []);
 

ui-common/src/components/UI/StatusCode.tsx

@@ -37,9 +37,9 @@ export function getClassification(statusCode: number): string {
 
     // 1 - 16 HTTP/2 (gRPC) status codes
     // 2xx - 5xx HTTP/1.x status codes
-    if ((statusCode >= 200 && statusCode <= 399) || statusCode === 0) {
+    if (statusCode >= 200 && statusCode <= 399) {
         classification = StatusCodeClassification.SUCCESS;
-    } else if (statusCode >= 400 || (statusCode >= 1 && statusCode <= 16)) {
+    } else if (statusCode >= 400) {
         classification = StatusCodeClassification.FAILURE;
     }