chore: message protocol backend

addon-dailer
liqiang 4 years ago
parent 6ab1b82680
commit c07adef09c

@ -1,14 +1,18 @@
package web package web
import ( import (
"encoding/json"
"strings"
"sync" "sync"
"github.com/gorilla/websocket" "github.com/gorilla/websocket"
"github.com/lqqyt2423/go-mitmproxy/flow" "github.com/lqqyt2423/go-mitmproxy/flow"
) )
type breakPointRule struct {
Method string `json:"method"`
URL string `json:"url"`
Action int `json:"action"` // 1 - change request 2 - change response 3 - both
}
type concurrentConn struct { type concurrentConn struct {
conn *websocket.Conn conn *websocket.Conn
mu sync.Mutex mu sync.Mutex
@ -16,7 +20,7 @@ type concurrentConn struct {
waitChans map[string]chan interface{} waitChans map[string]chan interface{}
waitChansMu sync.Mutex waitChansMu sync.Mutex
interceptUri string breakPointRules []*breakPointRule
} }
func newConn(c *websocket.Conn) *concurrentConn { func newConn(c *websocket.Conn) *concurrentConn {
@ -26,7 +30,7 @@ func newConn(c *websocket.Conn) *concurrentConn {
} }
} }
func (c *concurrentConn) writeMessage(msg *message, f *flow.Flow) { func (c *concurrentConn) writeMessage(msg *messageFlow, f *flow.Flow) {
if c.isIntercpt(f, msg) { if c.isIntercpt(f, msg) {
msg.waitIntercept = 1 msg.waitIntercept = 1
} }
@ -63,27 +67,15 @@ func (c *concurrentConn) readloop() {
continue continue
} }
if msg.mType == messageTypeChangeInterceptUri { if msgEdit, ok := msg.(*messageEdit); ok {
interceptUri := "" ch := c.initWaitChan(msgEdit.id.String())
if len(msg.content) > 0 { go func(m *messageEdit, ch chan<- interface{}) {
interceptUri = string(msg.content) ch <- m
} }(msgEdit, ch)
c.interceptUri = interceptUri } else if msgMeta, ok := msg.(*messageMeta); ok {
continue c.breakPointRules = msgMeta.breakPointRules
} } else {
log.Warn("invalid message, skip")
if msg.mType == messageTypeChangeRequest {
req := new(flow.Request)
err := json.Unmarshal(msg.content, req)
if err != nil {
log.Error(err)
continue
}
ch := c.initWaitChan(msg.id.String())
go func(req *flow.Request, ch chan<- interface{}) {
ch <- req
}(req, ch)
} }
} }
} }
@ -101,28 +93,24 @@ func (c *concurrentConn) initWaitChan(key string) chan interface{} {
} }
// 是否拦截 // 是否拦截
func (c *concurrentConn) isIntercpt(f *flow.Flow, after *message) bool { func (c *concurrentConn) isIntercpt(f *flow.Flow, after *messageFlow) bool {
if after.mType != messageTypeRequest { if after.mType != messageTypeRequest {
return false return false
} }
if c.interceptUri == "" {
return false
}
if strings.Contains(f.Request.URL.String(), c.interceptUri) {
return true
}
return false return false
} }
// 拦截 // 拦截
func (c *concurrentConn) waitIntercept(f *flow.Flow, after *message) { func (c *concurrentConn) waitIntercept(f *flow.Flow, after *messageFlow) {
log.Infof("waiting Intercept: %s\n", f.Request.URL) log.Infof("waiting Intercept: %s\n", f.Request.URL)
ch := c.initWaitChan(f.Id.String()) ch := c.initWaitChan(f.Id.String())
req := (<-ch).(*flow.Request) msg := (<-ch).(*messageEdit)
log.Infof("waited Intercept: %s\n", f.Request.URL) log.Infof("waited Intercept: %s\n", f.Request.URL)
f.Request.Method = req.Method // f.Request.Method = req.Method
f.Request.URL = req.URL // f.Request.URL = req.URL
f.Request.Header = req.Header // f.Request.Header = req.Header
log.Infof("waitIntercept: %v", msg)
} }

