Skip to content

Commit

Permalink
chore(martian): extract common code of proxyConn.handleConnectRequest…
Browse files Browse the repository at this point in the history
…() and proxyHandler.handleConnectRequest() to a function

Add Proxy.Connect() that wraps connect() and also handles ConnectFunc and TLS termination.

Fixes #445
  • Loading branch information
mmatczuk committed Sep 12, 2024
1 parent e3c2c2f commit 2345fa3
Show file tree
Hide file tree
Showing 3 changed files with 32 additions and 61 deletions.
31 changes: 2 additions & 29 deletions internal/martian/proxy_conn.go
Original file line number Diff line number Diff line change
Expand Up @@ -203,35 +203,8 @@ func (p *proxyConn) handleConnectRequest(req *http.Request) error {
return p.handleMITM(req)
}

var (
res *http.Response
crw io.ReadWriteCloser
cerr error
)
if p.ConnectFunc != nil {
res, crw, cerr = p.ConnectFunc(req)
}
if p.ConnectFunc == nil || errors.Is(cerr, ErrConnectFallback) {
var cconn net.Conn
res, cconn, cerr = p.connect(req)

if cconn != nil {
defer cconn.Close()
crw = cconn

if terminateTLS {
log.Debugf(ctx, "attempting to terminate TLS on CONNECT tunnel: %s", req.URL.Host)
tconn := tls.Client(cconn, p.clientTLSConfig())
if err := tconn.Handshake(); err == nil {
crw = tconn
} else {
log.Errorf(ctx, "failed to terminate TLS on CONNECT tunnel: %v", err)
cerr = err
}
}
}
}

log.Debugf(ctx, "attempting to establish CONNECT tunnel: %s", req.URL.Host)
res, crw, cerr := p.Connect(ctx, req, terminateTLS)
if res != nil {
defer res.Body.Close()
}
Expand Down
29 changes: 29 additions & 0 deletions internal/martian/proxy_connect.go
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@
package martian

import (
"context"
"crypto/tls"
"errors"
"fmt"
Expand Down Expand Up @@ -50,6 +51,34 @@ var ErrConnectFallback = errors.New("martian: connect fallback")
// If the returned net.Conn is not nil, the response must be not nil.
type ConnectFunc func(req *http.Request) (*http.Response, io.ReadWriteCloser, error)

func (p *Proxy) Connect(ctx context.Context, req *http.Request, terminateTLS bool) (res *http.Response, crw io.ReadWriteCloser, cerr error) {
if p.ConnectFunc != nil {
res, crw, cerr = p.ConnectFunc(req)
}
if p.ConnectFunc == nil || errors.Is(cerr, ErrConnectFallback) {
var cconn net.Conn
res, cconn, cerr = p.connect(req)

if cconn != nil {
defer cconn.Close()
crw = cconn

if terminateTLS {
log.Debugf(ctx, "attempting to terminate TLS on CONNECT tunnel: %s", req.URL.Host)
tconn := tls.Client(cconn, p.clientTLSConfig())
if err := tconn.Handshake(); err == nil {
crw = tconn
} else {
log.Errorf(ctx, "failed to terminate TLS on CONNECT tunnel: %v", err)
cerr = err
}
}
}
}

return
}

func (p *Proxy) connect(req *http.Request) (*http.Response, net.Conn, error) {
ctx := req.Context()

Expand Down
33 changes: 1 addition & 32 deletions internal/martian/proxy_handler.go
Original file line number Diff line number Diff line change
Expand Up @@ -18,11 +18,8 @@ package martian

import (
"context"
"crypto/tls"
"errors"
"fmt"
"io"
"net"
"net/http"
"strings"

Expand Down Expand Up @@ -113,35 +110,7 @@ func (p proxyHandler) handleConnectRequest(rw http.ResponseWriter, req *http.Req
}

log.Debugf(ctx, "attempting to establish CONNECT tunnel: %s", req.URL.Host)
var (
res *http.Response
crw io.ReadWriteCloser
cerr error
)
if p.ConnectFunc != nil {
res, crw, cerr = p.ConnectFunc(req)
}
if p.ConnectFunc == nil || errors.Is(cerr, ErrConnectFallback) {
var cconn net.Conn
res, cconn, cerr = p.connect(req)

if cconn != nil {
defer cconn.Close()
crw = cconn

if terminateTLS {
log.Debugf(ctx, "attempting to terminate TLS on CONNECT tunnel: %s", req.URL.Host)
tconn := tls.Client(cconn, p.clientTLSConfig())
if err := tconn.Handshake(); err == nil {
crw = tconn
} else {
log.Errorf(ctx, "failed to terminate TLS on CONNECT tunnel: %v", err)
cerr = err
}
}
}
}

res, crw, cerr := p.Connect(ctx, req, terminateTLS)
if res != nil {
defer res.Body.Close()
}
Expand Down

0 comments on commit 2345fa3

Please sign in to comment.