diff --git a/addon/web/message.go b/addon/web/message.go new file mode 100644 index 0000000..50f2ca1 --- /dev/null +++ b/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() +} diff --git a/addon/web/web.go b/addon/web/web.go index bb6acf0..fec9c94 100644 --- a/addon/web/web.go +++ b/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) - b, err := json.Marshal(msg) - if err != nil { - log.Error(err) - return - } + msg := msgFn() 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) + }) } diff --git a/flow/flow.go b/flow/flow.go index db249ac..9c2f6aa 100644 --- a/flow/flow.go +++ b/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.