diff --git a/web/client/src/App.tsx b/web/client/src/App.tsx index c517f2e..f7d1cf2 100644 --- a/web/client/src/App.tsx +++ b/web/client/src/App.tsx @@ -110,8 +110,13 @@ class App extends React.Component { if (msg.type === MessageType.CONN) { this.connMgr.add(msg.id, msg.content as IConnection) this.setState({ flows: this.state.flows }) - } else if (msg.type === MessageType.REQUEST) { + } + else if (msg.type === MessageType.CONN_CLOSE) { + this.connMgr.delete(msg.id) + } + else if (msg.type === MessageType.REQUEST) { const flow = new Flow(msg, this.connMgr) + flow.getConn() this.flowMgr.add(flow) let shouldScroll = false @@ -133,6 +138,7 @@ class App extends React.Component { else if (msg.type === MessageType.RESPONSE) { const flow = this.flowMgr.get(msg.id) if (!flow) return + flow.getConn() flow.addResponse(msg) this.setState({ flows: this.state.flows }) } diff --git a/web/client/src/lib/message.ts b/web/client/src/lib/message.ts index 3836443..00c5e00 100644 --- a/web/client/src/lib/message.ts +++ b/web/client/src/lib/message.ts @@ -3,6 +3,7 @@ import type { Flow, IFlowRequest, IRequest, IResponse } from './flow' export enum MessageType { CONN = 0, + CONN_CLOSE = 5, REQUEST = 1, REQUEST_BODY = 2, RESPONSE = 3, @@ -11,6 +12,7 @@ export enum MessageType { const allMessageBytes = [ MessageType.CONN, + MessageType.CONN_CLOSE, MessageType.REQUEST, MessageType.REQUEST_BODY, MessageType.RESPONSE, diff --git a/web/conn.go b/web/conn.go index cda992e..ec78d47 100644 --- a/web/conn.go +++ b/web/conn.go @@ -53,6 +53,15 @@ func (c *concurrentConn) trySendConnMessage(f *proxy.Flow) { func (c *concurrentConn) whenConnClose(connCtx *proxy.ConnContext) { delete(c.sendConnMessageMap, connCtx.Id().String()) + + msg := newMessageConnClose(connCtx) + c.mu.Lock() + err := c.conn.WriteMessage(websocket.BinaryMessage, msg.bytes()) + c.mu.Unlock() + if err != nil { + log.Error(err) + return + } } func (c *concurrentConn) writeMessage(msg *messageFlow, f *proxy.Flow) { diff --git a/web/message.go b/web/message.go index a81788c..fa9a1dc 100644 --- a/web/message.go +++ b/web/message.go @@ -31,6 +31,7 @@ type messageType byte const ( messageTypeConn messageType = 0 + messageTypeConnClose messageType = 5 messageTypeRequest messageType = 1 messageTypeRequestBody messageType = 2 messageTypeResponse messageType = 3 @@ -46,6 +47,7 @@ const ( var allMessageTypes = []messageType{ messageTypeConn, + messageTypeConnClose, messageTypeRequest, messageTypeRequestBody, messageTypeResponse, @@ -114,6 +116,13 @@ func newMessageFlow(mType messageType, f *proxy.Flow) *messageFlow { } } +func newMessageConnClose(connCtx *proxy.ConnContext) *messageFlow { + return &messageFlow{ + mType: messageTypeConnClose, + id: connCtx.Id(), + } +} + func (m *messageFlow) bytes() []byte { buf := bytes.NewBuffer(make([]byte, 0)) buf.WriteByte(byte(messageVersion))