Remove redundant Google auth from test workflow (#911)

* Remove google auth for test workflow

.github/workflows/test.yml

@@ -34,14 +34,6 @@ jobs:
         run: |
           sudo apt-get install libpcap-dev
 
-      - id: 'auth'
-        uses: 'google-github-actions/auth@v0'
-        with:
-          credentials_json: '${{ secrets.GCR_JSON_KEY }}'
-
-      - name: 'Set up Cloud SDK'
-        uses: 'google-github-actions/setup-gcloud@v0'
-
       - name: Check CLI modified files
         id: cli_modified_files
         run: devops/check_modified_files.sh cli/

tap/tlstapper/bpf-builder/build.sh

@@ -1,5 +1,7 @@
 #!/bin/bash
 
+pushd "$(dirname "$0")" || exit 1
+
 MIZU_HOME=$(realpath ../../../)
 
 docker build -t mizu-ebpf-builder . || exit 1
@@ -7,6 +9,7 @@ docker build -t mizu-ebpf-builder . || exit 1
 docker run --rm \
 	--name mizu-ebpf-builder \
 	-v $MIZU_HOME:/mizu \
+	-v $(go env GOPATH):/root/go \
 	-it mizu-ebpf-builder \
 	sh -c "
 		go generate tap/tlstapper/tls_tapper.go
@@ -15,3 +18,5 @@ docker run --rm \
 		chown $(id -u):$(id -g) tap/tlstapper/tlstapper_bpfel.go
 		chown $(id -u):$(id -g) tap/tlstapper/tlstapper_bpfel.o
 	" || exit 1
+
+popd

tap/tlstapper/bpf/fd_to_address_tracepoints.c

@@ -69,7 +69,7 @@ void sys_exit_accept4(struct sys_exit_accept4_ctx *ctx) {
 	
 	struct accept_info *infoPtr = bpf_map_lookup_elem(&accept_syscall_context, &id);
 	
-	if (infoPtr == 0) {
+	if (infoPtr == NULL) {
 		return;
 	}
 	
@@ -175,7 +175,7 @@ void sys_exit_connect(struct sys_exit_connect_ctx *ctx) {
 	
 	struct connect_info *infoPtr = bpf_map_lookup_elem(&connect_syscall_info, &id);
 	
-	if (infoPtr == 0) {
+	if (infoPtr == NULL) {
 		return;
 	}
 	

tap/tlstapper/bpf/fd_tracepoints.c

@@ -28,7 +28,7 @@ void sys_enter_read(struct sys_enter_read_ctx *ctx) {
 	
 	struct ssl_info *infoPtr = bpf_map_lookup_elem(&ssl_read_context, &id);
 	
-	if (infoPtr == 0) {
+	if (infoPtr == NULL) {
 		return;
 	}
 	
@@ -71,7 +71,7 @@ void sys_enter_write(struct sys_enter_write_ctx *ctx) {
 	
 	struct ssl_info *infoPtr = bpf_map_lookup_elem(&ssl_write_context, &id);
 	
-	if (infoPtr == 0) {
+	if (infoPtr == NULL) {
 		return;
 	}
 	

tap/tlstapper/bpf/include/maps.h

@@ -10,6 +10,12 @@ Copyright (C) UP9 Inc.
 #define FLAGS_IS_CLIENT_BIT (1 << 0)
 #define FLAGS_IS_READ_BIT (1 << 1)
 
+#define CHUNK_SIZE (1 << 12)
+#define MAX_CHUNKS_PER_OPERATION (8)
+
+// One minute in nano seconds. Chosen by gut feeling.
+#define SSL_INFO_MAX_TTL_NANO (1000000000l * 60l)
+
 // The same struct can be found in chunk.go
 //  
 //  Be careful when editing, alignment and padding should be exactly the same in go/c.
@@ -18,16 +24,18 @@ struct tlsChunk {
     __u32 pid;
     __u32 tgid;
     __u32 len;
+    __u32 start;
     __u32 recorded;
     __u32 fd;
     __u32 flags;
     __u8 address[16];
-    __u8 data[4096]; // Must be N^2
+    __u8 data[CHUNK_SIZE]; // Must be N^2
 };
 
 struct ssl_info {
     void* buffer;
     __u32 fd;
+    __u64 created_at_nano;
     
     // for ssl_write and ssl_read must be zero
     // for ssl_write_ex and ssl_read_ex save the *written/*readbytes pointer. 
@@ -53,10 +61,13 @@ struct fd_info {
 
 #define BPF_PERF_OUTPUT(_name) \
     BPF_MAP(_name, BPF_MAP_TYPE_PERF_EVENT_ARRAY, int, __u32, 1024)
+    
+#define BPF_LRU_HASH(_name, _key_type, _value_type) \
+    BPF_MAP(_name, BPF_MAP_TYPE_LRU_HASH, _key_type, _value_type, 16384)
 
 BPF_HASH(pids_map, __u32, __u32);
-BPF_HASH(ssl_write_context, __u64, struct ssl_info);
-BPF_HASH(ssl_read_context, __u64, struct ssl_info);
+BPF_LRU_HASH(ssl_write_context, __u64, struct ssl_info);
+BPF_LRU_HASH(ssl_read_context, __u64, struct ssl_info);
 BPF_HASH(file_descriptor_to_ipv4, __u64, struct fd_info);
 BPF_PERF_OUTPUT(chunks_buffer);
 

tap/tlstapper/bpf/openssl_uprobes.c

@@ -18,16 +18,166 @@ struct {
 	__type(value, struct tlsChunk);
 } heap SEC(".maps");
 
-static __always_inline int ssl_uprobe(void* ssl, void* buffer, int num, struct bpf_map_def* map_fd, size_t *count_ptr) {
-	__u64 id = bpf_get_current_pid_tgid();
+static __always_inline int get_count_bytes(struct pt_regs *ctx, struct ssl_info* info, __u64 id) {
+	int returnValue = PT_REGS_RC(ctx);
 	
-	if (!should_tap(id >> 32)) {
+	if (info->count_ptr == NULL) {
+		// ssl_read and ssl_write return the number of bytes written/read
+		//
+		return returnValue;
+	}
+	
+	// ssl_read_ex and ssl_write_ex return 1 for success
+	//
+	if (returnValue != 1) {
 		return 0;
 	}
 	
+	// ssl_read_ex and ssl_write_ex write the number of bytes to an arg named *count
+	//
+	size_t countBytes;
+	long err = bpf_probe_read(&countBytes, sizeof(size_t), (void*) info->count_ptr);
+	
+	if (err != 0) {
+		char msg[] = "Error reading bytes count of _ex (id: %ld) (err: %ld)";
+		bpf_trace_printk(msg, sizeof(msg), id, err);
+		return 0;
+	}
+	
+	return countBytes;
+}
+
+static __always_inline void add_address_to_chunk(struct tlsChunk* chunk, __u64 id, __u32 fd) {
+	__u32 pid = id >> 32;
+	__u64 key = (__u64) pid << 32 | fd;
+	
+	struct fd_info *fdinfo = bpf_map_lookup_elem(&file_descriptor_to_ipv4, &key);
+	
+	if (fdinfo == NULL) {
+		return;
+	}
+	
+	int err = bpf_probe_read(chunk->address, sizeof(chunk->address), fdinfo->ipv4_addr);
+	chunk->flags |= (fdinfo->flags & FLAGS_IS_CLIENT_BIT);
+	
+	if (err != 0) {
+		char msg[] = "Error reading from fd address %ld - %ld";
+		bpf_trace_printk(msg, sizeof(msg), id, err);
+	}
+}
+
+static __always_inline void send_chunk_part(struct pt_regs *ctx, __u8* buffer, __u64 id, 
+	struct tlsChunk* chunk, int start, int end) {
+	size_t recorded = MIN(end - start, sizeof(chunk->data));
+	
+	if (recorded <= 0) {
+		return;
+	}
+	
+	chunk->recorded = recorded;
+	chunk->start = start;
+	
+	// This ugly trick is for the ebpf verifier happiness
+	//
+	long err = 0;
+	if (chunk->recorded == sizeof(chunk->data)) {
+		err = bpf_probe_read(chunk->data, sizeof(chunk->data), buffer + start);
+	} else {
+		recorded &= (sizeof(chunk->data) - 1); // Buffer must be N^2
+		err = bpf_probe_read(chunk->data, recorded, buffer + start);
+	}
+	
+	if (err != 0) {
+		char msg[] = "Error reading from ssl buffer %ld - %ld";
+		bpf_trace_printk(msg, sizeof(msg), id, err);
+		return;
+	}
+	
+	bpf_perf_event_output(ctx, &chunks_buffer, BPF_F_CURRENT_CPU, chunk, sizeof(struct tlsChunk));
+}
+
+static __always_inline void send_chunk(struct pt_regs *ctx, __u8* buffer, __u64 id, struct tlsChunk* chunk) {
+	// ebpf loops must be bounded at compile time, we can't use (i < chunk->len / CHUNK_SIZE)
+	//
+	// 	https://lwn.net/Articles/794934/
+	// 
+	// If we want to compile in kernel older than 5.3, we should add "#pragma unroll" to this loop
+	// 
+	for (int i = 0; i < MAX_CHUNKS_PER_OPERATION; i++) {
+		if (chunk->len <= (CHUNK_SIZE * i)) {
+			break;
+		}
+		
+		send_chunk_part(ctx, buffer, id, chunk, CHUNK_SIZE * i, chunk->len);
+	}
+}
+
+static __always_inline void output_ssl_chunk(struct pt_regs *ctx, struct ssl_info* info, __u64 id, __u32 flags) {
+	int countBytes = get_count_bytes(ctx, info, id);
+	
+	if (countBytes <= 0) {
+		return;
+	}
+	
+	if (countBytes > (CHUNK_SIZE * MAX_CHUNKS_PER_OPERATION)) {
+		char msg[] = "Buffer too big %d (id: %ld)";
+		bpf_trace_printk(msg, sizeof(msg), countBytes, id);
+		return;
+	}
+	
+	struct tlsChunk* chunk;
+	int zero = 0;
+	
+	// If other thread, running on the same CPU get to this point at the same time like us (context switch)
+	//	the data will be corrupted - protection may be added in the future
+	//
+	chunk = bpf_map_lookup_elem(&heap, &zero);
+	
+	if (!chunk) {
+		char msg[] = "Unable to allocate chunk (id: %ld)";
+		bpf_trace_printk(msg, sizeof(msg), id);
+		return;
+	}
+	
+	chunk->flags = flags;
+	chunk->pid = id >> 32;
+	chunk->tgid = id;
+	chunk->len = countBytes;
+	chunk->fd = info->fd;
+	
+	add_address_to_chunk(chunk, id, chunk->fd);
+	send_chunk(ctx, info->buffer, id, chunk);
+}
+
+static __always_inline void ssl_uprobe(void* ssl, void* buffer, int num, struct bpf_map_def* map_fd, size_t *count_ptr) {
+	__u64 id = bpf_get_current_pid_tgid();
+	
+	if (!should_tap(id >> 32)) {
+		return;
+	}
+	
+	struct ssl_info *infoPtr = bpf_map_lookup_elem(map_fd, &id);
 	struct ssl_info info = {};
 	
-	info.fd = -1;
+	if (infoPtr == NULL) {
+		info.fd = -1;
+		info.created_at_nano = bpf_ktime_get_ns();
+	} else {
+		long err = bpf_probe_read(&info, sizeof(struct ssl_info), infoPtr);
+		
+		if (err != 0) {
+			char msg[] = "Error reading old ssl context (id: %ld) (err: %ld)";
+			bpf_trace_printk(msg, sizeof(msg), id, err);
+		}
+		
+		if ((bpf_ktime_get_ns() - info.created_at_nano) > SSL_INFO_MAX_TTL_NANO) {
+			// If the ssl info is too old, we don't want to use its info because it may be incorrect.
+			//
+			info.fd = -1;
+			info.created_at_nano = bpf_ktime_get_ns();
+		}
+	}
+	
 	info.count_ptr = count_ptr;
 	info.buffer = buffer;
 	
@@ -36,163 +186,90 @@ static __always_inline int ssl_uprobe(void* ssl, void* buffer, int num, struct b
 	if (err != 0) {
 		char msg[] = "Error putting ssl context (id: %ld) (err: %ld)";
 		bpf_trace_printk(msg, sizeof(msg), id, err);
-		return 0;
 	}
-	
-	return 0;
 }
 
-static __always_inline int ssl_uretprobe(struct pt_regs *ctx, struct bpf_map_def* map_fd, __u32 flags) {
+static __always_inline void ssl_uretprobe(struct pt_regs *ctx, struct bpf_map_def* map_fd, __u32 flags) {
 	__u64 id = bpf_get_current_pid_tgid();
 	
 	if (!should_tap(id >> 32)) {
-		return 0;
+		return;
 	}
 	
 	struct ssl_info *infoPtr = bpf_map_lookup_elem(map_fd, &id);
 	
-	if (infoPtr == 0) {
+	if (infoPtr == NULL) {
 		char msg[] = "Error getting ssl context info (id: %ld)";
 		bpf_trace_printk(msg, sizeof(msg), id);
-		return 0;
+		return;
 	}
 	
 	struct ssl_info info;
 	long err = bpf_probe_read(&info, sizeof(struct ssl_info), infoPtr);
 	
-	bpf_map_delete_elem(map_fd, &id);
+	// Do not clean map on purpose, sometimes there are two calls to ssl_read in a raw
+	//	while the first call actually goes to read from socket, and we get the chance
+	//	to find the fd. The other call already have all the information and we don't
+	//	have the chance to get the fd.
+	//
+	// There are two risks keeping the map items
+	//	1. It gets full - we solve it by using BPF_MAP_TYPE_LRU_HASH with hard limit
+	//	2. We get wrong info of an old call - we solve it by comparing the timestamp 
+	//		info before using it
+	//
+	// bpf_map_delete_elem(map_fd, &id);
 	
 	if (err != 0) {
 		char msg[] = "Error reading ssl context (id: %ld) (err: %ld)";
 		bpf_trace_printk(msg, sizeof(msg), id, err);
-		return 0;
+		return;
 	}
 	
 	if (info.fd == -1) {
 		char msg[] = "File descriptor is missing from ssl info (id: %ld)";
 		bpf_trace_printk(msg, sizeof(msg), id);
-		return 0;
+		return;
 	}
 	
-	int countBytes = PT_REGS_RC(ctx);
-	
-	if (info.count_ptr != 0) {
-		// ssl_read_ex and ssl_write_ex return 1 for success
-		//
-		if (countBytes != 1) {
-			return 0; 
-		}
-		
-		size_t tempCount;
-		long err = bpf_probe_read(&tempCount, sizeof(size_t), (void*) info.count_ptr);
-		
-		if (err != 0) {
-			char msg[] = "Error reading bytes count of _ex (id: %ld) (err: %ld)";
-			bpf_trace_printk(msg, sizeof(msg), id, err);
-			return 0;
-		}
-		
-		countBytes = tempCount;
-	}
-	
-	if (countBytes <= 0) {
-		return 0;
-	}
-	
-	struct tlsChunk* c;
-	int zero = 0;
-	
-	// If other thread, running on the same CPU get to this point at the same time like us
-	//	the data will be corrupted - protection may be added in the future
-	//	
-	c = bpf_map_lookup_elem(&heap, &zero);
-	
-	if (!c) {
-		char msg[] = "Unable to allocate chunk (id: %ld)";
-		bpf_trace_printk(msg, sizeof(msg), id);
-		return 0;
-	}
-	
-	size_t recorded = MIN(countBytes, sizeof(c->data));
-	
-	c->flags = flags;
-	c->pid = id >> 32;
-	c->tgid = id;
-	c->len = countBytes;
-	c->recorded = recorded;
-	c->fd = info.fd;
-	
-	// This ugly trick is for the ebpf verifier happiness
-	//
-	if (recorded == sizeof(c->data)) {
-		err = bpf_probe_read(c->data, sizeof(c->data), info.buffer);
-	} else {
-		recorded &= sizeof(c->data) - 1; // Buffer must be N^2
-		err = bpf_probe_read(c->data, recorded, info.buffer);
-	}
-	
-	if (err != 0) {
-		char msg[] = "Error reading from ssl buffer %ld - %ld";
-		bpf_trace_printk(msg, sizeof(msg), id, err);
-		return 0;
-	}
-	
-	__u32 pid = id >> 32;
-	__u32 fd = info.fd;
-	__u64 key = (__u64) pid << 32 | fd;
-	
-	struct fd_info *fdinfo = bpf_map_lookup_elem(&file_descriptor_to_ipv4, &key);
-	
-	if (fdinfo != 0) {
-		err = bpf_probe_read(c->address, sizeof(c->address), fdinfo->ipv4_addr);
-		c->flags |= (fdinfo->flags & FLAGS_IS_CLIENT_BIT);
-		
-		if (err != 0) {
-			char msg[] = "Error reading from fd address %ld - %ld";
-			bpf_trace_printk(msg, sizeof(msg), id, err);
-		}
-	}
-	
-	bpf_perf_event_output(ctx, &chunks_buffer, BPF_F_CURRENT_CPU, c, sizeof(struct tlsChunk));
-	return 0;
+	output_ssl_chunk(ctx, &info, id, flags);
 }
 
 SEC("uprobe/ssl_write")
-int BPF_KPROBE(ssl_write, void* ssl, void* buffer, int num) {
-	return ssl_uprobe(ssl, buffer, num, &ssl_write_context, 0);
+void BPF_KPROBE(ssl_write, void* ssl, void* buffer, int num) {
+	ssl_uprobe(ssl, buffer, num, &ssl_write_context, 0);
 }
 
 SEC("uretprobe/ssl_write")
-int BPF_KPROBE(ssl_ret_write) {
-	return ssl_uretprobe(ctx, &ssl_write_context, 0);
+void BPF_KPROBE(ssl_ret_write) {
+	ssl_uretprobe(ctx, &ssl_write_context, 0);
 }
 
 SEC("uprobe/ssl_read")
-int BPF_KPROBE(ssl_read, void* ssl, void* buffer, int num) {
-	return ssl_uprobe(ssl, buffer, num, &ssl_read_context, 0);
+void BPF_KPROBE(ssl_read, void* ssl, void* buffer, int num) {
+	ssl_uprobe(ssl, buffer, num, &ssl_read_context, 0);
 }
 
 SEC("uretprobe/ssl_read")
-int BPF_KPROBE(ssl_ret_read) {
-	return ssl_uretprobe(ctx, &ssl_read_context, FLAGS_IS_READ_BIT);
+void BPF_KPROBE(ssl_ret_read) {
+	ssl_uretprobe(ctx, &ssl_read_context, FLAGS_IS_READ_BIT);
 }
 
 SEC("uprobe/ssl_write_ex")
-int BPF_KPROBE(ssl_write_ex, void* ssl, void* buffer, size_t num, size_t *written) {
-	return ssl_uprobe(ssl, buffer, num, &ssl_write_context, written);
+void BPF_KPROBE(ssl_write_ex, void* ssl, void* buffer, size_t num, size_t *written) {
+	ssl_uprobe(ssl, buffer, num, &ssl_write_context, written);
 }
 
 SEC("uretprobe/ssl_write_ex")
-int BPF_KPROBE(ssl_ret_write_ex) {
-	return ssl_uretprobe(ctx, &ssl_write_context, 0);
+void BPF_KPROBE(ssl_ret_write_ex) {
+	ssl_uretprobe(ctx, &ssl_write_context, 0);
 }
 
 SEC("uprobe/ssl_read_ex")
-int BPF_KPROBE(ssl_read_ex, void* ssl, void* buffer, size_t num, size_t *readbytes) {
-	return ssl_uprobe(ssl, buffer, num, &ssl_read_context, readbytes);
+void BPF_KPROBE(ssl_read_ex, void* ssl, void* buffer, size_t num, size_t *readbytes) {
+	ssl_uprobe(ssl, buffer, num, &ssl_read_context, readbytes);
 }
 
 SEC("uretprobe/ssl_read_ex")
-int BPF_KPROBE(ssl_ret_read_ex) {
-	return ssl_uretprobe(ctx, &ssl_read_context, FLAGS_IS_READ_BIT);
+void BPF_KPROBE(ssl_ret_read_ex) {
+	ssl_uretprobe(ctx, &ssl_read_context, FLAGS_IS_READ_BIT);
 }

tap/tlstapper/chunk.go

@@ -16,14 +16,15 @@ const FLAGS_IS_READ_BIT uint32 = (1 << 1)
 //	Be careful when editing, alignment and padding should be exactly the same in go/c.
 //
 type tlsChunk struct {
-	Pid      uint32
-	Tgid     uint32
-	Len      uint32
-	Recorded uint32
-	Fd       uint32
-	Flags    uint32
-	Address  [16]byte
-	Data     [4096]byte
+	Pid      uint32	// process id
+	Tgid     uint32	// thread id inside the process
+	Len      uint32	// the size of the native buffer used to read/write the tls data (may be bigger than tlsChunk.Data[])
+	Start    uint32	// the start offset withing the native buffer
+	Recorded uint32	// number of bytes copied from the native buffer to tlsChunk.Data[]
+	Fd       uint32	// the file descriptor used to read/write the tls data (probably socket file descriptor)
+	Flags    uint32 // bitwise flags
+	Address  [16]byte // ipv4 address and port
+	Data     [4096]byte // actual tls data
 }
 
 func (c *tlsChunk) getAddress() (net.IP, uint16, error) {

tap/tlstapper/tls_poller.go

@@ -225,8 +225,8 @@ func (p *tlsPoller) logTls(chunk *tlsChunk, ip net.IP, port uint16) {
 
 	str := strings.ReplaceAll(strings.ReplaceAll(string(chunk.Data[0:chunk.Recorded]), "\n", " "), "\r", "")
 
-	logger.Log.Infof("PID: %v (tid: %v) (fd: %v) (client: %v) (addr: %v:%v) (fdaddr %v:%v>%v:%v) (recorded %v out of %v) - %v - %v",
+	logger.Log.Infof("PID: %v (tid: %v) (fd: %v) (client: %v) (addr: %v:%v) (fdaddr %v:%v>%v:%v) (recorded %v out of %v starting at %v) - %v - %v",
 		chunk.Pid, chunk.Tgid, chunk.Fd, flagsStr, ip, port,
 		srcIp, srcPort, dstIp, dstPort,
-		chunk.Recorded, chunk.Len, str, hex.EncodeToString(chunk.Data[0:chunk.Recorded]))
+		chunk.Recorded, chunk.Len, chunk.Start, str, hex.EncodeToString(chunk.Data[0:chunk.Recorded]))
 }

ui-common/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@up9/mizu-common",
-  "version": "1.0.130",
+  "version": "1.0.132",
   "description": "Made with create-react-library",
   "author": "",
   "license": "MIT",

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

@@ -5,10 +5,8 @@ import Moment from 'moment';
 import {EntryItem} from "./EntryListItem/EntryListItem";
 import down from "assets/downImg.svg";
 import spinner from 'assets/spinner.svg';
-
 import {RecoilState, useRecoilState, useRecoilValue} from "recoil";
 import entriesAtom from "../../recoil/entries";
-import wsConnectionAtom, {WsConnectionStatus} from "../../recoil/wsConnection";
 import queryAtom from "../../recoil/query";
 import TrafficViewerApiAtom from "../../recoil/TrafficViewerApi";
 import TrafficViewerApi from "./TrafficViewerApi";
@@ -32,14 +30,15 @@ interface EntriesListProps {
     truncatedTimestamp: number;
     setTruncatedTimestamp: any;
     scrollableRef: any;
+    ws: any;
 }
 
-export const EntriesList: React.FC<EntriesListProps> = ({listEntryREF, onSnapBrokenEvent, isSnappedToBottom, setIsSnappedToBottom, queriedCurrent, setQueriedCurrent, queriedTotal, setQueriedTotal, startTime, noMoreDataTop, setNoMoreDataTop, leftOffTop, setLeftOffTop, openWebSocket, leftOffBottom, truncatedTimestamp, setTruncatedTimestamp, scrollableRef}) => {
+export const EntriesList: React.FC<EntriesListProps> = ({listEntryREF, onSnapBrokenEvent, isSnappedToBottom, setIsSnappedToBottom, queriedCurrent, setQueriedCurrent, queriedTotal, setQueriedTotal, startTime, noMoreDataTop, setNoMoreDataTop, leftOffTop, setLeftOffTop, openWebSocket, leftOffBottom, truncatedTimestamp, setTruncatedTimestamp, scrollableRef, ws}) => {
 
     const [entries, setEntries] = useRecoilState(entriesAtom);
-    const wsConnection = useRecoilValue(wsConnectionAtom);
     const query = useRecoilValue(queryAtom);
-    const isWsConnectionClosed = wsConnection === WsConnectionStatus.Closed;
+    const isWsConnectionClosed = ws?.current?.readyState !== WebSocket.OPEN;
+
     const trafficViewerApi = useRecoilValue(TrafficViewerApiAtom as RecoilState<TrafficViewerApi>)
 
     const [loadMoreTop, setLoadMoreTop] = useState(false);

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

@@ -1,4 +1,4 @@
-import React, { useCallback, useEffect, useMemo, useRef, useState } from "react";
+import React, { useEffect, useMemo, useRef, useState } from "react";
 import { Filters } from "./Filters";
 import { EntriesList } from "./EntriesList";
 import { makeStyles } from "@material-ui/core";
@@ -14,7 +14,6 @@ import debounce from 'lodash/debounce';
 import { RecoilRoot, RecoilState, useRecoilState, useRecoilValue, useSetRecoilState } from "recoil";
 import entriesAtom from "../../recoil/entries";
 import focusedEntryIdAtom from "../../recoil/focusedEntryId";
-import websocketConnectionAtom, { WsConnectionStatus } from "../../recoil/wsConnection";
 import queryAtom from "../../recoil/query";
 import { TLSWarning } from "../TLSWarning/TLSWarning";
 import trafficViewerApiAtom from "../../recoil/TrafficViewerApi"
@@ -60,12 +59,12 @@ export const TrafficViewer : React.FC<TrafficViewerProps> = ({setAnalyzeStatus,
 
   const [entries, setEntries] = useRecoilState(entriesAtom);
   const [focusedEntryId, setFocusedEntryId] = useRecoilState(focusedEntryIdAtom);
-  const [wsConnection, setWsConnection] = useRecoilState(websocketConnectionAtom);
   const query = useRecoilValue(queryAtom);
   const setTrafficViewerApiState = useSetRecoilState(trafficViewerApiAtom as RecoilState<TrafficViewerApi>)
   const [tappingStatus, setTappingStatus] = useRecoilState(tappingStatusAtom);
   const [noMoreDataTop, setNoMoreDataTop] = useState(false);
   const [isSnappedToBottom, setIsSnappedToBottom] = useState(true);
+  const [forceRender, setForceRender] = useState(0);
 
   const [queryBackgroundColor, setQueryBackgroundColor] = useState("#f5f5f5");
 
@@ -121,22 +120,26 @@ export const TrafficViewer : React.FC<TrafficViewerProps> = ({setAnalyzeStatus,
       setLeftOffTop(null);
       setNoMoreDataTop(false);
     }
-    ws.current = new WebSocket(webSocketUrl);
-    ws.current.onopen = () => {
-      setWsConnection(WsConnectionStatus.Connected);
+    try {
+      ws.current = new WebSocket(webSocketUrl);
       sendQueryWhenWsOpen(query);
-    }
-    ws.current.onclose = () => {
-      setWsConnection(WsConnectionStatus.Closed);
-    }
-    ws.current.onerror = (event) => {
-      console.error("WebSocket error:", event);
-      if (query) {
-        openWebSocket(`(${query}) and leftOff(${leftOffBottom})`, false);
-      } else {
-        openWebSocket(`leftOff(${leftOffBottom})`, false);
+
+      ws.current.onclose = () => {
+        if(window.location.pathname === "/")
+          setForceRender(forceRender + 1);
       }
-    }
+      ws.current.onerror = (event) => {
+        console.error("WebSocket error:", event);
+        if (ws?.current?.readyState === WebSocket.OPEN) {
+          ws.current.close();
+        }
+        if (query) {
+          openWebSocket(`(${query}) and leftOff(${leftOffBottom})`, false);
+        } else {
+          openWebSocket(`leftOff(${leftOffBottom})`, false);
+        }
+      }
+    } catch (e) {}
   }
 
   const sendQueryWhenWsOpen = (query) => {
@@ -150,7 +153,9 @@ export const TrafficViewer : React.FC<TrafficViewerProps> = ({setAnalyzeStatus,
   }
 
   const closeWebSocket = () => {
-    ws.current && ws.current.close()
+    if(ws?.current?.readyState === WebSocket.OPEN) {
+      ws.current.close();
+    }
   }
 
   if (ws.current) {
@@ -229,8 +234,9 @@ export const TrafficViewer : React.FC<TrafficViewerProps> = ({setAnalyzeStatus,
   }, []);
 
   const toggleConnection = () => {
-    ws.current.close();
-    if (wsConnection !== WsConnectionStatus.Connected) {
+    if(ws?.current?.readyState === WebSocket.OPEN) {
+      ws?.current?.close();
+    } else {
       if (query) {
         openWebSocket(`(${query}) and leftOff(-1)`, true);
       } else {
@@ -241,6 +247,12 @@ export const TrafficViewer : React.FC<TrafficViewerProps> = ({setAnalyzeStatus,
     }
   }
 
+  useEffect(() => {
+    return () => {
+      ws.current.close();
+    };
+  }, []);
+
   const onTLSDetected = (destAddress: string) => {
     addressesWithTLS.add(destAddress);
     setAddressesWithTLS(new Set(addressesWithTLS));
@@ -251,8 +263,8 @@ export const TrafficViewer : React.FC<TrafficViewerProps> = ({setAnalyzeStatus,
   };
 
   const getConnectionIndicator = () => {
-    switch (wsConnection) {
-      case WsConnectionStatus.Connected:
+    switch (ws?.current?.readyState) {
+      case WebSocket.OPEN:
         return <div className={`${TrafficViewerStyles.indicatorContainer} ${TrafficViewerStyles.greenIndicatorContainer}`}>
           <div className={`${TrafficViewerStyles.indicator} ${TrafficViewerStyles.greenIndicator}`} />
         </div>
@@ -264,8 +276,8 @@ export const TrafficViewer : React.FC<TrafficViewerProps> = ({setAnalyzeStatus,
   }
 
   const getConnectionTitle = () => {
-    switch (wsConnection) {
-      case WsConnectionStatus.Connected:
+    switch (ws?.current?.readyState) {
+      case WebSocket.OPEN:
         return "streaming live traffic"
       default:
         return "streaming paused";
@@ -274,7 +286,7 @@ export const TrafficViewer : React.FC<TrafficViewerProps> = ({setAnalyzeStatus,
 
   const onSnapBrokenEvent = () => {
     setIsSnappedToBottom(false);
-    if (wsConnection === WsConnectionStatus.Connected) {
+    if (ws?.current?.readyState === WebSocket.OPEN) {
       ws.current.close();
     }
   }
@@ -284,9 +296,9 @@ export const TrafficViewer : React.FC<TrafficViewerProps> = ({setAnalyzeStatus,
       {tappingStatus && isShowStatusBar && <StatusBar />}
       <div className={TrafficViewerStyles.TrafficPageHeader}>
         <div className={TrafficViewerStyles.TrafficPageStreamStatus}>
-          <img className={TrafficViewerStyles.playPauseIcon} style={{ visibility: wsConnection === WsConnectionStatus.Connected ? "visible" : "hidden" }} alt="pause"
+          <img className={TrafficViewerStyles.playPauseIcon} style={{ visibility: ws?.current?.readyState === WebSocket.OPEN ? "visible" : "hidden" }} alt="pause"
             src={pauseIcon} onClick={toggleConnection} />
-          <img className={TrafficViewerStyles.playPauseIcon} style={{ position: "absolute", visibility: wsConnection === WsConnectionStatus.Connected ? "hidden" : "visible" }} alt="play"
+          <img className={TrafficViewerStyles.playPauseIcon} style={{ position: "absolute", visibility: ws?.current?.readyState === WebSocket.OPEN ? "hidden" : "visible" }} alt="play"
             src={playIcon} onClick={toggleConnection} />
           <div className={TrafficViewerStyles.connectionText}>
             {getConnectionTitle()}
@@ -322,6 +334,7 @@ export const TrafficViewer : React.FC<TrafficViewerProps> = ({setAnalyzeStatus,
               truncatedTimestamp={truncatedTimestamp}
               setTruncatedTimestamp={setTruncatedTimestamp}
               scrollableRef={scrollableRef}
+              ws={ws}
             />
           </div>
         </div>

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

@@ -15,7 +15,7 @@ interface Props {
     classes?: any,
     tabs: Tab[],
     currentTab: string,
-    color: string,
+    color?: string,
     onChange: (string) => void,
     leftAligned?: boolean,
     dark?: boolean,

ui-common/src/recoil/wsConnection/atom.ts

@@ -1,8 +0,0 @@
-import { atom } from "recoil";
-
-const wsConnectionAtom = atom({
-    key: "wsConnectionAtom",
-    default: 0
-});
-
-export default wsConnectionAtom;

ui-common/src/recoil/wsConnection/index.ts

@@ -1,10 +0,0 @@
-import atom from "./atom";
-
-enum WsConnectionStatus {
-    Closed,
-    Connected,
-}
-
-export {WsConnectionStatus};
-
-export default atom

ui/package.json

@@ -13,7 +13,7 @@
     "@types/jest": "^26.0.22",
     "@types/node": "^12.20.10",
     "@uiw/react-textarea-code-editor": "^1.4.12",
-    "@up9/mizu-common": "^1.0.130",
+    "@up9/mizu-common": "1.0.131",
     "axios": "^0.25.0",
     "core-js": "^3.20.2",
     "craco-babel-loader": "^1.0.3",