web addon add connection message

2 years ago · 9bd9ca04d0
parent 0c26763147
commit 9bd9ca04d0
6 changed files with 138 additions and 13 deletions
--- a/proxy/connection.go
+++ b/proxy/connection.go
@ -3,6 +3,7 @@ package proxy
 import (
 	"context"
 	"crypto/tls"
+	"encoding/json"
 	"net"
 	"net/http"

@ -25,6 +26,14 @@ func newClientConn(c net.Conn) *ClientConn {
 	}
 }

+func (c *ClientConn) MarshalJSON() ([]byte, error) {
+	m := make(map[string]interface{})
+	m["id"] = c.Id
+	m["tls"] = c.Tls
+	m["address"] = c.Conn.RemoteAddr().String()
+	return json.Marshal(m)
+}
+
 // server connection
 type ServerConn struct {
 	Id      uuid.UUID
@ -45,6 +54,14 @@ func newServerConn() *ServerConn {
 	}
 }

+func (c *ServerConn) MarshalJSON() ([]byte, error) {
+	m := make(map[string]interface{})
+	m["id"] = c.Id
+	m["address"] = c.Address
+	m["peername"] = c.Conn.RemoteAddr().String()
+	return json.Marshal(m)
+}
+
 func (c *ServerConn) TlsState() *tls.ConnectionState {
 	<-c.tlsHandshaked
 	return c.tlsState
@ -55,8 +72,8 @@ var connContextKey = new(struct{})

 // connection context
 type ConnContext struct {
-	ClientConn *ClientConn
-	ServerConn *ServerConn
+	ClientConn *ClientConn `json:"clientConn"`
+	ServerConn *ServerConn `json:"serverConn"`

 	proxy    *Proxy
 	pipeConn *pipeConn
@ -70,6 +87,10 @@ func newConnContext(c net.Conn, proxy *Proxy) *ConnContext {
 	}
 }

+func (connCtx *ConnContext) Id() uuid.UUID {
+	return connCtx.ClientConn.Id
+}
+
 func (connCtx *ConnContext) initHttpServerConn() {
 	if connCtx.ServerConn != nil {
 		return
--- a/web/client/src/components/ViewFlow.tsx
+++ b/web/client/src/components/ViewFlow.tsx
@ -17,7 +17,7 @@ interface Iprops {
 }

 interface IState {
-  flowTab: 'Headers' | 'Preview' | 'Response' | 'Hexview'
+  flowTab: 'Headers' | 'Preview' | 'Response' | 'Hexview' | 'Detail'
  copied: boolean
  requestBodyViewTab: 'Raw' | 'Preview'
  responseBodyLineBreak: boolean
@ -28,7 +28,7 @@ class ViewFlow extends React.Component<Iprops, IState> {
    super(props)

    this.state = {
-      flowTab: 'Headers',
+      flowTab: 'Detail',
      copied: false,
      requestBodyViewTab: 'Raw',
      responseBodyLineBreak: false,
@ -88,6 +88,13 @@ class ViewFlow extends React.Component<Iprops, IState> {
    return <pre>{flow.hexviewResponseBody()}</pre>
  }

+  detail() {
+    const { flow } = this.props
+    if (!flow) return null
+
+    return <div>detail todo</div>
+  }
+
  render() {
    if (!this.props.flow) return null

@ -109,6 +116,7 @@ class ViewFlow extends React.Component<Iprops, IState> {
      <div className="flow-detail">
        <div className="header-tabs">
          <span onClick={() => { this.props.onClose() }}>x</span>
+          <span className={flowTab === 'Detail' ? 'selected' : undefined} onClick={() => { this.setState({ flowTab: 'Detail' }) }}>Detail</span>
          <span className={flowTab === 'Headers' ? 'selected' : undefined} onClick={() => { this.setState({ flowTab: 'Headers' }) }}>Headers</span>
          <span className={flowTab === 'Preview' ? 'selected' : undefined} onClick={() => { this.setState({ flowTab: 'Preview' }) }}>Preview</span>
          <span className={flowTab === 'Response' ? 'selected' : undefined} onClick={() => { this.setState({ flowTab: 'Response' }) }}>Response</span>
@ -281,6 +289,11 @@ class ViewFlow extends React.Component<Iprops, IState> {
            !(flowTab === 'Hexview') ? null :
              <div>{this.hexview()}</div>
          }
+
+          {
+            !(flowTab === 'Detail') ? null :
+              <div>{this.detail()}</div>
+          }
        </div>

      </div>
--- a/web/client/src/message.ts
+++ b/web/client/src/message.ts
@ -1,6 +1,7 @@
 import { arrayBufferToBase64, bufHexView, getSize, isTextBody } from './utils'

 export enum MessageType {
+  CONN = 0,
  REQUEST = 1,
  REQUEST_BODY = 2,
  RESPONSE = 3,
@ -9,6 +10,20 @@ export enum MessageType {

 export type Header = Record<string, string[]>

+export interface IConnection {
+  id: string
+  clientConn: {
+    id: string
+    tls: boolean
+    address: string
+  }
+  serverConn: {
+    id: string
+    address: string
+    peername: string
+  }
+}
+
 export interface IRequest {
  method: string
  url: string
@ -17,6 +32,11 @@ export interface IRequest {
  body?: ArrayBuffer
 }

+interface IFlowRequest {
+  connId: string
+  request: IRequest
+}
+
 export interface IResponse {
  statusCode: number
  header: Header
@ -27,7 +47,7 @@ export interface IMessage {
  type: MessageType
  id: string
  waitIntercept: boolean
-  content?: ArrayBuffer | IRequest | IResponse
+  content?: ArrayBuffer | IFlowRequest | IResponse
 }

 export interface IFlowPreview {
@ -51,6 +71,7 @@ interface IPreviewBody {
 export class Flow {
  public no: number
  public id: string
+  public connId: string
  public waitIntercept: boolean
  public request: IRequest
  public response: IResponse | null = null
@ -84,7 +105,10 @@ export class Flow {
    this.no = ++Flow.curNo
    this.id = msg.id
    this.waitIntercept = msg.waitIntercept
-    this.request = msg.content as IRequest
+
+    const flowRequestMsg = msg.content as IFlowRequest
+    this.connId = flowRequestMsg.connId
+    this.request = flowRequestMsg.request

    this.url = new URL(this.request.url)
    this.path = this.url.pathname + this.url.search
@ -255,6 +279,7 @@ export class Flow {
 }

 const allMessageBytes = [
+  MessageType.CONN,
  MessageType.REQUEST,
  MessageType.REQUEST_BODY,
  MessageType.RESPONSE,
@ -262,7 +287,7 @@ const allMessageBytes = [
 ]


-// type: 1/2/3/4
+// type: 0/1/2/3/4
 // messageFlow
 // version 1 byte + type 1 byte + id 36 byte + waitIntercept 1 byte + content left bytes
 export const parseMessage = (data: ArrayBuffer): IMessage | null => {
--- a/web/conn.go
+++ b/web/conn.go
@ -19,6 +19,8 @@ type concurrentConn struct {
 	conn *websocket.Conn
 	mu   sync.Mutex

+	sendConnMessageMap map[string]bool
+
 	waitChans   map[string]chan interface{}
 	waitChansMu sync.Mutex

@ -27,11 +29,32 @@ type concurrentConn struct {

 func newConn(c *websocket.Conn) *concurrentConn {
 	return &concurrentConn{
-		conn:      c,
-		waitChans: make(map[string]chan interface{}),
+		conn:               c,
+		sendConnMessageMap: make(map[string]bool),
+		waitChans:          make(map[string]chan interface{}),
+	}
+}
+
+func (c *concurrentConn) trySendConnMessage(f *proxy.Flow) {
+	key := f.ConnContext.Id().String()
+	if send := c.sendConnMessageMap[key]; send {
+		return
+	}
+	c.sendConnMessageMap[key] = true
+	msg := newMessageFlow(messageTypeConn, f)
+	c.mu.Lock()
+	err := c.conn.WriteMessage(websocket.BinaryMessage, msg.bytes())
+	c.mu.Unlock()
+	if err != nil {
+		log.Error(err)
+		return
 	}
 }

+func (c *concurrentConn) whenConnClose(connCtx *proxy.ConnContext) {
+	delete(c.sendConnMessageMap, connCtx.Id().String())
+}
+
 func (c *concurrentConn) writeMessage(msg *messageFlow, f *proxy.Flow) {
 	if c.isIntercpt(f, msg) {
 		msg.waitIntercept = 1
--- a/web/message.go
+++ b/web/message.go
@ -13,7 +13,7 @@ import (

 // message:

-// type: 1/2/3/4
+// type: 0/1/2/3/4
 // messageFlow
 // version 1 byte + type 1 byte + id 36 byte + waitIntercept 1 byte + content left bytes

@ -30,6 +30,7 @@ const messageVersion = 1
 type messageType byte

 const (
+	messageTypeConn         messageType = 0
 	messageTypeRequest      messageType = 1
 	messageTypeRequestBody  messageType = 2
 	messageTypeResponse     messageType = 3
@ -44,6 +45,7 @@ const (
 )

 var allMessageTypes = []messageType{
+	messageTypeConn,
 	messageTypeRequest,
 	messageTypeRequestBody,
 	messageTypeResponse,
@ -79,8 +81,13 @@ func newMessageFlow(mType messageType, f *proxy.Flow) *messageFlow {
 	var content []byte
 	var err error = nil

-	if mType == messageTypeRequest {
-		content, err = json.Marshal(f.Request)
+	if mType == messageTypeConn {
+		content, err = json.Marshal(f.ConnContext)
+	} else if mType == messageTypeRequest {
+		m := make(map[string]interface{})
+		m["request"] = f.Request
+		m["connId"] = f.ConnContext.Id().String()
+		content, err = json.Marshal(m)
 	} else if mType == messageTypeRequestBody {
 		content = f.Request.Body
 	} else if mType == messageTypeResponse {
@ -95,9 +102,14 @@ func newMessageFlow(mType messageType, f *proxy.Flow) *messageFlow {
 		panic(err)
 	}

+	id := f.Id
+	if mType == messageTypeConn {
+		id = f.ConnContext.Id()
+	}
+
 	return &messageFlow{
 		mType:   mType,
-		id:      f.Id,
+		id:      id,
 		content: content,
 	}
 }
--- a/web/web.go
+++ b/web/web.go
@ -92,6 +92,19 @@ func (web *WebAddon) removeConn(conn *concurrentConn) {
 	web.conns = append(web.conns[:index], web.conns[index+1:]...)
 }

+func (web *WebAddon) forEachConn(do func(c *concurrentConn)) bool {
+	web.connsMu.RLock()
+	conns := web.conns
+	web.connsMu.RUnlock()
+	if len(conns) == 0 {
+		return false
+	}
+	for _, c := range conns {
+		do(c)
+	}
+	return true
+}
+
 func (web *WebAddon) sendFlow(f *proxy.Flow, msgFn func() *messageFlow) bool {
 	web.connsMu.RLock()
 	conns := web.conns
@ -110,6 +123,12 @@ func (web *WebAddon) sendFlow(f *proxy.Flow, msgFn func() *messageFlow) bool {
 }

 func (web *WebAddon) Requestheaders(f *proxy.Flow) {
+	if f.ConnContext.ClientConn.Tls {
+		web.forEachConn(func(c *concurrentConn) {
+			c.trySendConnMessage(f)
+		})
+	}
+
 	web.sendFlow(f, func() *messageFlow {
 		return newMessageFlow(messageTypeRequest, f)
 	})
@ -122,6 +141,12 @@ func (web *WebAddon) Request(f *proxy.Flow) {
 }

 func (web *WebAddon) Responseheaders(f *proxy.Flow) {
+	if !f.ConnContext.ClientConn.Tls {
+		web.forEachConn(func(c *concurrentConn) {
+			c.trySendConnMessage(f)
+		})
+	}
+
 	web.sendFlow(f, func() *messageFlow {
 		return newMessageFlow(messageTypeResponse, f)
 	})
@ -132,3 +157,9 @@ func (web *WebAddon) Response(f *proxy.Flow) {
 		return newMessageFlow(messageTypeResponseBody, f)
 	})
 }
+
+func (web *WebAddon) ServerDisconnected(connCtx *proxy.ConnContext) {
+	web.forEachConn(func(c *concurrentConn) {
+		c.whenConnClose(connCtx)
+	})
+}