Wormhole/internal/proxy/manager.go

package proxy

import (
	"bufio"
	"context"
	"crypto/tls"
	"fmt"
	"log"
	"net"
	"net/http"
	"net/http/httputil"
	"net/url"
	"strings"
	"sync"

	"github.com/wormhole/wormhole/internal/auth"
	"github.com/wormhole/wormhole/internal/bridge"
	"github.com/wormhole/wormhole/internal/cache"
	"github.com/wormhole/wormhole/internal/metrics"
	"github.com/wormhole/wormhole/internal/protocol"
	"github.com/wormhole/wormhole/internal/security"
	"github.com/wormhole/wormhole/internal/store"
)

type Manager struct {
	cache    *cache.Manager
	store    *store.Store
	bridge   *bridge.Bridge
	security *security.Service
	resAuth  *auth.ResourceAuth
	metrics  *metrics.Collector

	mu            sync.Mutex
	services      map[uint]Service
	httpProxySrv  *http.Server
	httpsProxySrv *http.Server
}

type Service interface {
	Start() error
	Stop() error
	ID() uint
}

func NewManager(c *cache.Manager, st *store.Store, b *bridge.Bridge, sec *security.Service, ra *auth.ResourceAuth, mc *metrics.Collector) *Manager {
	return &Manager{
		cache:    c,
		store:    st,
		bridge:   b,
		security: sec,
		resAuth:  ra,
		metrics:  mc,
		services: make(map[uint]Service),
	}
}

func (m *Manager) StartAll() error {
	tunnels := m.cache.ListEnabledTunnels()
	for _, t := range tunnels {
		if err := m.StartTunnel(t.ID); err != nil {
			log.Printf("[proxy] start tunnel %d: %v", t.ID, err)
		}
	}
	return nil
}

func (m *Manager) StartTunnel(id uint) error {
	t, ok := m.cache.GetTunnel(id)
	if !ok {
		return nil
	}
	m.mu.Lock()
	defer m.mu.Unlock()
	if old, exists := m.services[id]; exists {
		_ = old.Stop()
		delete(m.services, id)
	}
	var svc Service
	switch t.Mode {
	case "udp":
		svc = newUDPServer(m, t)
	default:
		if t.VisitorAuthEnabled {
			svc = newHTTPTunnelServer(m, t)
		} else {
			svc = newTCPServer(m, t)
		}
	}
	if err := svc.Start(); err != nil {
		return err
	}
	m.services[id] = svc
	t.RunStatus = true
	_ = m.store.UpdateTunnel(t)
	m.cache.SetTunnelRunStatus(id, true)
	return nil
}

func (m *Manager) StopTunnel(id uint) error {
	m.mu.Lock()
	defer m.mu.Unlock()
	if svc, ok := m.services[id]; ok {
		_ = svc.Stop()
		delete(m.services, id)
	}
	if t, ok := m.cache.GetTunnel(id); ok {
		t.RunStatus = false
		_ = m.store.UpdateTunnel(t)
		m.cache.SetTunnelRunStatus(id, false)
	}
	return nil
}

func (m *Manager) StartHTTPServer(httpPort, httpsPort int, certFile, keyFile string) error {
	if httpPort > 0 {
		srv := &http.Server{
			Addr:    formatPort(httpPort),
			Handler: m.httpHandler("http"),
		}
		m.httpProxySrv = srv
		go func() {
			log.Printf("[proxy] http host server on %s", srv.Addr)
			if err := srv.ListenAndServe(); err != nil && err != http.ErrServerClosed {
				log.Printf("[proxy] http server error: %v", err)
			}
		}()
	}
	if httpsPort > 0 {
		srv := &http.Server{
			Addr:    formatPort(httpsPort),
			Handler: m.httpHandler("https"),
		}
		m.httpsProxySrv = srv
		go func() {
			log.Printf("[proxy] https host server on %s", srv.Addr)
			if certFile != "" && keyFile != "" {
				if err := srv.ListenAndServeTLS(certFile, keyFile); err != nil && err != http.ErrServerClosed {
					log.Printf("[proxy] https server error: %v", err)
				}
			} else {
				log.Printf("[proxy] https port enabled but tls cert/key not configured, skipping TLS listener")
			}
		}()
	}
	return nil
}

func (m *Manager) Shutdown(ctx context.Context) error {
	m.mu.Lock()
	services := make([]Service, 0, len(m.services))
	for _, svc := range m.services {
		services = append(services, svc)
	}
	m.mu.Unlock()

	for _, svc := range services {
		_ = svc.Stop()
	}
	var err error
	if m.httpProxySrv != nil {
		if e := m.httpProxySrv.Shutdown(ctx); e != nil {
			err = e
		}
	}
	if m.httpsProxySrv != nil {
		if e := m.httpsProxySrv.Shutdown(ctx); e != nil {
			err = e
		}
	}
	return err
}

func formatPort(port int) string {
	return net.JoinHostPort("", itoa(port))
}

func itoa(n int) string {
	if n == 0 {
		return "0"
	}
	var b [12]byte
	i := len(b)
	for n > 0 {
		i--
		b[i] = byte('0' + n%10)
		n /= 10
	}
	return string(b[i:])
}

