warmbly/internal/api/handler/socket.go

package handler

import (
	"context"
	"net"
	"net/http"
	"net/http/httputil"
	"net/url"
	"os"
	"strings"
	"time"

	"github.com/gin-gonic/gin"
	"github.com/google/uuid"
	"github.com/warmbly/warmbly/internal/api/middleware"
	"github.com/warmbly/warmbly/internal/app/socket"
	"github.com/warmbly/warmbly/internal/app/token"
	"github.com/warmbly/warmbly/internal/config"
	"github.com/warmbly/warmbly/internal/errx"
)

func (h *Handler) GenerateWebsocket(c *gin.Context) {
	userID := middleware.GetUserID(c)
	uid, err := uuid.Parse(userID)
	if err != nil {
		errx.Handle(c, errx.ErrUser)
		return
	}

	token, xerr := h.SocketService.GenerateWebsocketToken(c.Request.Context(), uid)
	if xerr != nil {
		errx.Handle(c, xerr)
		return
	}

	wsURL := token
	endpoint := h.websocketEndpoint()
	if endpoint != "" {
		u, err := url.Parse(endpoint)
		if err != nil {
			errx.Handle(c, errx.InternalError())
			return
		}

		q := u.Query()
		q.Set("token", token)
		u.RawQuery = q.Encode()
		wsURL = u.String()
	}

	response := gin.H{
		"url":        wsURL,
		"expires_in": socket.SocketTTL.Seconds(),
	}
	if endpoint != "" {
		response["proxy_path"] = "/realtime/socket/websocket?token=" + url.QueryEscape(token)
	}
	c.JSON(http.StatusOK, response)
}

func (h *Handler) websocketEndpoint() string {
	if h.WebsocketURI != "" {
		return config.NormalizeWebsocketURL(h.WebsocketURI)
	}
	return config.WebsocketURL()
}

// Browser connections use the API origin already permitted by the dashboard CSP.
func (h *Handler) ProxyWebsocket(c *gin.Context) {
	endpoint := h.websocketEndpoint()
	if endpoint == "" {
		c.Status(http.StatusServiceUnavailable)
		return
	}
	if !strings.EqualFold(c.GetHeader("Upgrade"), "websocket") {
		c.Status(http.StatusBadRequest)
		return
	}
	claims, xerr := h.TokenService.VerifyTokenFor(token.PurposeWebSocket, c.Query("token"))
	if xerr != nil {
		c.Status(http.StatusUnauthorized)
		return
	}
	proof, err := h.TokenService.GenerateWebsocketProxyProof(claims, c.ClientIP())
	if err != nil {
		c.Status(http.StatusInternalServerError)
		return
	}
	target, err := url.Parse(endpoint)
	if err != nil {
		c.Status(http.StatusServiceUnavailable)
		return
	}
	target.Scheme = strings.Replace(target.Scheme, "ws", "http", 1)
	proxy := &httputil.ReverseProxy{
		Transport: websocketProxyTransport(target),
		Rewrite: func(request *httputil.ProxyRequest) {
			request.SetURL(target)
			request.Out.URL.Path = target.Path
			request.Out.URL.RawPath = target.RawPath
			query := target.Query()
			query.Set("token", c.Query("token"))
			query.Set("vsn", c.DefaultQuery("vsn", "1.0.0"))
			request.Out.URL.RawQuery = query.Encode()
			request.Out.Header.Del("Authorization")
			request.Out.Header.Del("Cookie")
			request.Out.Header.Del("X-Warmbly-Token")
			request.Out.Header.Set("X-Warmbly-Proxy-Proof", proof)
		},
		ErrorHandler: func(w http.ResponseWriter, _ *http.Request, _ error) {
			w.WriteHeader(http.StatusBadGateway)
		},
	}
	proxy.ServeHTTP(c.Writer, c.Request)
}

func websocketProxyTransport(target *url.URL) http.RoundTripper {
	transport := http.DefaultTransport.(*http.Transport).Clone()
	transport.ResponseHeaderTimeout = 10 * time.Second
	// Old compose files advertise localhost to browsers, not to the API container.
	host := target.Hostname()
	ip := net.ParseIP(host)
	if host == "localhost" || (ip != nil && ip.IsLoopback()) {
		if _, err := os.Stat("/.dockerenv"); err == nil {
			dialer := &net.Dialer{Timeout: 5 * time.Second, KeepAlive: 30 * time.Second}
			transport.DialContext = func(ctx context.Context, network, address string) (net.Conn, error) {
				connection, err := dialer.DialContext(ctx, network, address)
				if err == nil {
					return connection, nil
				}
				port := target.Port()
				if port == "" {
					port = "80"
					if target.Scheme == "https" {
						port = "443"
					}
				}
				return dialer.DialContext(ctx, network, net.JoinHostPort("realtime", port))
			}
		}
	}
	return transport
}