refactor: organize data forwarder interface

Author: cyunrei

cmd/main.go

@@ -85,13 +85,13 @@ func parseOptionsToRules() []rules.Rule {
 
 func startForwarding(r rules.Rule) error {
 	fc := forward.NewForwardingConfig().WithSourceAddr(r.SourceAddr).
-		WithDestinationAddr(r.DestinationAddr).WithProtocol(r.Protocol).
-		WithUDPDataForwarder(forward.NewSimpleUDPDataForwarder().
-			SetBufferSize(r.UDPBufferSize).SetDeadlineSecond(r.UDPTimeoutSecond))
-	if r.BandwidthLimit == forward.DefaultTCPBandwidthLimit {
-		fc.WithTCPDataForwarder(forward.NewSimpleTCPDataForwarder())
-	} else {
-		fc.WithTCPDataForwarder(forward.NewTrafficControlTCPDataForwarder().SetBandwidthLimit(r.BandwidthLimit))
+		WithDestinationAddr(r.DestinationAddr).WithProtocol(r.Protocol)
+	switch r.Protocol {
+	case "tcp":
+		fc.WithDataForwarder(forward.NewTCPDataForwarder().SetBandwidthLimit(r.BandwidthLimit))
+	case "udp":
+		fc.WithDataForwarder(forward.NewUDPDataForwarder().SetBandwidthLimit(r.BandwidthLimit).
+			SetDeadlineSecond(r.UDPTimeoutSecond).SetBufferSize(r.UDPBufferSize))
 	}
 	return fc.StartPortForwarding()
 }

pkg/forward/constants.go

@@ -0,0 +1,6 @@
+package forward
+
+const DefaultBandwidthLimit uint64 = 0
+
+const DefaultUDPBufferSize uint64 = 1024
+const DefaultUDPDeadlineSecond uint64 = 5

pkg/forward/forward.go

@@ -8,11 +8,10 @@ import (
 )
 
 type ForwardingConfig struct {
-	SourceAddr       string
-	DestinationAddr  string
-	Protocol         string
-	TCPDataForwarder TCPDataForwarder
-	UDPDataForwarder UDPDataForwarder
+	SourceAddr      string
+	DestinationAddr string
+	Protocol        string
+	DataForwarder   DataForwarder
 }
 
 func NewForwardingConfig() *ForwardingConfig {
@@ -34,13 +33,8 @@ func (f *ForwardingConfig) WithProtocol(protocol string) *ForwardingConfig {
 	return f
 }
 
-func (f *ForwardingConfig) WithTCPDataForwarder(tcpDataForwarder TCPDataForwarder) *ForwardingConfig {
-	f.TCPDataForwarder = tcpDataForwarder
-	return f
-}
-
-func (f *ForwardingConfig) WithUDPDataForwarder(udpDataForwarder UDPDataForwarder) *ForwardingConfig {
-	f.UDPDataForwarder = udpDataForwarder
+func (f *ForwardingConfig) WithDataForwarder(dataForwarder DataForwarder) *ForwardingConfig {
+	f.DataForwarder = dataForwarder
 	return f
 }
 
@@ -83,7 +77,7 @@ func (f *ForwardingConfig) startTCPPortForwarding() error {
 			continue
 		}
 		go func() {
-			f.TCPDataForwarder.Forward(localConn, remoteConn)
+			f.DataForwarder.Forward(localConn, remoteConn)
 			localConn.Close()
 			log.Printf("TCP connection disconnected from %s\n", localConn.RemoteAddr())
 			remoteConn.Close()
@@ -114,7 +108,7 @@ func (f *ForwardingConfig) startUDPPortForwarding() error {
 	}
 	defer remoteConn.Close()
 
-	f.UDPDataForwarder.Forward(*localConn, *remoteConn)
+	f.DataForwarder.Forward(&*localConn, &*remoteConn)
 
 	return nil
 }

pkg/forward/tcp.go

@@ -6,9 +6,23 @@ import (
 	"net"
 )
 
-type SimpleTCPDataForwarder struct{}
+type TCPDataForwarder struct {
+	BandwidthLimit uint64
+}
+
+func NewTCPDataForwarder() *TCPDataForwarder {
+	return &TCPDataForwarder{BandwidthLimit: DefaultBandwidthLimit}
+}
+
+func (f *TCPDataForwarder) Forward(sourceConn, destinationConn net.Conn) error {
+	if f.BandwidthLimit != DefaultBandwidthLimit {
+		return f.ForwardWithTrafficControl(sourceConn, destinationConn)
+	} else {
+		return f.ForwardWithNormal(sourceConn, destinationConn)
+	}
+}
 
-func (f *SimpleTCPDataForwarder) Forward(sourceConn, destinationConn net.Conn) error {
+func (f *TCPDataForwarder) ForwardWithNormal(sourceConn, destinationConn net.Conn) error {
 	done := make(chan *ForwardingError, 2)
 
 	go func() {
@@ -31,26 +45,7 @@ func (f *SimpleTCPDataForwarder) Forward(sourceConn, destinationConn net.Conn) e
 	return nil
 }
 
-func NewSimpleTCPDataForwarder() *SimpleTCPDataForwarder {
-	return &SimpleTCPDataForwarder{}
-}
-
-func NewTrafficControlTCPDataForwarder() *TrafficControlTCPDataForwarder {
-	return &TrafficControlTCPDataForwarder{BandwidthLimit: DefaultTCPBandwidthLimit}
-}
-
-const DefaultTCPBandwidthLimit uint64 = 0
-
-type TrafficControlTCPDataForwarder struct {
-	BandwidthLimit uint64
-}
-
-func (f *TrafficControlTCPDataForwarder) SetBandwidthLimit(bandwidthLimit uint64) *TrafficControlTCPDataForwarder {
-	f.BandwidthLimit = bandwidthLimit
-	return f
-}
-
-func (f *TrafficControlTCPDataForwarder) Forward(sourceConn, destinationConn net.Conn) error {
+func (f *TCPDataForwarder) ForwardWithTrafficControl(sourceConn, destinationConn net.Conn) error {
 	done := make(chan *ForwardingError, 2)
 
 	go func() {
@@ -76,3 +71,8 @@ func (f *TrafficControlTCPDataForwarder) Forward(sourceConn, destinationConn net
 
 	return nil
 }
+
+func (f *TCPDataForwarder) SetBandwidthLimit(bandwidthLimit uint64) *TCPDataForwarder {
+	f.BandwidthLimit = bandwidthLimit
+	return f
+}

pkg/forward/types.go

@@ -2,10 +2,8 @@ package forward
 
 import "net"
 
-type TCPDataForwarder interface {
+type DataForwarder interface {
 	Forward(sourceConn, destinationConn net.Conn) error
-}
-
-type UDPDataForwarder interface {
-	Forward(sourceConn, destinationConn net.UDPConn)
+	ForwardWithNormal(sourceConn, destinationConn net.Conn) error
+	ForwardWithTrafficControl(sourceConn, destinationConn net.Conn) error
 }

pkg/forward/udp.go

@@ -6,36 +6,31 @@ import (
 	"time"
 )
 
-const DefaultUDPBufferSize uint64 = 1024
-const DefaultUDPDeadlineSecond uint64 = 5
-
-type SimpleUDPDataForwarder struct {
+type UDPDataForwarder struct {
 	BufferSize     uint64
+	BandwidthLimit uint64
 	DeadlineSecond time.Duration
 }
 
-func NewSimpleUDPDataForwarder() *SimpleUDPDataForwarder {
-	return &SimpleUDPDataForwarder{
+func NewUDPDataForwarder() *UDPDataForwarder {
+	return &UDPDataForwarder{
+		BandwidthLimit: DefaultBandwidthLimit,
 		BufferSize:     DefaultUDPBufferSize,
 		DeadlineSecond: time.Duration(DefaultUDPDeadlineSecond),
 	}
 }
 
-func (f *SimpleUDPDataForwarder) SetBufferSize(size uint64) *SimpleUDPDataForwarder {
-	f.BufferSize = size
-	return f
-}
-
-func (f *SimpleUDPDataForwarder) SetDeadlineSecond(second uint64) *SimpleUDPDataForwarder {
-	f.DeadlineSecond = time.Duration(second)
-	return f
+func (f *UDPDataForwarder) Forward(sourceConn, destinationConn net.Conn) error {
+	return f.ForwardWithNormal(sourceConn, destinationConn)
 }
 
-func (f *SimpleUDPDataForwarder) Forward(sourceConn, destinationConn net.UDPConn) {
+func (f *UDPDataForwarder) ForwardWithNormal(sourceConn, destinationConn net.Conn) error {
+	sourceUDPConn, _ := sourceConn.(*net.UDPConn)
+	destinationUDPConn, _ := destinationConn.(*net.UDPConn)
 	sourceConnBuffer := make([]byte, f.BufferSize)
 	for {
 		sourceConn.SetReadDeadline(time.Now().Add(f.DeadlineSecond * time.Second))
-		n, sourceConnAddr, err := sourceConn.ReadFromUDP(sourceConnBuffer)
+		n, sourceConnAddr, err := sourceUDPConn.ReadFromUDP(sourceConnBuffer)
 		if err != nil {
 			continue
 		}
@@ -51,17 +46,36 @@ func (f *SimpleUDPDataForwarder) Forward(sourceConn, destinationConn net.UDPConn
 
 			destinationConnBuffer := make([]byte, f.BufferSize)
 			destinationConn.SetReadDeadline(time.Now().Add(f.DeadlineSecond * time.Second))
-			m, _, err := destinationConn.ReadFromUDP(destinationConnBuffer)
+			m, _, err := destinationUDPConn.ReadFromUDP(destinationConnBuffer)
 			var netErr net.Error
 			if errors.As(err, &netErr) && netErr.Timeout() {
 				return
 			}
 
-			_, err = sourceConn.WriteToUDP(destinationConnBuffer[:m], sourceConnAddr)
+			_, err = sourceUDPConn.WriteToUDP(destinationConnBuffer[:m], sourceConnAddr)
 			if err != nil {
 				return
 			}
 
 		}(data, sourceConnAddr)
 	}
 }
+
+func (f *UDPDataForwarder) ForwardWithTrafficControl(sourceConn, destinationConn net.Conn) error {
+	return f.ForwardWithNormal(sourceConn, destinationConn)
+}
+
+func (f *UDPDataForwarder) SetBandwidthLimit(bandwidthLimit uint64) *UDPDataForwarder {
+	f.BandwidthLimit = bandwidthLimit
+	return f
+}
+
+func (f *UDPDataForwarder) SetBufferSize(size uint64) *UDPDataForwarder {
+	f.BufferSize = size
+	return f
+}
+
+func (f *UDPDataForwarder) SetDeadlineSecond(second uint64) *UDPDataForwarder {
+	f.DeadlineSecond = time.Duration(second)
+	return f
+}