agent/upgrade.go

Rama main · réplica de lectura del repositorio original.

package agent

import (
	"bufio"
	"net"
	"net/http"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/go-chi/chi/v5"
	"golang.org/x/xerrors"

	"cdr.dev/slog/v3"
	"github.com/coder/coder/v2/coderd/httpapi"
	"github.com/coder/coder/v2/coderd/tracing"
	"github.com/coder/coder/v2/coderd/util/syncmap"
	"github.com/coder/coder/v2/codersdk"
)

type upgradedConn struct {
	net.Conn
	clientSessionID string
	r               *bufio.Reader
}

func (c *upgradedConn) Read(p []byte) (int, error) { return c.r.Read(p) }
func (c *upgradedConn) ClientSessionID() string    { return c.clientSessionID }

type httpUpgradeListener struct {
	mu     sync.Mutex
	addr   net.Addr
	conn   chan *upgradedConn
	closed chan struct{}
}

func (ln *httpUpgradeListener) Accept() (net.Conn, error) {
	var c *upgradedConn
	select {
	case c = <-ln.conn:
	case <-ln.closed:
		return nil, xerrors.Errorf("tcp upgrade listener: %w", net.ErrClosed)
	}
	return c, nil
}

func (ln *httpUpgradeListener) Addr() net.Addr { return ln.addr }
func (ln *httpUpgradeListener) Close() error {
	ln.mu.Lock()
	defer ln.mu.Unlock()
	select {
	case <-ln.closed:
	default:
		close(ln.closed)
	}
	return nil
}

type httpUpgrader struct {
	listeners syncmap.Map[uint16, *httpUpgradeListener]
	logger    slog.Logger
	addr      net.Addr
}

// listen listens for upgrades on the specified port.
func (up *httpUpgrader) listen(port uint16) net.Listener {
	listener, _ := up.listeners.LoadOrStore(port, &httpUpgradeListener{
		addr:   up.addr,
		closed: make(chan struct{}),
		conn:   make(chan *upgradedConn),
	})
	return listener
}

// handler upgrades an HTTP connection based on the port number, passing along
// any client session ID that exists on the request context.
func (up *httpUpgrader) handler(rw http.ResponseWriter, r *http.Request) {
	ctx := r.Context()
	port, err := strconv.ParseUint(chi.URLParam(r, "port"), 10, 16)
	if err != nil {
		httpapi.Write(ctx, rw, http.StatusNotFound, codersdk.Response{
			Message: "Invalid port",
			Detail:  err.Error(),
		})
		return
	}

	listener, ok := up.listeners.Load(uint16(port))
	if !ok {
		httpapi.Write(ctx, rw, http.StatusNotFound, codersdk.Response{
			Message: "Unsupported port for upgrading",
		})
		return
	}

	if r.Header.Get("upgrade") == "" {
		httpapi.Write(ctx, rw, http.StatusBadRequest, codersdk.Response{
			Message: "Missing upgrade header",
		})
		return
	}

	if !strings.EqualFold(r.Header.Get("upgrade"), "tcp") {
		httpapi.Write(ctx, rw, http.StatusBadRequest, codersdk.Response{
			Message: "Invalid upgrade header",
		})
		return
	}

	clientSessionID := tracing.ClientSessionID(r)
	if clientSessionID == "" {
		up.logger.Warn(ctx, "client session ID is missing")
	}

	hijacker, ok := rw.(http.Hijacker)
	if !ok {
		httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{
			Message: "Hijacking not supported",
		})
		return
	}

	conn, brw, err := hijacker.Hijack()
	if err != nil {
		httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{
			Message: "Failed to upgrade connection",
			Detail:  err.Error(),
		})
		return
	}

	_, err = conn.Write([]byte(strings.Join([]string{
		"HTTP/1.1 101 Switching Protocols",
		"Connection: Upgrade",
		"Upgrade: tcp",
	}, "\r\n") + "\r\n\r\n"))
	if err != nil {
		_ = conn.Close()
		return
	}

	uconn := &upgradedConn{
		Conn:            conn,
		clientSessionID: clientSessionID,
		r:               brw.Reader,
	}

	t := time.NewTimer(30 * time.Second)
	defer t.Stop()
	select {
	case listener.conn <- uconn:
		return
	case <-listener.closed:
		up.logger.Info(ctx, "listener closed; closing connection")
	case <-t.C:
		up.logger.Info(ctx, "listener timed out accepting; closing connection")
	}
	_ = conn.Close()
}