pkg/revdial/revdial.go
Rama master · réplica de lectura del repositorio original.
// Copyright 2019 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.
// Package revdial implements a Dialer and Listener which work together
// to turn an accepted connection (for instance, a Hijacked HTTP request) into
// a Dialer which can then create net.Conns connecting back to the original
// dialer, which then gets a net.Listener accepting those conns.
//
// This is basically a very minimal SOCKS5 client & server.
//
// The motivation is that sometimes you want to run a server on a
// machine deep inside a NAT. Rather than connecting to the machine
// directly (which you can't, because of the NAT), you have the
// sequestered machine connect out to a public machine. Both sides
// then use revdial and the public machine can become a client for the
// NATed machine.
package revdial
import (
"bufio"
"context"
"crypto/rand"
"encoding/hex"
"encoding/json"
"errors"
"fmt"
"net"
"net/http"
"net/url"
"strings"
"sync"
"time"
"github.com/gorilla/websocket"
"github.com/shellhub-io/shellhub/pkg/clock"
"github.com/shellhub-io/shellhub/pkg/uuid"
"github.com/shellhub-io/shellhub/pkg/wsconnadapter"
log "github.com/sirupsen/logrus"
)
// ErrDialerClosed is returned once a dialer can no longer serve connections. It is terminal for the
// dialer, and the caller reconnects.
var ErrDialerClosed = errors.New("revdial.Dialer closed")
const (
dialerUniqParam = "revdial.dialer"
dialUUIDParam = "uuid"
)
const dialerKeepAliveTimeout = 35 * time.Second
// The Dialer can create new connections.
type Dialer struct {
conn net.Conn // hijacked client conn
path string // e.g. "/revdial"
uniqID string
pickupPath string // path + uniqID: "/revdial?revdial.dialer="+uniqID
pendingMu sync.Mutex
pending map[string]*pendingDial
connReady chan string
donec chan struct{}
closeOnce sync.Once
logger *log.Entry
clock clock.Clock
}
type pendingDial struct {
conn chan net.Conn
failed chan error
}
var dialers = sync.Map{}
// NewDialer returns the side of the connection which will initiate
// new connections. This will typically be the side which did the HTTP
// Hijack. The connection is (typically) the hijacked HTTP client
// connection. The connPath is the HTTP path and optional query (but
// without scheme or host) on the dialer where the ConnHandler is
// mounted. The dialer reads clk to set its write deadlines.
func NewDialer(logger *log.Entry, c net.Conn, connPath string, clk clock.Clock) *Dialer {
d := &Dialer{
path: connPath,
uniqID: newUniqID(),
conn: c,
donec: make(chan struct{}),
connReady: make(chan string, 8),
pending: make(map[string]*pendingDial),
logger: logger,
clock: clk,
}
join := "?"
if strings.Contains(connPath, "?") {
join = "&"
}
d.pickupPath = connPath + join + dialerUniqParam + "=" + d.uniqID
d.register()
go d.serve() //nolint:errcheck
d.logger.Debug("new dialer connection")
return d
}
func newUniqID() string {
buf := make([]byte, 16)
rand.Read(buf)
return hex.EncodeToString(buf)
}
func (d *Dialer) register() {
dialers.Store(d.uniqID, d)
}
func (d *Dialer) unregister() {
dialers.Delete(d.uniqID)
}
// Done returns a channel which is closed when d is closed (either by
// this process on purpose, by a local error, or close or error from
// the peer).
func (d *Dialer) Done() <-chan struct{} { return d.donec }
// Close closes the Dialer. It never blocks, returns nil, and every later call is a no-op.
func (d *Dialer) Close() error {
d.closeOnce.Do(d.close)
return nil
}
func (d *Dialer) close() {
d.logger.Debug("dialer connection closed")
d.unregister()
_ = d.conn.Close()
close(d.donec)
}
// Dial creates a new connection back to the Listener. It returns the connection the device
// dialled back for this call, the error the device reported when it could not, ErrDialerClosed
// once the dialer is closed, or ctx's error when ctx ends first. A connection that arrives
// after Dial returned is closed.
func (d *Dialer) Dial(ctx context.Context) (net.Conn, error) {
id := uuid.Generate()
pending := &pendingDial{
conn: make(chan net.Conn, 1),
failed: make(chan error, 1),
}
d.addPending(id, pending)
defer d.removePending(id, pending)
select {
case d.connReady <- id:
d.logger.Debug("message true to conn ready channel")
case <-d.donec:
d.logger.Debug("dial done")
return nil, ErrDialerClosed
case <-ctx.Done():
d.logger.Debug("dial done due context cancellation")
return nil, ctx.Err()
}
select {
case c := <-pending.conn:
d.logger.Debug("new incoming connection")
return c, nil
case err := <-pending.failed:
d.logger.Debug("failed to pick-up connection")
return nil, err
case <-d.donec:
d.logger.Debug("dial done on pick-up")
return nil, ErrDialerClosed
case <-ctx.Done():
d.logger.Debug("dial done on pick-up due context cancellation")
return nil, ctx.Err()
}
}
func (d *Dialer) addPending(id string, pending *pendingDial) {
d.pendingMu.Lock()
defer d.pendingMu.Unlock()
d.pending[id] = pending
}
func (d *Dialer) removePending(id string, pending *pendingDial) {
d.pendingMu.Lock()
delete(d.pending, id)
d.pendingMu.Unlock()
pending.abandon()
}
func (d *Dialer) matchConn(id string, c net.Conn) {
d.pendingMu.Lock()
defer d.pendingMu.Unlock()
pending, ok := d.pending[id]
if !ok {
d.logger.WithField("uuid", id).Debug("closing a connection no dial is waiting for")
_ = c.Close()
return
}
if !pending.deliver(c) {
d.logger.WithField("uuid", id).Debug("closing a second connection for the same dial")
_ = c.Close()
}
}
func (d *Dialer) failPickup(connPath string, err error) {
parsed, parseErr := url.ParseRequestURI(connPath)
if parseErr != nil {
d.logger.WithError(parseErr).WithField("connPath", connPath).Debug("pick-up failed for a path that names no dial")
return
}
id := parsed.Query().Get(dialUUIDParam)
d.pendingMu.Lock()
defer d.pendingMu.Unlock()
pending, ok := d.pending[id]
if !ok {
d.logger.WithError(err).WithField("uuid", id).Debug("pick-up failed for a dial that is no longer waiting")
return
}
pending.fail(err)
}
func (p *pendingDial) deliver(c net.Conn) bool {
select {
case p.conn <- c:
return true
default:
return false
}
}
func (p *pendingDial) fail(err error) {
select {
case p.failed <- err:
default:
}
}
func (p *pendingDial) abandon() {
select {
case c := <-p.conn:
_ = c.Close()
default:
}
}
func (d *Dialer) serve() error {
defer d.Close() //nolint:errcheck
go func() {
defer d.Close() //nolint:errcheck
defer d.logger.Debug("dialer serve done")
br := bufio.NewReader(d.conn)
for {
line, err := br.ReadSlice('\n')
if err != nil {
d.logger.WithError(err).Trace("failed to read the agent's command")
unexpectedError := websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseAbnormalClosure)
if !errors.Is(err, net.ErrClosed) && unexpectedError {
d.logger.WithError(err).Error("revdial.Dialer failed to read")
}
return
}
var msg controlMsg
if err := json.Unmarshal(line, &msg); err != nil {
d.logger.WithError(err).WithField("line", line).Printf("revdial.Dialer read invalid JSON")
return
}
switch msg.Command {
case "pickup-failed":
d.failPickup(msg.ConnPath, fmt.Errorf("revdial listener failed to pick up connection: %v", msg.Err))
case "keep-alive":
default:
log.WithField("message", msg.Command).Debug("unknown message received")
}
}
}()
for {
if err := d.sendMessage(controlMsg{Command: "keep-alive"}); err != nil {
d.logger.WithError(err).Debug("failed to send keep-alive message to device")
return err
}
t := time.NewTimer(30 * time.Second)
select {
case <-t.C:
continue
case uuid := <-d.connReady:
t.Stop()
if err := d.sendMessage(controlMsg{
Command: "conn-ready",
ConnPath: d.pickupPath + "&" + dialUUIDParam + "=" + uuid,
}); err != nil {
d.logger.WithError(err).Debug("failed to send conn-ready message to device")
return err
}
case <-d.donec:
t.Stop()
return ErrDialerClosed
}
}
}
func (d *Dialer) sendMessage(m controlMsg) error {
if err := d.conn.SetWriteDeadline(d.clock.Now().Add(10 * time.Second)); err != nil {
d.logger.WithError(err).Debug("failed to set the write dead line to device")
return err
}
j, err := json.Marshal(m)
if err != nil {
d.logger.WithError(err).Error("failed to encode the control message")
return err
}
j = append(j, '\n')
if _, err := d.conn.Write(j); err != nil {
d.logger.WithError(err).Debug("failed to write on the connection")
return err
}
return d.conn.SetWriteDeadline(time.Time{})
}
// NewListener returns a new Listener, accepting connections which
// arrive from the provided server connection, which should be after
// any necessary authentication (usually after an HTTP exchange).
//
// The provided dialServer func is responsible for connecting back to
// the server and doing TLS setup.
func NewListener(serverConn net.Conn, dialServer func(context.Context, string) (*websocket.Conn, *http.Response, error)) *Listener {
ln := &Listener{
sc: serverConn,
dial: dialServer,
connc: make(chan net.Conn, acceptQueueSize),
donec: make(chan struct{}),
}
go ln.run()
return ln
}
var _ net.Listener = (*Listener)(nil)
// Listener is a net.Listener, returning new connections which arrive
// from a corresponding Dialer.
type Listener struct {
sc net.Conn
connc chan net.Conn
donec chan struct{}
dial func(context.Context, string) (*websocket.Conn, *http.Response, error)
writec chan<- []byte
mu sync.Mutex
closed bool
}
type controlMsg struct {
Command string `json:"command,omitempty"` // "keep-alive", "conn-ready", "pickup-failed"
ConnPath string `json:"connPath,omitempty"` // conn pick-up URL path for "conn-url", "pickup-failed"
Err string `json:"err,omitempty"`
}
func (ln *Listener) run() {
done := func() {
_ = ln.Close()
}
var onceDefer sync.Once
defer onceDefer.Do(done)
closeTimer := time.AfterFunc(dialerKeepAliveTimeout, done)
writec := make(chan []byte, controlWriteQueueSize)
ln.writec = writec
go func() {
defer onceDefer.Do(done)
for {
select {
case <-ln.donec:
return
case msg := <-writec:
if _, err := ln.sc.Write(msg); err != nil {
log.Printf("revdial.Listener: error writing message to server: %v", err)
return
}
}
}
}()
go func() {
defer onceDefer.Do(done)
br := bufio.NewReader(ln.sc)
for {
line, err := br.ReadSlice('\n')
if err != nil {
return
}
var msg controlMsg
if err := json.Unmarshal(line, &msg); err != nil {
log.Printf("revdial.Listener read invalid JSON: %q: %v", line, err)
return
}
switch msg.Command {
case "keep-alive":
closeTimer.Reset(dialerKeepAliveTimeout)
case "conn-ready":
go ln.grabConn(msg.ConnPath)
default:
}
}
}()
for {
ln.sendMessage(controlMsg{Command: "keep-alive"})
t := time.NewTimer(30 * time.Second)
select {
case <-t.C:
continue
case <-ln.donec:
t.Stop()
return
}
}
}
func (ln *Listener) sendMessage(m controlMsg) {
j, err := json.Marshal(m)
if err != nil {
log.WithError(err).Error("failed to encode the control message")
return
}
j = append(j, '\n')
select {
case ln.writec <- j:
case <-ln.donec:
}
}
func (ln *Listener) grabConn(path string) {
ctx, cancel := context.WithTimeout(context.Background(), 20*time.Second)
defer cancel()
wsConn, resp, err := ln.dial(ctx, path) //nolint:bodyclose // gorilla/websocket documents that the handshake response body need not be closed.
if err != nil {
ln.sendMessage(controlMsg{Command: "pickup-failed", ConnPath: path, Err: err.Error()})
return
}
failPickup := func(err error) {
_ = wsConn.Close()
log.Printf("revdial.Listener: failed to pick up connection to %s: %v", path, err)
ln.sendMessage(controlMsg{Command: "pickup-failed", ConnPath: path, Err: err.Error()})
}
if resp.StatusCode != http.StatusSwitchingProtocols {
failPickup(fmt.Errorf("non-101 response %v", resp.Status))
return
}
err = ln.enqueue(wsconnadapter.New(wsConn))
switch {
case errors.Is(err, ErrListenerClosed):
_ = wsConn.Close()
case err != nil:
failPickup(err)
}
}
const (
acceptQueueSize = 8
controlWriteQueueSize = 8
)
var errTooManyPickups = errors.New("too many picked-up connections are waiting to be accepted")
func (ln *Listener) enqueue(c net.Conn) error {
ln.mu.Lock()
defer ln.mu.Unlock()
if ln.closed {
return ErrListenerClosed
}
select {
case ln.connc <- c:
return nil
default:
return errTooManyPickups
}
}
// Accept blocks and returns a new connection, or ErrListenerClosed once the Listener is closed:
// by Close, or because its control connection to the server failed, carried a malformed message,
// or stopped delivering keep-alives.
func (ln *Listener) Accept() (net.Conn, error) {
select {
case c := <-ln.connc:
return c, nil
case <-ln.donec:
return nil, ErrListenerClosed
}
}
// ErrListenerClosed is returned by Accept once the Listener is closed. Accept lists what closes it.
var ErrListenerClosed = errors.New("revdial: Listener closed")
// Close closes the Listener, making future Accept calls return
// ErrListenerClosed, and closes every picked-up connection that no Accept has
// taken yet. A connection an Accept running concurrently takes belongs to that
// caller. It never blocks and always returns nil.
func (ln *Listener) Close() error {
ln.mu.Lock()
defer ln.mu.Unlock()
if ln.closed {
return nil
}
go func() { _ = ln.sc.Close() }()
ln.closed = true
close(ln.donec)
for {
select {
case c := <-ln.connc:
_ = c.Close()
default:
return nil
}
}
}
// Addr returns a dummy address. This exists only to conform to the
// net.Listener interface.
func (ln *Listener) Addr() net.Addr { return fakeAddr{} }
type fakeAddr struct{}
func (fakeAddr) Network() string { return "revdial" }
func (fakeAddr) String() string { return "revdialconn" }
// ConnHandler returns the HTTP handler that needs to be mounted somewhere
// that the Listeners can dial out and get to. A dialer to connect to it
// is given to NewListener and the path to reach it is given to NewDialer
// to use in messages to the listener. It answers an unknown dialer with 400,
// and when the upgrade fails it logs the cause through the dialer's logger and
// writes nothing, because the upgrader has already written the response.
func ConnHandler(upgrader websocket.Upgrader) http.Handler {
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
dialerUniq := r.FormValue(dialerUniqParam)
uuid := r.FormValue(dialUUIDParam)
d, _ := dialers.Load(dialerUniq)
dialer, ok := d.(*Dialer)
if !ok {
http.Error(w, "unknown dialer", http.StatusBadRequest)
return
}
wsConn, err := upgrader.Upgrade(w, r, nil)
if err != nil {
dialer.logger.WithError(err).WithField("uuid", uuid).Warn("failed to upgrade the pickup connection")
return
}
dialer.matchConn(uuid, wsconnadapter.New(wsConn))
})
}