feat(forwarder): split forwarder to different file

core/udp.go

@@ -0,0 +1,170 @@
+package core
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"io"
+	"log/slog"
+	"net"
+	"sync"
+	"time"
+
+	"tailscale.com/tsnet"
+)
+
+const udpForwardIdleTimeout = 2 * time.Minute
+
+func runUDPForwarder(ctx context.Context, srv *tsnet.Server, rule ForwardRule, logger *slog.Logger) {
+	ip := getSelfTsnetAddr(srv)
+	ln, err := srv.Listen("udp", fmt.Sprintf("%s:%d", ip.String(), rule.TailscalePort))
+	if err != nil {
+		logger.Error("failed to listen", "error", err)
+		return
+	}
+	logger.Debug("listening", slog.String("on", fmt.Sprintf("tailscale:%s:%d", ip.String(), rule.TailscalePort)))
+
+	go func() {
+		<-ctx.Done()
+		ln.Close()
+	}()
+
+	for {
+		conn, err := ln.Accept()
+		if err != nil {
+			if ctx.Err() != nil {
+				return
+			}
+			logger.Error("accept error", "error", err)
+			continue
+		}
+		go handleUDPForward(ctx, srv, conn, rule, logger)
+	}
+}
+
+func handleUDPForward(ctx context.Context, srv *tsnet.Server, conn net.Conn, rule ForwardRule, logger *slog.Logger) {
+	remoteAddrStr := conn.RemoteAddr().String()
+	clog := logger.With(slog.String("remote", remoteAddrStr))
+
+	lc, err := srv.LocalClient()
+	if err == nil {
+		who, err := lc.WhoIs(ctx, remoteAddrStr)
+		if err == nil {
+			clog = clog.With(slog.String("user", who.UserProfile.LoginName))
+		}
+	}
+
+	connType := getConnType(ctx, srv, remoteAddrStr)
+	clog.Info("accepted connection",
+		slog.String("conn_type", connType),
+		slog.String("local_addr", rule.LocalAddr),
+	)
+
+	localConn, err := dialUDP(ctx, rule.LocalAddr)
+	if err != nil {
+		clog.Error("failed to dial local", "error", err)
+		conn.Close()
+		return
+	}
+
+	stop := context.AfterFunc(ctx, func() {
+		conn.Close()
+		localConn.Close()
+	})
+	defer stop()
+
+	remoteIP, _, _ := net.SplitHostPort(remoteAddrStr)
+
+	var toTs, toLocal int64
+	var wg sync.WaitGroup
+	wg.Add(2)
+
+	go func() {
+		defer wg.Done()
+		buf := make([]byte, 65535)
+		for {
+			_ = conn.SetReadDeadline(time.Now().Add(udpForwardIdleTimeout))
+			n, err := conn.Read(buf)
+			if err != nil {
+				if netErr, ok := err.(net.Error); ok && netErr.Timeout() {
+					clog.Debug("udp forward idle timeout on ts side")
+				} else if !errors.Is(err, net.ErrClosed) && !errors.Is(err, io.EOF) {
+					clog.Debug("udp forward read from ts", "error", err)
+				}
+				localConn.Close()
+				return
+			}
+			_ = conn.SetReadDeadline(time.Time{})
+			toLocal += int64(n)
+			clog.Debug("inbound udp packet",
+				slog.String("from_ip", remoteIP),
+				slog.String("to_ip", rule.LocalAddr),
+				slog.Int("pkg_size", n),
+			)
+			if _, err := localConn.Write(buf[:n]); err != nil {
+				conn.Close()
+				return
+			}
+		}
+	}()
+
+	go func() {
+		defer wg.Done()
+		buf := make([]byte, 65535)
+		for {
+			_ = localConn.SetReadDeadline(time.Now().Add(udpForwardIdleTimeout))
+			n, err := localConn.Read(buf)
+			if err != nil {
+				if netErr, ok := err.(net.Error); ok && netErr.Timeout() {
+					clog.Debug("udp forward idle timeout on local side")
+				} else if !errors.Is(err, net.ErrClosed) && !errors.Is(err, io.EOF) {
+					clog.Debug("udp forward read from local", "error", err)
+				}
+				conn.Close()
+				return
+			}
+			_ = localConn.SetReadDeadline(time.Time{})
+			toTs += int64(n)
+			localIP, _, _ := net.SplitHostPort(localConn.RemoteAddr().String())
+			clog.Debug("outbound udp packet",
+				slog.String("from_ip", localIP),
+				slog.String("to_ip", remoteIP),
+				slog.Int("pkg_size", n),
+			)
+			if _, err := conn.Write(buf[:n]); err != nil {
+				localConn.Close()
+				return
+			}
+		}
+	}()
+
+	wg.Wait()
+	clog.Info("connection closed", slog.Int64("ts_rx_bytes", toLocal), slog.Int64("ts_tx_bytes", toTs))
+}
+
+func runUDPConnector(ctx context.Context, srv *tsnet.Server, rule ConnectRule, logger *slog.Logger) {
+	bindIP := rule.BindIP()
+	addr := fmt.Sprintf("%s:%d", bindIP, rule.LocalPort)
+	addrUDP, err := net.ResolveUDPAddr("udp", addr)
+	if err != nil {
+		logger.Error("failed to resolve local addr", "error", err)
+		return
+	}
+
+	pc, err := net.ListenUDP("udp", addrUDP)
+	if err != nil {
+		logger.Error("failed to listen locally", "error", err)
+		return
+	}
+	logger.Info("listening", slog.String("on", addr))
+
+	relay := &udpRelay{
+		listenConn: pc,
+		dialAddr:   rule.DstAddr,
+		logger:     logger,
+		direction:  "tailscale",
+		srv:        srv,
+		sessions:   make(map[string]*udpSession),
+	}
+	relay.run(ctx)
+}