Wormhole/internal/proxy/tcp_udp.go

package proxy

import (
	"log"
	"net"
	"sync"
	"time"

	"github.com/wormhole/wormhole/internal/protocol"
	"github.com/wormhole/wormhole/internal/store"
)

type tcpServer struct {
	manager *Manager
	tunnel  *store.Tunnel
	ln      net.Listener
	stop    chan struct{}
	wg      sync.WaitGroup
}

func newTCPServer(m *Manager, t *store.Tunnel) Service {
	return &tcpServer{
		manager: m,
		tunnel:  t,
		stop:    make(chan struct{}),
	}
}

func (s *tcpServer) ID() uint { return s.tunnel.ID }

func (s *tcpServer) Start() error {
	addr := net.JoinHostPort("0.0.0.0", itoa(s.tunnel.ListenPort))
	ln, err := net.Listen("tcp", addr)
	if err != nil {
		return err
	}
	s.ln = ln
	log.Printf("[proxy] tcp tunnel %d listening on %s -> %s", s.tunnel.ID, addr, s.tunnel.TargetAddr)
	go s.acceptLoop()
	return nil
}

func (s *tcpServer) acceptLoop() {
	for {
		conn, err := s.ln.Accept()
		if err != nil {
			select {
			case <-s.stop:
				return
			default:
				continue
			}
		}
		s.wg.Add(1)
		go func(c net.Conn) {
			defer s.wg.Done()
			s.manager.handleTCPConnection(s.tunnel, c)
		}(conn)
	}
}

func (s *tcpServer) Stop() error {
	close(s.stop)
	if s.ln != nil {
		_ = s.ln.Close()
	}
	s.wg.Wait()
	return nil
}

type udpSession struct {
	stream net.Conn
	last   time.Time
}

type udpServer struct {
	manager *Manager
	tunnel  *store.Tunnel
	conn    *net.UDPConn
	stop    chan struct{}
	mu      sync.Mutex
	sessions map[string]*udpSession
}

func newUDPServer(m *Manager, t *store.Tunnel) Service {
	return &udpServer{
		manager:  m,
		tunnel:   t,
		stop:     make(chan struct{}),
		sessions: make(map[string]*udpSession),
	}
}

func (s *udpServer) ID() uint { return s.tunnel.ID }

func (s *udpServer) Start() error {
	addr := net.JoinHostPort("0.0.0.0", itoa(s.tunnel.ListenPort))
	udpAddr, err := net.ResolveUDPAddr("udp", addr)
	if err != nil {
		return err
	}
	conn, err := net.ListenUDP("udp", udpAddr)
	if err != nil {
		return err
	}
	s.conn = conn
	log.Printf("[proxy] udp tunnel %d listening on %s -> %s", s.tunnel.ID, addr, s.tunnel.TargetAddr)
	go s.readLoop()
	go s.sweeper()
	return nil
}

func (s *udpServer) readLoop() {
	buf := make([]byte, 65535)
	for {
		n, remote, err := s.conn.ReadFromUDP(buf)
		if err != nil {
			select {
			case <-s.stop:
				return
			default:
				continue
			}
		}
		ip := remote.IP.String()
		s.manager.metrics.RecordRequestIP("tunnel", s.tunnel.ID, ip, ip)
		if !s.manager.security.Allow(ip, s.tunnel.ClientID, "tunnel", s.tunnel.ID) {
			continue
		}
		key := remote.String()
		s.mu.Lock()
		sess, ok := s.sessions[key]
		if !ok {
			sess = s.createSession(remote)
			if sess == nil {
				s.mu.Unlock()
				continue
			}
			s.sessions[key] = sess
			go s.runSession(key, sess, remote)
		}
		sess.last = time.Now()
		stream := sess.stream
		s.mu.Unlock()
		_, _ = stream.Write(buf[:n])
	}
}

func (s *udpServer) createSession(remote *net.UDPAddr) *udpSession {
	if !s.manager.bridge.IsOnline(s.tunnel.ClientID) {
		return nil
	}
	mode, _ := resolveProxyConfig(s.manager.store, s.tunnel.ClientID, s.tunnel.IPForwardMode, s.tunnel.CustomHeaders)
	meta := &protocol.LinkMeta{
		ConnType:   protocol.ConnTypeUDP,
		Target:     s.tunnel.TargetAddr,
		RemoteAddr: forwardIPForTCP(remote.String(), mode),
	}
	stream, err := s.manager.bridge.OpenStream(s.tunnel.ClientID, meta)
	if err != nil {
		return nil
	}
	return &udpSession{stream: stream, last: time.Now()}
}

func (s *udpServer) runSession(key string, sess *udpSession, remote *net.UDPAddr) {
	buf := make([]byte, 65535)
	for {
		n, err := sess.stream.Read(buf)
		if err != nil {
			break
		}
		_, _ = s.conn.WriteToUDP(buf[:n], remote)
		inlet := int64(n)
		s.manager.metrics.AddFlow("tunnel", s.tunnel.ID, s.tunnel.ClientID, inlet, 0)
	}
	s.mu.Lock()
	delete(s.sessions, key)
	s.mu.Unlock()
	_ = sess.stream.Close()
}

func (s *udpServer) sweeper() {
	ticker := time.NewTicker(time.Minute)
	defer ticker.Stop()
	for {
		select {
		case <-ticker.C:
			now := time.Now()
			s.mu.Lock()
			for k, sess := range s.sessions {
				if now.Sub(sess.last) > 120*time.Second {
					_ = sess.stream.Close()
					delete(s.sessions, k)
				}
			}
			s.mu.Unlock()
		case <-s.stop:
			return
		}
	}
}

func (s *udpServer) Stop() error {
	close(s.stop)
	if s.conn != nil {
		_ = s.conn.Close()
	}
	s.mu.Lock()
	for _, sess := range s.sessions {
		_ = sess.stream.Close()
	}
	s.sessions = make(map[string]*udpSession)
	s.mu.Unlock()
	return nil
}