This is a fork of:
https://github.com/roglew/puppy
You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
837 lines
20 KiB
837 lines
20 KiB
package main |
|
|
|
import ( |
|
"crypto/tls" |
|
"encoding/base64" |
|
"fmt" |
|
"log" |
|
"net" |
|
"net/http" |
|
"sync" |
|
"time" |
|
|
|
"github.com/gorilla/websocket" |
|
) |
|
|
|
var getNextSubId = IdCounter() |
|
var getNextStorageId = IdCounter() |
|
|
|
// Working on using this for webui |
|
type proxyWebUIHandler func(http.ResponseWriter, *http.Request, *InterceptingProxy) |
|
|
|
type savedStorage struct { |
|
storage MessageStorage |
|
description string |
|
} |
|
|
|
type InterceptingProxy struct { |
|
slistener *ProxyListener |
|
server *http.Server |
|
mtx sync.Mutex |
|
logger *log.Logger |
|
proxyStorage int |
|
netDial NetDialer |
|
|
|
usingProxy bool |
|
proxyHost string |
|
proxyPort int |
|
proxyIsSOCKS bool |
|
proxyCreds *ProxyCredentials |
|
|
|
requestInterceptor RequestInterceptor |
|
responseInterceptor ResponseInterceptor |
|
wSInterceptor WSInterceptor |
|
scopeChecker RequestChecker |
|
scopeQuery MessageQuery |
|
|
|
reqSubs []*ReqIntSub |
|
rspSubs []*RspIntSub |
|
wsSubs []*WSIntSub |
|
|
|
httpHandlers map[string]proxyWebUIHandler |
|
|
|
messageStorage map[int]*savedStorage |
|
} |
|
|
|
type ProxyCredentials struct { |
|
Username string |
|
Password string |
|
} |
|
|
|
type RequestInterceptor func(req *ProxyRequest) (*ProxyRequest, error) |
|
type ResponseInterceptor func(req *ProxyRequest, rsp *ProxyResponse) (*ProxyResponse, error) |
|
type WSInterceptor func(req *ProxyRequest, rsp *ProxyResponse, msg *ProxyWSMessage) (*ProxyWSMessage, error) |
|
|
|
type ReqIntSub struct { |
|
id int |
|
Interceptor RequestInterceptor |
|
} |
|
|
|
type RspIntSub struct { |
|
id int |
|
Interceptor ResponseInterceptor |
|
} |
|
|
|
type WSIntSub struct { |
|
id int |
|
Interceptor WSInterceptor |
|
} |
|
|
|
func (creds *ProxyCredentials) SerializeHeader() string { |
|
toEncode := []byte(fmt.Sprintf("%s:%s", creds.Username, creds.Password)) |
|
encoded := base64.StdEncoding.EncodeToString(toEncode) |
|
return fmt.Sprintf("Basic %s", encoded) |
|
} |
|
|
|
func NewInterceptingProxy(logger *log.Logger) *InterceptingProxy { |
|
var iproxy InterceptingProxy |
|
iproxy.messageStorage = make(map[int]*savedStorage) |
|
iproxy.slistener = NewProxyListener(logger) |
|
iproxy.server = newProxyServer(logger, &iproxy) |
|
iproxy.logger = logger |
|
iproxy.httpHandlers = make(map[string]proxyWebUIHandler) |
|
|
|
go func() { |
|
iproxy.server.Serve(iproxy.slistener) |
|
}() |
|
return &iproxy |
|
} |
|
|
|
func (iproxy *InterceptingProxy) Close() { |
|
// Closes all associated listeners, but does not shut down the server because there is no way to gracefully shut down an http server yet :| |
|
// Will throw errors when the server finally shuts down and tries to call iproxy.slistener.Close a second time |
|
iproxy.mtx.Lock() |
|
defer iproxy.mtx.Unlock() |
|
iproxy.slistener.Close() |
|
//iproxy.server.Close() // Coming eventually... I hope |
|
} |
|
|
|
func (iproxy *InterceptingProxy) SetCACertificate(caCert *tls.Certificate) { |
|
if iproxy.slistener == nil { |
|
panic("intercepting proxy does not have a proxy listener") |
|
} |
|
iproxy.slistener.SetCACertificate(caCert) |
|
} |
|
|
|
func (iproxy *InterceptingProxy) GetCACertificate() *tls.Certificate { |
|
return iproxy.slistener.GetCACertificate() |
|
} |
|
|
|
func (iproxy *InterceptingProxy) AddListener(l net.Listener) { |
|
iproxy.mtx.Lock() |
|
defer iproxy.mtx.Unlock() |
|
iproxy.slistener.AddListener(l) |
|
} |
|
|
|
func (iproxy *InterceptingProxy) RemoveListener(l net.Listener) { |
|
iproxy.mtx.Lock() |
|
defer iproxy.mtx.Unlock() |
|
iproxy.slistener.RemoveListener(l) |
|
} |
|
|
|
func (iproxy *InterceptingProxy) GetMessageStorage(id int) (MessageStorage, string) { |
|
iproxy.mtx.Lock() |
|
defer iproxy.mtx.Unlock() |
|
savedStorage, ok := iproxy.messageStorage[id] |
|
if !ok { |
|
return nil, "" |
|
} |
|
return savedStorage.storage, savedStorage.description |
|
} |
|
|
|
func (iproxy *InterceptingProxy) AddMessageStorage(storage MessageStorage, description string) int { |
|
iproxy.mtx.Lock() |
|
defer iproxy.mtx.Unlock() |
|
id := getNextStorageId() |
|
iproxy.messageStorage[id] = &savedStorage{storage, description} |
|
return id |
|
} |
|
|
|
func (iproxy *InterceptingProxy) CloseMessageStorage(id int) { |
|
iproxy.mtx.Lock() |
|
defer iproxy.mtx.Unlock() |
|
savedStorage, ok := iproxy.messageStorage[id] |
|
if !ok { |
|
return |
|
} |
|
delete(iproxy.messageStorage, id) |
|
savedStorage.storage.Close() |
|
} |
|
|
|
type SavedStorage struct { |
|
Id int |
|
Storage MessageStorage |
|
Description string |
|
} |
|
|
|
func (iproxy *InterceptingProxy) ListMessageStorage() []*SavedStorage { |
|
iproxy.mtx.Lock() |
|
defer iproxy.mtx.Unlock() |
|
|
|
r := make([]*SavedStorage, 0) |
|
for id, ss := range iproxy.messageStorage { |
|
r = append(r, &SavedStorage{id, ss.storage, ss.description}) |
|
} |
|
return r |
|
} |
|
|
|
func (iproxy *InterceptingProxy) getRequestSubs() []*ReqIntSub { |
|
iproxy.mtx.Lock() |
|
defer iproxy.mtx.Unlock() |
|
return iproxy.reqSubs |
|
} |
|
|
|
func (iproxy *InterceptingProxy) getResponseSubs() []*RspIntSub { |
|
iproxy.mtx.Lock() |
|
defer iproxy.mtx.Unlock() |
|
return iproxy.rspSubs |
|
} |
|
|
|
func (iproxy *InterceptingProxy) getWSSubs() []*WSIntSub { |
|
iproxy.mtx.Lock() |
|
defer iproxy.mtx.Unlock() |
|
return iproxy.wsSubs |
|
} |
|
|
|
func (iproxy *InterceptingProxy) LoadScope(storageId int) error { |
|
// Try and set the scope |
|
savedStorage, ok := iproxy.messageStorage[storageId] |
|
if !ok { |
|
return fmt.Errorf("proxy has no associated storage") |
|
} |
|
iproxy.logger.Println("loading scope") |
|
if scope, err := savedStorage.storage.LoadQuery("__scope"); err == nil { |
|
if err := iproxy.setScopeQuery(scope); err != nil { |
|
iproxy.logger.Println("error setting scope:", err.Error()) |
|
} |
|
} else { |
|
iproxy.logger.Println("error loading scope:", err.Error()) |
|
} |
|
return nil |
|
} |
|
|
|
func (iproxy *InterceptingProxy) GetScopeChecker() RequestChecker { |
|
iproxy.mtx.Lock() |
|
defer iproxy.mtx.Unlock() |
|
return iproxy.scopeChecker |
|
} |
|
|
|
func (iproxy *InterceptingProxy) SetScopeChecker(checker RequestChecker) error { |
|
iproxy.mtx.Lock() |
|
defer iproxy.mtx.Unlock() |
|
savedStorage, ok := iproxy.messageStorage[iproxy.proxyStorage] |
|
if !ok { |
|
savedStorage = nil |
|
} |
|
iproxy.scopeChecker = checker |
|
iproxy.scopeQuery = nil |
|
emptyQuery := make(MessageQuery, 0) |
|
if savedStorage != nil { |
|
savedStorage.storage.SaveQuery("__scope", emptyQuery) // Assume it clears it I guess |
|
} |
|
return nil |
|
} |
|
|
|
func (iproxy *InterceptingProxy) GetScopeQuery() MessageQuery { |
|
iproxy.mtx.Lock() |
|
defer iproxy.mtx.Unlock() |
|
return iproxy.scopeQuery |
|
} |
|
|
|
func (iproxy *InterceptingProxy) SetScopeQuery(query MessageQuery) error { |
|
iproxy.mtx.Lock() |
|
defer iproxy.mtx.Unlock() |
|
return iproxy.setScopeQuery(query) |
|
} |
|
|
|
func (iproxy *InterceptingProxy) setScopeQuery(query MessageQuery) error { |
|
checker, err := CheckerFromMessageQuery(query) |
|
if err != nil { |
|
return err |
|
} |
|
savedStorage, ok := iproxy.messageStorage[iproxy.proxyStorage] |
|
if !ok { |
|
savedStorage = nil |
|
} |
|
iproxy.scopeChecker = checker |
|
iproxy.scopeQuery = query |
|
if savedStorage != nil { |
|
if err = savedStorage.storage.SaveQuery("__scope", query); err != nil { |
|
return fmt.Errorf("could not save scope to storage: %s", err.Error()) |
|
} |
|
} |
|
|
|
return nil |
|
} |
|
|
|
func (iproxy *InterceptingProxy) SetNetDial(dialer NetDialer) { |
|
iproxy.mtx.Lock() |
|
defer iproxy.mtx.Unlock() |
|
iproxy.netDial = dialer |
|
} |
|
|
|
func (iproxy *InterceptingProxy) NetDial() NetDialer { |
|
iproxy.mtx.Lock() |
|
defer iproxy.mtx.Unlock() |
|
return iproxy.netDial |
|
} |
|
|
|
func (iproxy *InterceptingProxy) ClearUpstreamProxy() { |
|
iproxy.mtx.Lock() |
|
defer iproxy.mtx.Unlock() |
|
iproxy.usingProxy = false |
|
iproxy.proxyHost = "" |
|
iproxy.proxyPort = 0 |
|
iproxy.proxyIsSOCKS = false |
|
} |
|
|
|
func (iproxy *InterceptingProxy) SetUpstreamProxy(proxyHost string, proxyPort int, creds *ProxyCredentials) { |
|
iproxy.mtx.Lock() |
|
defer iproxy.mtx.Unlock() |
|
iproxy.usingProxy = true |
|
iproxy.proxyHost = proxyHost |
|
iproxy.proxyPort = proxyPort |
|
iproxy.proxyIsSOCKS = false |
|
iproxy.proxyCreds = creds |
|
} |
|
|
|
func (iproxy *InterceptingProxy) SetUpstreamSOCKSProxy(proxyHost string, proxyPort int, creds *ProxyCredentials) { |
|
iproxy.mtx.Lock() |
|
defer iproxy.mtx.Unlock() |
|
iproxy.usingProxy = true |
|
iproxy.proxyHost = proxyHost |
|
iproxy.proxyPort = proxyPort |
|
iproxy.proxyIsSOCKS = true |
|
iproxy.proxyCreds = creds |
|
} |
|
|
|
func (iproxy *InterceptingProxy) ClearScope() error { |
|
iproxy.mtx.Lock() |
|
defer iproxy.mtx.Unlock() |
|
iproxy.scopeChecker = nil |
|
iproxy.scopeChecker = nil |
|
emptyQuery := make(MessageQuery, 0) |
|
savedStorage, ok := iproxy.messageStorage[iproxy.proxyStorage] |
|
if !ok { |
|
savedStorage = nil |
|
} |
|
if savedStorage != nil { |
|
if err := savedStorage.storage.SaveQuery("__scope", emptyQuery); err != nil { |
|
return fmt.Errorf("could not clear scope in storage: %s", err.Error()) |
|
} |
|
} |
|
return nil |
|
} |
|
|
|
func (iproxy *InterceptingProxy) SubmitRequest(req *ProxyRequest) error { |
|
oldDial := req.NetDial |
|
defer func() { req.NetDial = oldDial }() |
|
req.NetDial = iproxy.NetDial() |
|
|
|
if iproxy.usingProxy { |
|
if iproxy.proxyIsSOCKS { |
|
return SubmitRequestSOCKSProxy(req, iproxy.proxyHost, iproxy.proxyPort, iproxy.proxyCreds) |
|
} else { |
|
return SubmitRequestProxy(req, iproxy.proxyHost, iproxy.proxyPort, iproxy.proxyCreds) |
|
} |
|
} |
|
return SubmitRequest(req) |
|
} |
|
|
|
func (iproxy *InterceptingProxy) WSDial(req *ProxyRequest) (*WSSession, error) { |
|
oldDial := req.NetDial |
|
defer func() { req.NetDial = oldDial }() |
|
req.NetDial = iproxy.NetDial() |
|
|
|
if iproxy.usingProxy { |
|
if iproxy.proxyIsSOCKS { |
|
return WSDialSOCKSProxy(req, iproxy.proxyHost, iproxy.proxyPort, iproxy.proxyCreds) |
|
} else { |
|
return WSDialProxy(req, iproxy.proxyHost, iproxy.proxyPort, iproxy.proxyCreds) |
|
} |
|
} |
|
return WSDial(req) |
|
} |
|
|
|
func (iproxy *InterceptingProxy) AddReqIntSub(f RequestInterceptor) *ReqIntSub { |
|
iproxy.mtx.Lock() |
|
defer iproxy.mtx.Unlock() |
|
|
|
sub := &ReqIntSub{ |
|
id: getNextSubId(), |
|
Interceptor: f, |
|
} |
|
iproxy.reqSubs = append(iproxy.reqSubs, sub) |
|
return sub |
|
} |
|
|
|
func (iproxy *InterceptingProxy) RemoveReqIntSub(sub *ReqIntSub) { |
|
iproxy.mtx.Lock() |
|
defer iproxy.mtx.Unlock() |
|
|
|
for i, checkSub := range iproxy.reqSubs { |
|
if checkSub.id == sub.id { |
|
iproxy.reqSubs = append(iproxy.reqSubs[:i], iproxy.reqSubs[i+1:]...) |
|
return |
|
} |
|
} |
|
} |
|
|
|
func (iproxy *InterceptingProxy) AddRspIntSub(f ResponseInterceptor) *RspIntSub { |
|
iproxy.mtx.Lock() |
|
defer iproxy.mtx.Unlock() |
|
|
|
sub := &RspIntSub{ |
|
id: getNextSubId(), |
|
Interceptor: f, |
|
} |
|
iproxy.rspSubs = append(iproxy.rspSubs, sub) |
|
return sub |
|
} |
|
|
|
func (iproxy *InterceptingProxy) RemoveRspIntSub(sub *RspIntSub) { |
|
iproxy.mtx.Lock() |
|
defer iproxy.mtx.Unlock() |
|
|
|
for i, checkSub := range iproxy.rspSubs { |
|
if checkSub.id == sub.id { |
|
iproxy.rspSubs = append(iproxy.rspSubs[:i], iproxy.rspSubs[i+1:]...) |
|
return |
|
} |
|
} |
|
} |
|
|
|
func (iproxy *InterceptingProxy) AddWSIntSub(f WSInterceptor) *WSIntSub { |
|
iproxy.mtx.Lock() |
|
defer iproxy.mtx.Unlock() |
|
|
|
sub := &WSIntSub{ |
|
id: getNextSubId(), |
|
Interceptor: f, |
|
} |
|
iproxy.wsSubs = append(iproxy.wsSubs, sub) |
|
return sub |
|
} |
|
|
|
func (iproxy *InterceptingProxy) RemoveWSIntSub(sub *WSIntSub) { |
|
iproxy.mtx.Lock() |
|
defer iproxy.mtx.Unlock() |
|
|
|
for i, checkSub := range iproxy.wsSubs { |
|
if checkSub.id == sub.id { |
|
iproxy.wsSubs = append(iproxy.wsSubs[:i], iproxy.wsSubs[i+1:]...) |
|
return |
|
} |
|
} |
|
} |
|
|
|
func (iproxy *InterceptingProxy) SetProxyStorage(storageId int) error { |
|
iproxy.mtx.Lock() |
|
defer iproxy.mtx.Unlock() |
|
|
|
iproxy.proxyStorage = storageId |
|
|
|
_, ok := iproxy.messageStorage[iproxy.proxyStorage] |
|
if !ok { |
|
return fmt.Errorf("no storage with id %d", storageId) |
|
} |
|
|
|
iproxy.LoadScope(storageId) |
|
return nil |
|
} |
|
|
|
func (iproxy *InterceptingProxy) GetProxyStorage() MessageStorage { |
|
iproxy.mtx.Lock() |
|
defer iproxy.mtx.Unlock() |
|
|
|
savedStorage, ok := iproxy.messageStorage[iproxy.proxyStorage] |
|
if !ok { |
|
return nil |
|
} |
|
return savedStorage.storage |
|
} |
|
|
|
func (iproxy *InterceptingProxy) AddHTTPHandler(host string, handler proxyWebUIHandler) { |
|
iproxy.mtx.Lock() |
|
defer iproxy.mtx.Unlock() |
|
iproxy.httpHandlers[host] = handler |
|
} |
|
|
|
func (iproxy *InterceptingProxy) GetHTTPHandler(host string) (proxyWebUIHandler, error) { |
|
iproxy.mtx.Lock() |
|
defer iproxy.mtx.Unlock() |
|
handler, ok := iproxy.httpHandlers[host] |
|
if !ok { |
|
return nil, fmt.Errorf("no handler for host %s", host) |
|
} |
|
return handler, nil |
|
} |
|
|
|
func (iproxy *InterceptingProxy) RemoveHTTPHandler(host string) { |
|
iproxy.mtx.Lock() |
|
defer iproxy.mtx.Unlock() |
|
delete(iproxy.httpHandlers, host) |
|
} |
|
|
|
func ParseProxyRequest(r *http.Request) (*ProxyRequest, error) { |
|
host, port, useTLS, err := DecodeRemoteAddr(r.RemoteAddr) |
|
if err != nil { |
|
return nil, nil |
|
} |
|
pr := NewProxyRequest(r, host, port, useTLS) |
|
return pr, nil |
|
} |
|
|
|
func BlankResponse(w http.ResponseWriter) { |
|
w.Header().Set("Connection", "close") |
|
w.Header().Set("Cache-control", "no-cache") |
|
w.Header().Set("Pragma", "no-cache") |
|
w.Header().Set("Cache-control", "no-store") |
|
w.Header().Set("X-Frame-Options", "DENY") |
|
w.WriteHeader(200) |
|
} |
|
|
|
func ErrResponse(w http.ResponseWriter, err error) { |
|
http.Error(w, err.Error(), http.StatusInternalServerError) |
|
} |
|
|
|
func (iproxy *InterceptingProxy) ServeHTTP(w http.ResponseWriter, r *http.Request) { |
|
handler, err := iproxy.GetHTTPHandler(r.Host) |
|
if err == nil { |
|
handler(w, r, iproxy) |
|
return |
|
} |
|
|
|
req, _ := ParseProxyRequest(r) |
|
iproxy.logger.Println("Received request to", req.FullURL().String()) |
|
req.StripProxyHeaders() |
|
|
|
ms := iproxy.GetProxyStorage() |
|
scopeChecker := iproxy.GetScopeChecker() |
|
|
|
// Helper functions |
|
checkScope := func(req *ProxyRequest) bool { |
|
if scopeChecker != nil { |
|
return scopeChecker(req) |
|
} |
|
return true |
|
} |
|
|
|
saveIfExists := func(req *ProxyRequest) error { |
|
if ms != nil && checkScope(req) { |
|
if err := UpdateRequest(ms, req); err != nil { |
|
return err |
|
} |
|
} |
|
return nil |
|
} |
|
|
|
/* |
|
functions to mangle messages using the iproxy's manglers |
|
each return the new message, whether it was modified, and an error |
|
*/ |
|
|
|
mangleRequest := func(req *ProxyRequest) (*ProxyRequest, bool, error) { |
|
newReq := req.Clone() |
|
reqSubs := iproxy.getRequestSubs() |
|
for _, sub := range reqSubs { |
|
var err error = nil |
|
newReq, err = sub.Interceptor(newReq) |
|
if err != nil { |
|
e := fmt.Errorf("error with request interceptor: %s", err) |
|
return nil, false, e |
|
} |
|
if newReq == nil { |
|
break |
|
} |
|
} |
|
|
|
if newReq != nil { |
|
newReq.StartDatetime = time.Now() |
|
if !req.Eq(newReq) { |
|
iproxy.logger.Println("Request modified by interceptor") |
|
return newReq, true, nil |
|
} |
|
} else { |
|
return nil, true, nil |
|
} |
|
return req, false, nil |
|
} |
|
|
|
mangleResponse := func(req *ProxyRequest, rsp *ProxyResponse) (*ProxyResponse, bool, error) { |
|
reqCopy := req.Clone() |
|
newRsp := rsp.Clone() |
|
rspSubs := iproxy.getResponseSubs() |
|
iproxy.logger.Printf("%d interceptors", len(rspSubs)) |
|
for _, sub := range rspSubs { |
|
iproxy.logger.Println("mangling rsp...") |
|
var err error = nil |
|
newRsp, err = sub.Interceptor(reqCopy, newRsp) |
|
if err != nil { |
|
e := fmt.Errorf("error with response interceptor: %s", err) |
|
return nil, false, e |
|
} |
|
if newRsp == nil { |
|
break |
|
} |
|
} |
|
|
|
if newRsp != nil { |
|
if !rsp.Eq(newRsp) { |
|
iproxy.logger.Println("Response for", req.FullURL(), "modified by interceptor") |
|
// it was mangled |
|
return newRsp, true, nil |
|
} |
|
} else { |
|
// it was dropped |
|
return nil, true, nil |
|
} |
|
|
|
// it wasn't changed |
|
return rsp, false, nil |
|
} |
|
|
|
mangleWS := func(req *ProxyRequest, rsp *ProxyResponse, ws *ProxyWSMessage) (*ProxyWSMessage, bool, error) { |
|
newMsg := ws.Clone() |
|
reqCopy := req.Clone() |
|
rspCopy := rsp.Clone() |
|
wsSubs := iproxy.getWSSubs() |
|
for _, sub := range wsSubs { |
|
var err error = nil |
|
newMsg, err = sub.Interceptor(reqCopy, rspCopy, newMsg) |
|
if err != nil { |
|
e := fmt.Errorf("error with ws interceptor: %s", err) |
|
return nil, false, e |
|
} |
|
if newMsg == nil { |
|
break |
|
} |
|
} |
|
|
|
if newMsg != nil { |
|
if !ws.Eq(newMsg) { |
|
newMsg.Timestamp = time.Now() |
|
newMsg.Direction = ws.Direction |
|
iproxy.logger.Println("Message modified by interceptor") |
|
return newMsg, true, nil |
|
} |
|
} else { |
|
return nil, true, nil |
|
} |
|
return ws, false, nil |
|
} |
|
|
|
req.StartDatetime = time.Now() |
|
|
|
if checkScope(req) { |
|
if err := saveIfExists(req); err != nil { |
|
ErrResponse(w, err) |
|
return |
|
} |
|
newReq, mangled, err := mangleRequest(req) |
|
if err != nil { |
|
ErrResponse(w, err) |
|
return |
|
} |
|
if mangled { |
|
if newReq == nil { |
|
req.ServerResponse = nil |
|
if err := saveIfExists(req); err != nil { |
|
ErrResponse(w, err) |
|
return |
|
} |
|
BlankResponse(w) |
|
return |
|
} |
|
newReq.Unmangled = req |
|
req = newReq |
|
req.StartDatetime = time.Now() |
|
if err := saveIfExists(req); err != nil { |
|
ErrResponse(w, err) |
|
return |
|
} |
|
} |
|
} |
|
|
|
if req.IsWSUpgrade() { |
|
iproxy.logger.Println("Detected websocket request. Upgrading...") |
|
|
|
rc, err := iproxy.WSDial(req) |
|
if err != nil { |
|
iproxy.logger.Println("error dialing ws server:", err) |
|
http.Error(w, fmt.Sprintf("error dialing websocket server: %s", err.Error()), http.StatusInternalServerError) |
|
return |
|
} |
|
defer rc.Close() |
|
req.EndDatetime = time.Now() |
|
if err := saveIfExists(req); err != nil { |
|
ErrResponse(w, err) |
|
return |
|
} |
|
|
|
var upgrader = websocket.Upgrader{ |
|
CheckOrigin: func(r *http.Request) bool { |
|
return true |
|
}, |
|
} |
|
|
|
lc, err := upgrader.Upgrade(w, r, nil) |
|
if err != nil { |
|
iproxy.logger.Println("error upgrading connection:", err) |
|
http.Error(w, fmt.Sprintf("error upgrading connection: %s", err.Error()), http.StatusInternalServerError) |
|
return |
|
} |
|
defer lc.Close() |
|
|
|
var wg sync.WaitGroup |
|
var reqMtx sync.Mutex |
|
addWSMessage := func(req *ProxyRequest, wsm *ProxyWSMessage) { |
|
reqMtx.Lock() |
|
defer reqMtx.Unlock() |
|
req.WSMessages = append(req.WSMessages, wsm) |
|
} |
|
|
|
// Get messages from server |
|
wg.Add(1) |
|
go func() { |
|
for { |
|
mtype, msg, err := rc.ReadMessage() |
|
if err != nil { |
|
lc.WriteMessage(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseNormalClosure, "")) |
|
iproxy.logger.Println("error with receiving server message:", err) |
|
wg.Done() |
|
return |
|
} |
|
pws, err := NewProxyWSMessage(mtype, msg, ToClient) |
|
if err != nil { |
|
iproxy.logger.Println("error creating ws object:", err.Error()) |
|
continue |
|
} |
|
pws.Timestamp = time.Now() |
|
|
|
if checkScope(req) { |
|
newMsg, mangled, err := mangleWS(req, req.ServerResponse, pws) |
|
if err != nil { |
|
iproxy.logger.Println("error mangling ws:", err) |
|
return |
|
} |
|
if mangled { |
|
if newMsg == nil { |
|
continue |
|
} else { |
|
newMsg.Unmangled = pws |
|
pws = newMsg |
|
pws.Request = nil |
|
} |
|
} |
|
} |
|
|
|
addWSMessage(req, pws) |
|
if err := saveIfExists(req); err != nil { |
|
iproxy.logger.Println("error saving request:", err) |
|
continue |
|
} |
|
lc.WriteMessage(pws.Type, pws.Message) |
|
} |
|
}() |
|
|
|
// Get messages from client |
|
wg.Add(1) |
|
go func() { |
|
for { |
|
mtype, msg, err := lc.ReadMessage() |
|
if err != nil { |
|
rc.WriteMessage(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseNormalClosure, "")) |
|
iproxy.logger.Println("error with receiving client message:", err) |
|
wg.Done() |
|
return |
|
} |
|
pws, err := NewProxyWSMessage(mtype, msg, ToServer) |
|
if err != nil { |
|
iproxy.logger.Println("error creating ws object:", err.Error()) |
|
continue |
|
} |
|
pws.Timestamp = time.Now() |
|
|
|
if checkScope(req) { |
|
newMsg, mangled, err := mangleWS(req, req.ServerResponse, pws) |
|
if err != nil { |
|
iproxy.logger.Println("error mangling ws:", err) |
|
return |
|
} |
|
if mangled { |
|
if newMsg == nil { |
|
continue |
|
} else { |
|
newMsg.Unmangled = pws |
|
pws = newMsg |
|
pws.Request = nil |
|
} |
|
} |
|
} |
|
|
|
addWSMessage(req, pws) |
|
if err := saveIfExists(req); err != nil { |
|
iproxy.logger.Println("error saving request:", err) |
|
continue |
|
} |
|
rc.WriteMessage(pws.Type, pws.Message) |
|
} |
|
}() |
|
wg.Wait() |
|
iproxy.logger.Println("Websocket session complete!") |
|
} else { |
|
err := iproxy.SubmitRequest(req) |
|
if err != nil { |
|
http.Error(w, fmt.Sprintf("error submitting request: %s", err.Error()), http.StatusInternalServerError) |
|
return |
|
} |
|
req.EndDatetime = time.Now() |
|
if err := saveIfExists(req); err != nil { |
|
ErrResponse(w, err) |
|
return |
|
} |
|
|
|
if checkScope(req) { |
|
newRsp, mangled, err := mangleResponse(req, req.ServerResponse) |
|
if err != nil { |
|
http.Error(w, err.Error(), http.StatusInternalServerError) |
|
return |
|
} |
|
if mangled { |
|
if newRsp == nil { |
|
req.ServerResponse = nil |
|
if err := saveIfExists(req); err != nil { |
|
ErrResponse(w, err) |
|
return |
|
} |
|
BlankResponse(w) |
|
return |
|
} |
|
newRsp.Unmangled = req.ServerResponse |
|
req.ServerResponse = newRsp |
|
if err := saveIfExists(req); err != nil { |
|
ErrResponse(w, err) |
|
return |
|
} |
|
} |
|
} |
|
|
|
for k, v := range req.ServerResponse.Header { |
|
for _, vv := range v { |
|
w.Header().Add(k, vv) |
|
} |
|
} |
|
w.WriteHeader(req.ServerResponse.StatusCode) |
|
w.Write(req.ServerResponse.BodyBytes()) |
|
return |
|
} |
|
} |
|
|
|
func newProxyServer(logger *log.Logger, iproxy *InterceptingProxy) *http.Server { |
|
server := &http.Server{ |
|
Handler: iproxy, |
|
ErrorLog: logger, |
|
} |
|
return server |
|
}
|
|
|