@ -2,12 +2,28 @@ package web
import ( import (
"bytes" "bytes"
"encoding/binary"
"encoding/json" "encoding/json"
"errors"
"github.com/lqqyt2423/go-mitmproxy/flow" "github.com/lqqyt2423/go-mitmproxy/flow"
uuid "github.com/satori/go.uuid" uuid "github.com/satori/go.uuid"
) )
// message:
// type: 1/2/3/4
// messageFlow
// version 1 byte + type 1 byte + id 36 byte + waitIntercept 1 byte + content left bytes
// type: 11/12
// messageEdit
// version 1 byte + type 1 byte + id 36 byte + header len 4 byte + header content bytes + body len 4 byte + [body content bytes]
// type: 21
// messageMeta
// version 1 byte + type 1 byte + content left bytes
const messageVersion = 1 const messageVersion = 1
type messageType byte type messageType byte
@ -18,9 +34,10 @@ const (
messageTypeResponse messageType = 3 messageTypeResponse messageType = 3
messageTypeResponseBody messageType = 4 messageTypeResponseBody messageType = 4
messageTypeChangeRequest messageType = 11 messageTypeChangeRequest messageType = 11
messageTypeChangeResponse messageType = 12
messageTypeChangeInterceptUri messageType = 21 messageTypeChangeBreakPointRules messageType = 21
) )
var allMessageTypes = []messageType{ var allMessageTypes = []messageType{
@ -29,7 +46,8 @@ var allMessageTypes = []messageType{
messageTypeResponse, messageTypeResponse,
messageTypeResponseBody, messageTypeResponseBody,
messageTypeChangeRequest, messageTypeChangeRequest,
messageTypeChangeInterceptUri, messageTypeChangeResponse,
messageTypeChangeBreakPointRules,
} }
func validMessageType(t byte) bool { func validMessageType(t byte) bool {
@ -41,73 +59,207 @@ func validMessageType(t byte) bool {
return false return false
} }
type message struct { type message interface {
bytes() []byte
}
type messageFlow struct {
mType messageType mType messageType
id uuid.UUID id uuid.UUID
waitIntercept byte waitIntercept byte
content []byte content []byte
} }
func newMessage(mType messageType, id uuid.UUID, content []byte) *message { func newMessageFlow(mType messageType, f *flow.Flow) *messageFlow {
return &message{ var content []byte
var err error = nil
if mType == messageTypeRequest {
content, err = json.Marshal(f.Request)
} else if mType == messageTypeRequestBody {
content = f.Request.Body
} else if mType == messageTypeResponse {
content, err = json.Marshal(f.Response)
} else if mType == messageTypeResponseBody {
content, err = f.Response.DecodedBody()
} else {
panic(errors.New("invalid message type"))
}
if err != nil {
panic(err)
}
return &messageFlow{
mType: mType, mType: mType,
id: id, id: f.Id,
content: content, content: content,
} }
} }
func parseMessage(data []byte) *message { func (m *messageFlow) bytes() []byte {
if len(data) < 39 { buf := bytes.NewBuffer(make([]byte, 0))
buf.WriteByte(byte(messageVersion))
buf.WriteByte(byte(m.mType))
buf.WriteString(m.id.String()) // len: 36
buf.WriteByte(m.waitIntercept)
buf.Write(m.content)
return buf.Bytes()
}
type messageEdit struct {
mType messageType
id uuid.UUID
request *flow.Request
response *flow.Response
}
func parseMessageEdit(data []byte) *messageEdit {
// 2 + 36 + 4 + 4
if len(data) < 46 {
return nil return nil
} }
if data[0] != messageVersion {
mType := (messageType)(data[1])
id, err := uuid.FromString(string(data[2:38]))
if err != nil {
return nil return nil
} }
if !validMessageType(data[1]) {
hl := (int)(binary.BigEndian.Uint32(data[38:42]))
if 42+hl+4 > len(data) {
return nil return nil
} }
headerContent := data[42 : 42+hl]
id, err := uuid.FromString(string(data[3:39])) bl := (int)(binary.BigEndian.Uint32(data[42+hl : 42+hl+4]))
if err != nil { if 42+hl+4+bl != len(data) {
return nil
}
bodyContent := data[42+hl+4:]
msg := &messageEdit{
mType: mType,
id: id,
}
if mType == messageTypeChangeRequest {
req := new(flow.Request)
err := json.Unmarshal(headerContent, req)
if err != nil {
return nil
}
req.Body = bodyContent
msg.request = req
} else if mType == messageTypeChangeResponse {
res := new(flow.Response)
err := json.Unmarshal(headerContent, res)
if err != nil {
return nil
}
res.Body = bodyContent
msg.response = res
} else {
return nil return nil
} }
msg := newMessage(messageType(data[1]), id, data[39:])
msg.waitIntercept = data[2]
return msg return msg
} }
func newMessageRequest(f *flow.Flow) *message { func (m *messageEdit) bytes() []byte {
content, err := json.Marshal(f.Request) buf := bytes.NewBuffer(make([]byte, 0))
if err != nil { buf.WriteByte(byte(messageVersion))
panic(err) buf.WriteByte(byte(m.mType))
buf.WriteString(m.id.String()) // len: 36
if m.mType == messageTypeChangeRequest {
headerContent, err := json.Marshal(m.request)
if err != nil {
panic(err)
}
hl := make([]byte, 4)
binary.BigEndian.PutUint32(hl, (uint32)(len(headerContent)))
buf.Write(hl)
bodyContent := m.request.Body
bl := make([]byte, 4)
binary.BigEndian.PutUint32(bl, (uint32)(len(bodyContent)))
buf.Write(bl)
buf.Write(bodyContent)
} else if m.mType == messageTypeChangeResponse {
headerContent, err := json.Marshal(m.response)
if err != nil {
panic(err)
}
hl := make([]byte, 4)
binary.BigEndian.PutUint32(hl, (uint32)(len(headerContent)))
buf.Write(hl)
bodyContent := m.response.Body
bl := make([]byte, 4)
binary.BigEndian.PutUint32(bl, (uint32)(len(bodyContent)))
buf.Write(bl)
buf.Write(bodyContent)
} }
return newMessage(messageTypeRequest, f.Id, content)
return buf.Bytes()
} }
func newMessageRequestBody(f *flow.Flow) *message { type messageMeta struct {
return newMessage(messageTypeRequestBody, f.Id, f.Request.Body) mType messageType
breakPointRules []*breakPointRule
} }
func newMessageResponse(f *flow.Flow) *message { func parseMessageMeta(data []byte) *messageMeta {
content, err := json.Marshal(f.Response) content := data[2:]
rules := make([]*breakPointRule, 0)
err := json.Unmarshal(content, &rules)
if err != nil { if err != nil {
panic(err) return nil
} }
return newMessage(messageTypeResponse, f.Id, content)
}
func newMessageResponseBody(f *flow.Flow) *message { return &messageMeta{
body, _ := f.Response.DecodedBody() mType: messageType(data[1]),
return newMessage(messageTypeResponseBody, f.Id, body) breakPointRules: rules,
}
} }
func (m *message) bytes() []byte { func (m *messageMeta) bytes() []byte {
buf := bytes.NewBuffer(make([]byte, 0)) buf := bytes.NewBuffer(make([]byte, 0))
buf.WriteByte(byte(messageVersion)) buf.WriteByte(byte(messageVersion))
buf.WriteByte(byte(m.mType)) buf.WriteByte(byte(m.mType))
buf.WriteByte(m.waitIntercept)
buf.WriteString(m.id.String()) // len: 36 content, err := json.Marshal(m.breakPointRules)
buf.Write(m.content) if err != nil {
panic(err)
}
buf.Write(content)
return buf.Bytes() return buf.Bytes()
} }
func parseMessage(data []byte) message {
if len(data) < 2 {
return nil
}
if data[0] != messageVersion {
return nil
}
if !validMessageType(data[1]) {
return nil
}
mType := (messageType)(data[1])
if mType == messageTypeChangeRequest || mType == messageTypeChangeResponse {
return parseMessageEdit(data)
} else if mType == messageTypeChangeBreakPointRules {
return parseMessageMeta(data)
} else {
log.Warnf("invalid message type %v", mType)
return nil
}
}

@ -96,7 +96,7 @@ func (web *WebAddon) removeConn(conn *concurrentConn) {
web.conns = append(web.conns[:index], web.conns[index+1:]...) web.conns = append(web.conns[:index], web.conns[index+1:]...)
} }
func (web *WebAddon) sendFlow(f *flow.Flow, msgFn func() *message) bool { func (web *WebAddon) sendFlow(f *flow.Flow, msgFn func() *messageFlow) bool {
web.connsMu.RLock() web.connsMu.RLock()
conns := web.conns conns := web.conns
web.connsMu.RUnlock() web.connsMu.RUnlock()
@ -114,25 +114,25 @@ func (web *WebAddon) sendFlow(f *flow.Flow, msgFn func() *message) bool {
} }
func (web *WebAddon) Requestheaders(f *flow.Flow) { func (web *WebAddon) Requestheaders(f *flow.Flow) {
web.sendFlow(f, func() *message { web.sendFlow(f, func() *messageFlow {
return newMessageRequest(f) return newMessageFlow(messageTypeRequest, f)
}) })
} }
func (web *WebAddon) Request(f *flow.Flow) { func (web *WebAddon) Request(f *flow.Flow) {
web.sendFlow(f, func() *message { web.sendFlow(f, func() *messageFlow {
return newMessageRequestBody(f) return newMessageFlow(messageTypeRequestBody, f)
}) })
} }
func (web *WebAddon) Responseheaders(f *flow.Flow) { func (web *WebAddon) Responseheaders(f *flow.Flow) {
web.sendFlow(f, func() *message { web.sendFlow(f, func() *messageFlow {
return newMessageResponse(f) return newMessageFlow(messageTypeResponse, f)
}) })
} }
func (web *WebAddon) Response(f *flow.Flow) { func (web *WebAddon) Response(f *flow.Flow) {
web.sendFlow(f, func() *message { web.sendFlow(f, func() *messageFlow {
return newMessageResponseBody(f) return newMessageFlow(messageTypeResponseBody, f)
}) })
} }

@ -28,7 +28,6 @@ func (req *Request) MarshalJSON() ([]byte, error) {
r["url"] = req.URL.String() r["url"] = req.URL.String()
r["proto"] = req.Proto r["proto"] = req.Proto
r["header"] = req.Header r["header"] = req.Header
r["body"] = req.Body
return json.Marshal(r) return json.Marshal(r)
} }

Loading…
Cancel
Save