Cleanup on the Noise protocol file

This commit is contained in:
Juan Font Alonso 2022-03-29 16:16:05 +02:00
parent d77cb3ba21
commit 38418e940f

View file

@ -2,8 +2,6 @@ package headscale
import ( import (
"encoding/base64" "encoding/base64"
"errors"
"fmt"
"net/http" "net/http"
"github.com/gin-gonic/gin" "github.com/gin-gonic/gin"
@ -12,13 +10,12 @@ import (
"golang.org/x/net/http2/h2c" "golang.org/x/net/http2/h2c"
"tailscale.com/control/controlbase" "tailscale.com/control/controlbase"
"tailscale.com/net/netutil" "tailscale.com/net/netutil"
"tailscale.com/types/key"
) )
const ( const (
errWrongConnectionUpgrade = Error("wrong connection upgrade") errWrongConnectionUpgrade = Error("wrong connection upgrade")
errCannotHijack = Error("cannot hijack connection") errCannotHijack = Error("cannot hijack connection")
errNetClosing = Error("net is closing") errNoiseHandshakeFailed = Error("noise handshake failed")
) )
const ( const (
@ -32,14 +29,6 @@ const (
handshakeHeaderName = "X-Tailscale-Handshake" handshakeHeaderName = "X-Tailscale-Handshake"
) )
type serverResult struct {
err error
clientAddr string
version int
peer key.MachinePublic
conn *controlbase.Conn
}
// NoiseUpgradeHandler is to upgrade the connection and hijack the net.Conn // NoiseUpgradeHandler is to upgrade the connection and hijack the net.Conn
// in order to use the Noise-based TS2021 protocol. Listens in /ts2021 // in order to use the Noise-based TS2021 protocol. Listens in /ts2021
func (h *Headscale) NoiseUpgradeHandler(ctx *gin.Context) { func (h *Headscale) NoiseUpgradeHandler(ctx *gin.Context) {
@ -72,29 +61,29 @@ func (h *Headscale) getNoiseConnection(ctx *gin.Context) (*controlbase.Conn, err
next := ctx.GetHeader("Upgrade") next := ctx.GetHeader("Upgrade")
if next == "" { if next == "" {
ctx.String(http.StatusBadRequest, "missing next protocol") ctx.String(http.StatusBadRequest, "missing next protocol")
return nil, errors.New("no next protocol in HTTP request") return nil, errWrongConnectionUpgrade
} }
if next != upgradeHeaderValue { if next != upgradeHeaderValue {
ctx.String(http.StatusBadRequest, "unknown next protocol") ctx.String(http.StatusBadRequest, "unknown next protocol")
return nil, fmt.Errorf("client requested unhandled next protocol %q", next) return nil, errWrongConnectionUpgrade
} }
initB64 := ctx.GetHeader(handshakeHeaderName) initB64 := ctx.GetHeader(handshakeHeaderName)
if initB64 == "" { if initB64 == "" {
ctx.String(http.StatusBadRequest, "missing Tailscale handshake header") ctx.String(http.StatusBadRequest, "missing Tailscale handshake header")
return nil, errors.New("no tailscale handshake header in HTTP request") return nil, errWrongConnectionUpgrade
} }
init, err := base64.StdEncoding.DecodeString(initB64) init, err := base64.StdEncoding.DecodeString(initB64)
if err != nil { if err != nil {
ctx.String(http.StatusBadRequest, "invalid tailscale handshake header") ctx.String(http.StatusBadRequest, "invalid tailscale handshake header")
return nil, fmt.Errorf("decoding base64 handshake header: %v", err) return nil, errWrongConnectionUpgrade
} }
hijacker, ok := ctx.Writer.(http.Hijacker) hijacker, ok := ctx.Writer.(http.Hijacker)
if !ok { if !ok {
log.Error().Caller().Err(err).Msgf("Hijack failed") log.Error().Caller().Err(err).Msgf("Hijack failed")
ctx.String(http.StatusInternalServerError, "HTTP does not support general TCP support") ctx.String(http.StatusInternalServerError, "HTTP does not support general TCP support")
return nil, errors.New("can't hijack client connection") return nil, errCannotHijack
} }
// This is what changes from the original AcceptHTTP() function. // This is what changes from the original AcceptHTTP() function.
@ -109,18 +98,18 @@ func (h *Headscale) getNoiseConnection(ctx *gin.Context) (*controlbase.Conn, err
log.Error().Caller().Err(err).Msgf("Hijack failed") log.Error().Caller().Err(err).Msgf("Hijack failed")
ctx.String(http.StatusInternalServerError, "HTTP does not support general TCP support") ctx.String(http.StatusInternalServerError, "HTTP does not support general TCP support")
return nil, errors.New("can't hijack client connection") return nil, errCannotHijack
} }
if err := conn.Flush(); err != nil { if err := conn.Flush(); err != nil {
netConn.Close() netConn.Close()
return nil, fmt.Errorf("flushing hijacked HTTP buffer: %w", err) return nil, errCannotHijack
} }
netConn = netutil.NewDrainBufConn(netConn, conn.Reader) netConn = netutil.NewDrainBufConn(netConn, conn.Reader)
nc, err := controlbase.Server(ctx.Request.Context(), netConn, *h.noisePrivateKey, init) nc, err := controlbase.Server(ctx.Request.Context(), netConn, *h.noisePrivateKey, init)
if err != nil { if err != nil {
netConn.Close() netConn.Close()
return nil, fmt.Errorf("noise handshake failed: %w", err) return nil, errNoiseHandshakeFailed
} }
return nc, nil return nc, nil