func (m *Manager) httpHandler(scheme string) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		directIP := auth.DirectRemoteIP(r)

		host := m.cache.MatchHost(r.Host, r.URL.Path, scheme)
		if host == nil {
			http.Error(w, "host not found", http.StatusNotFound)
			return
		}
		ip := auth.ExtractRemoteIP(r)
		m.metrics.RecordRequestIP("host", host.ID, ip, directIP)
		if !m.security.Allow(directIP, host.ClientID, "host", host.ID) {
			http.Error(w, "forbidden", http.StatusForbidden)
			return
		}
		if !m.resAuth.CheckHTTP(w, r, "host", host.ID) {
			return
		}
		if host.AutoHTTPS && scheme == "http" {
			target := "https://" + r.Host + r.URL.RequestURI()
			http.Redirect(w, r, target, http.StatusMovedPermanently)
			return
		}
		m.proxyHTTP(w, r, host, "host")
	})
}

func (m *Manager) proxyHTTP(w http.ResponseWriter, r *http.Request, host *store.Host, resourceType string) {
	targetURL, err := url.Parse("http://" + host.TargetAddr)
	if err != nil {
		http.Error(w, "bad target", http.StatusBadGateway)
		return
	}
	if !m.bridge.IsOnline(host.ClientID) {
		http.Error(w, "client offline", http.StatusBadGateway)
		return
	}
	mode, headers := resolveProxyConfig(m.store, host.ClientID, host.IPForwardMode, host.CustomHeaders)
	forwardIP := forwardIPForHTTP(r, mode)
	meta := &protocol.LinkMeta{
		ConnType:   protocol.ConnTypeHTTP,
		Target:     host.TargetAddr,
		RemoteAddr: forwardIP,
	}
	stream, err := m.bridge.OpenStream(host.ClientID, meta)
	if err != nil {
		http.Error(w, "tunnel error", http.StatusBadGateway)
		return
	}
	defer stream.Close()

	transport := &http.Transport{
		TLSClientConfig: &tls.Config{InsecureSkipVerify: true},
		DialContext: func(ctx context.Context, network, addr string) (net.Conn, error) {
			return stream, nil
		},
	}
	proxy := httputil.NewSingleHostReverseProxy(targetURL)
	proxy.Transport = transport
	originalDirector := proxy.Director
	proxy.Director = func(req *http.Request) {
		originalDirector(req)
		req.Host = targetURL.Host
		req.URL.Host = targetURL.Host
		req.URL.Scheme = targetURL.Scheme
		applyOutboundHeaders(req, r, mode, headers)
	}
	proxy.ErrorHandler = func(rw http.ResponseWriter, req *http.Request, e error) {
		http.Error(rw, "proxy error", http.StatusBadGateway)
	}
	if isWebSocketUpgrade(r) {
		proxy.ServeHTTP(w, r)
		return
	}
	rec := &responseRecorder{ResponseWriter: w, status: 200}
	proxy.ServeHTTP(rec, r)
	inlet := int64(rec.written)
	export := int64(r.ContentLength)
	if export < 0 {
		export = 0
	}
	m.metrics.AddFlow(resourceType, host.ID, host.ClientID, inlet, export)
}

func (m *Manager) proxyTunnelHTTP(w http.ResponseWriter, r *http.Request, t *store.Tunnel) {
	host := &store.Host{
		ID:            t.ID,
		ClientID:      t.ClientID,
		TargetAddr:    t.TargetAddr,
		IPForwardMode: t.IPForwardMode,
		CustomHeaders: t.CustomHeaders,
	}
	m.proxyHTTP(w, r, host, "tunnel")
}

type responseRecorder struct {
	http.ResponseWriter
	status  int
	written int
}

func (r *responseRecorder) Write(b []byte) (int, error) {
	n, err := r.ResponseWriter.Write(b)
	r.written += n
	return n, err
}

func (r *responseRecorder) Hijack() (net.Conn, *bufio.ReadWriter, error) {
	if h, ok := r.ResponseWriter.(http.Hijacker); ok {
		return h.Hijack()
	}
	return nil, nil, fmt.Errorf("hijack not supported")
}

func (r *responseRecorder) Flush() {
	if f, ok := r.ResponseWriter.(http.Flusher); ok {
		f.Flush()
	}
}

func isWebSocketUpgrade(r *http.Request) bool {
	return strings.EqualFold(r.Header.Get("Upgrade"), "websocket")
}

func (m *Manager) handleTCPConnection(t *store.Tunnel, clientConn net.Conn) {
	defer clientConn.Close()
	ip := auth.AddrHost(clientConn.RemoteAddr())
	m.metrics.RecordRequestIP("tunnel", t.ID, ip, ip)
	if !m.security.Allow(ip, t.ClientID, "tunnel", t.ID) {
		return
	}
	if !m.resAuth.CheckTCP(clientConn.RemoteAddr(), "tunnel", t.ID) {
		return
	}
	if !m.bridge.IsOnline(t.ClientID) {
		return
	}
	mode, _ := resolveProxyConfig(m.store, t.ClientID, t.IPForwardMode, t.CustomHeaders)
	meta := &protocol.LinkMeta{
		ConnType:   protocol.ConnTypeTCP,
		Target:     t.TargetAddr,
		RemoteAddr: forwardIPForTCP(clientConn.RemoteAddr().String(), mode),
	}
	stream, err := m.bridge.OpenStream(t.ClientID, meta)
	if err != nil {
		return
	}
	defer stream.Close()
	inlet, export := bridge.CopyBoth(clientConn, stream)
	m.metrics.AddFlow("tunnel", t.ID, t.ClientID, inlet, export)
}