refactor: optimize code (#557)

pkg/core/gvisortunendpoint.go

@@ -25,11 +25,8 @@ func (h *gvisorTCPHandler) readFromEndpointWriteToTCPConn(ctx context.Context, c
 		pktBuffer := endpoint.ReadContext(ctx)
 		if pktBuffer != nil {
 			sniffer.LogPacket("[gVISOR] ", sniffer.DirectionSend, pktBuffer.NetworkProtocolNumber, pktBuffer)
-			data := pktBuffer.ToView().AsSlice()
-			buf := config.LPool.Get().([]byte)[:]
-			n := copy(buf, data)
-			_, err := tcpConn.Write(buf[:n+2])
-			config.LPool.Put(buf)
+			buf := pktBuffer.ToView().AsSlice()
+			_, err := tcpConn.Write(buf)
 			if err != nil {
 				plog.G(ctx).Errorf("[TUN-GVISOR] Failed to write data to tun device: %v", err)
 			}
@@ -123,12 +120,7 @@ func (h *gvisorTCPHandler) handlePacket(ctx context.Context, buf []byte, length
 		}
 	} else if config.RouterIP.Equal(dst) || config.RouterIP6.Equal(dst) {
 		plog.G(ctx).Debugf("[TCP-GVISOR] Forward to TUN device, SRC: %s, DST: %s, Protocol: %s, Length: %d", src, dst, protocol, length)
-		util.SafeWrite(h.packetChan, &Packet{
-			length: length,
-			data:   buf[:],
-			src:    src,
-			dst:    dst,
-		}, func(v *Packet) {
+		util.SafeWrite(h.packetChan, NewPacket(buf[:], length, src, dst), func(v *Packet) {
 			config.LPool.Put(v.data[:])
 			plog.G(context.Background()).Errorf("[TCP-GVISOR] Drop packet, SRC: %s, DST: %s, Protocol: %s, Length: %d", src, dst, protocol, v.length)
 		})

pkg/core/gvisorudphandler.go

@@ -68,7 +68,11 @@ func (c *gvisorUDPConnOverTCP) Read(b []byte) (int, error) {
 }
 
 func (c *gvisorUDPConnOverTCP) Write(b []byte) (int, error) {
-	packet := newDatagramPacket(b, len(b)-2)
+	buf := config.LPool.Get().([]byte)[:]
+	n := copy(buf, b)
+	defer config.LPool.Put(buf)
+
+	packet := newDatagramPacket(buf, n)
 	if err := packet.Write(c.Conn); err != nil {
 		return 0, err
 	}

pkg/core/tcphandler.go

@@ -98,12 +98,7 @@ func (h *UDPOverTCPHandler) handlePacket(ctx context.Context, tcpConn net.Conn,
 		config.LPool.Put(datagram.Data[:])
 	} else {
 		plog.G(ctx).Debugf("[TCP] Forward to TUN device, SRC: %s, DST: %s, Protocol: %s, Length: %d", src, dst, layers.IPProtocol(protocol).String(), datagram.DataLength)
-		util.SafeWrite(h.packetChan, &Packet{
-			data:   datagram.Data,
-			length: int(datagram.DataLength),
-			src:    src,
-			dst:    dst,
-		}, func(v *Packet) {
+		util.SafeWrite(h.packetChan, NewPacket(datagram.Data, int(datagram.DataLength), src, dst), func(v *Packet) {
 			plog.G(context.Background()).Errorf("Stuck packet, SRC: %s, DST: %s, Protocol: %s, Length: %d", src, dst, layers.IPProtocol(protocol).String(), v.length)
 			h.packetChan <- v
 		})
@@ -160,7 +155,11 @@ func (c *UDPConnOverTCP) ReadFrom(b []byte) (int, net.Addr, error) {
 }
 
 func (c *UDPConnOverTCP) WriteTo(b []byte, _ net.Addr) (int, error) {
-	packet := newDatagramPacket(b, len(b)-2)
+	buf := config.LPool.Get().([]byte)[:]
+	n := copy(buf, b)
+	defer config.LPool.Put(buf)
+
+	packet := newDatagramPacket(buf, n)
 	if err := packet.Write(c.Conn); err != nil {
 		return 0, err
 	}

pkg/core/tunhandler.go

@@ -99,12 +99,7 @@ func (d *Device) readFromTUN(ctx context.Context) {
 		}
 
 		plog.G(ctx).Debugf("[TUN] SRC: %s, DST: %s, Protocol: %s, Length: %d", src, dst, layers.IPProtocol(protocol).String(), n)
-		util.SafeWrite(d.tunInbound, &Packet{
-			data:   buf[:],
-			length: n,
-			src:    src,
-			dst:    dst,
-		}, func(v *Packet) {
+		util.SafeWrite(d.tunInbound, NewPacket(buf[:], n, src, dst), func(v *Packet) {
 			config.LPool.Put(v.data[:])
 			plog.G(context.Background()).Errorf("Drop packet, SRC: %s, DST: %s, Protocol: %s, Length: %d", v.src, v.dst, layers.IPProtocol(protocol).String(), v.length)
 		})
@@ -227,12 +222,7 @@ func (p *Peer) readFromConn(ctx context.Context) {
 		}
 		p.addToRouteMapUDP(ctx, src, from)
 		plog.G(context.Background()).Errorf("[TUN] SRC: %s, DST: %s, Protocol: %s, Length: %d", src, dst, layers.IPProtocol(protocol).String(), n)
-		p.tunInbound <- &Packet{
-			data:   buf[:],
-			length: n,
-			src:    src,
-			dst:    dst,
-		}
+		p.tunInbound <- NewPacket(buf[:], n, src, dst)
 	}
 }
 

pkg/core/tunhandlerclient.go

@@ -89,7 +89,7 @@ func handlePacketClient(ctx context.Context, tunInbound <-chan *Packet, tunOutbo
 				})
 				continue
 			}
-			_, err := packetConn.WriteTo(packet.data[:packet.length+2], remoteAddr)
+			_, err := packetConn.WriteTo(packet.data[:packet.length], remoteAddr)
 			config.LPool.Put(packet.data[:])
 			if err != nil {
 				util.SafeWrite(errChan, errors.Wrap(err, fmt.Sprintf("failed to write packet to remote %s", remoteAddr)))
@@ -113,7 +113,7 @@ func handlePacketClient(ctx context.Context, tunInbound <-chan *Packet, tunOutbo
 				config.LPool.Put(buf[:])
 				continue
 			}
-			util.SafeWrite(tunOutbound, &Packet{data: buf[:], length: n}, func(v *Packet) {
+			util.SafeWrite(tunOutbound, NewPacket(buf[:], n, nil, nil), func(v *Packet) {
 				config.LPool.Put(v.data[:])
 				plog.G(context.Background()).Errorf("Drop packet, LocalAddr: %s, Remote: %s, Length: %d", packetConn.LocalAddr(), remoteAddr, v.length)
 			})