websocket message

addon/web/message.go

@@ -0,0 +1,60 @@
+package web
+
+import (
+	"bytes"
+	"encoding/json"
+
+	"github.com/lqqyt2423/go-mitmproxy/flow"
+	uuid "github.com/satori/go.uuid"
+)
+
+const messageVersion = 1
+
+const (
+	messageTypeRequest      = 1
+	messageTypeResponse     = 2
+	messageTypeResponseBody = 3
+)
+
+type message struct {
+	messageType int
+	id          uuid.UUID
+	content     []byte
+}
+
+func newMessage(messageType int, id uuid.UUID, content []byte) *message {
+	return &message{
+		messageType: messageType,
+		id:          id,
+		content:     content,
+	}
+}
+
+func newMessageRequest(f *flow.Flow) *message {
+	content, err := json.Marshal(f.Request)
+	if err != nil {
+		panic(err)
+	}
+	return newMessage(messageTypeRequest, f.Id, content)
+}
+
+func newMessageResponse(f *flow.Flow) *message {
+	content, err := json.Marshal(f.Response)
+	if err != nil {
+		panic(err)
+	}
+	return newMessage(messageTypeResponse, f.Id, content)
+}
+
+func newMessageResponseBody(f *flow.Flow) *message {
+	return newMessage(messageTypeResponseBody, f.Id, f.Response.Body)
+}
+
+func (m *message) bytes() []byte {
+	buf := bytes.NewBuffer(make([]byte, 0))
+	buf.WriteByte(byte(messageVersion))
+	buf.WriteByte(byte(m.messageType))
+	buf.WriteString(m.id.String()) // len: 36
+	buf.Write(m.content)
+	return buf.Bytes()
+}

addon/web/web.go

@@ -1,7 +1,6 @@
 package web
 
 import (
-	"encoding/json"
 	"net/http"
 	"sync"
 
@@ -57,18 +56,6 @@ type WebAddon struct {
 	connsMu sync.RWMutex
 }
 
-type message struct {
-	On   string     `json:"on"`
-	Flow *flow.Flow `json:"flow"`
-}
-
-func newMessage(on string, f *flow.Flow) *message {
-	return &message{
-		On:   on,
-		Flow: f,
-	}
-}
-
 func NewWebAddon() *WebAddon {
 	web := new(WebAddon)
 	web.addr = ":9081"
@@ -119,7 +106,7 @@ func (web *WebAddon) removeConn(conn *websocket.Conn) {
 	web.conns = append(web.conns[:index], web.conns[index+1:]...)
 }
 
-func (web *WebAddon) sendFlow(on string, f *flow.Flow) {
+func (web *WebAddon) sendFlow(msgFn func() *message) {
 	web.connsMu.RLock()
 	conns := web.conns
 	web.connsMu.RUnlock()
@@ -128,23 +115,28 @@ func (web *WebAddon) sendFlow(on string, f *flow.Flow) {
 		return
 	}
 
-	msg := newMessage(on, f)
+	msg := msgFn()
-	b, err := json.Marshal(msg)
-	if err != nil {
-		log.Error(err)
-		return
-	}
 	for _, c := range conns {
 		c.mu.Lock()
-		c.conn.WriteMessage(websocket.TextMessage, b)
+		c.conn.WriteMessage(websocket.BinaryMessage, msg.bytes())
 		c.mu.Unlock()
 	}
 }
 
 func (web *WebAddon) Request(f *flow.Flow) {
-	web.sendFlow("request", f)
+	web.sendFlow(func() *message {
+		return newMessageRequest(f)
+	})
+}
+
+func (web *WebAddon) Responseheaders(f *flow.Flow) {
+	web.sendFlow(func() *message {
+		return newMessageResponse(f)
+	})
 }
 
 func (web *WebAddon) Response(f *flow.Flow) {
-	web.sendFlow("response", f)
+	web.sendFlow(func() *message {
+		return newMessageResponseBody(f)
+	})
 }

flow/flow.go

@@ -48,7 +48,7 @@ func (r *Request) Raw() *http.Request {
 type Response struct {
 	StatusCode int         `json:"statusCode"`
 	Header     http.Header `json:"header"`
-	Body       []byte      `json:"body"`
+	Body       []byte      `json:"-"`
 
 	decodedBody []byte
 	decoded     bool // decoded reports whether the response was sent compressed but was decoded to decodedBody.