Fix more duplication

Author: lixmal

client/firewall/uspfilter/nat.go

@@ -26,6 +26,10 @@ const (
 	// Port offsets in TCP/UDP headers
 	sourcePortOffset      = 0
 	destinationPortOffset = 2
+
+	// IP address offsets in IPv4 header
+	sourceIPOffset      = 12
+	destinationIPOffset = 16
 )
 
 // ipv4Checksum calculates IPv4 header checksum.
@@ -374,17 +378,17 @@ func (m *Manager) translateInboundReverse(packetData []byte, d *decoder) bool {
 	return true
 }
 
-// rewritePacketDestination replaces the destination IP in the packet and updates checksums.
-func (m *Manager) rewritePacketDestination(packetData []byte, d *decoder, newIP netip.Addr) error {
+// rewritePacketIP replaces an IP address (source or destination) in the packet and updates checksums.
+func (m *Manager) rewritePacketIP(packetData []byte, d *decoder, newIP netip.Addr, ipOffset int) error {
 	if len(packetData) < 20 || d.decoded[0] != layers.LayerTypeIPv4 || !newIP.Is4() {
 		return ErrIPv4Only
 	}
 
-	var oldDst [4]byte
-	copy(oldDst[:], packetData[16:20])
-	newDst := newIP.As4()
+	var oldIP [4]byte
+	copy(oldIP[:], packetData[ipOffset:ipOffset+4])
+	newIPBytes := newIP.As4()
 
-	copy(packetData[16:20], newDst[:])
+	copy(packetData[ipOffset:ipOffset+4], newIPBytes[:])
 
 	ipHeaderLen := int(d.ip4.IHL) * 4
 	if ipHeaderLen < 20 || ipHeaderLen > len(packetData) {
@@ -398,9 +402,9 @@ func (m *Manager) rewritePacketDestination(packetData []byte, d *decoder, newIP
 	if len(d.decoded) > 1 {
 		switch d.decoded[1] {
 		case layers.LayerTypeTCP:
-			m.updateTCPChecksum(packetData, ipHeaderLen, oldDst[:], newDst[:])
+			m.updateTCPChecksum(packetData, ipHeaderLen, oldIP[:], newIPBytes[:])
 		case layers.LayerTypeUDP:
-			m.updateUDPChecksum(packetData, ipHeaderLen, oldDst[:], newDst[:])
+			m.updateUDPChecksum(packetData, ipHeaderLen, oldIP[:], newIPBytes[:])
 		case layers.LayerTypeICMPv4:
 			m.updateICMPChecksum(packetData, ipHeaderLen)
 		}
@@ -409,39 +413,14 @@ func (m *Manager) rewritePacketDestination(packetData []byte, d *decoder, newIP
 	return nil
 }
 
+// rewritePacketDestination replaces the destination IP in the packet and updates checksums.
+func (m *Manager) rewritePacketDestination(packetData []byte, d *decoder, newIP netip.Addr) error {
+	return m.rewritePacketIP(packetData, d, newIP, destinationIPOffset)
+}
+
 // rewritePacketSource replaces the source IP address in the packet and updates checksums.
 func (m *Manager) rewritePacketSource(packetData []byte, d *decoder, newIP netip.Addr) error {
-	if len(packetData) < 20 || d.decoded[0] != layers.LayerTypeIPv4 || !newIP.Is4() {
-		return ErrIPv4Only
-	}
-
-	var oldSrc [4]byte
-	copy(oldSrc[:], packetData[12:16])
-	newSrc := newIP.As4()
-
-	copy(packetData[12:16], newSrc[:])
-
-	ipHeaderLen := int(d.ip4.IHL) * 4
-	if ipHeaderLen < 20 || ipHeaderLen > len(packetData) {
-		return errInvalidIPHeaderLength
-	}
-
-	binary.BigEndian.PutUint16(packetData[10:12], 0)
-	ipChecksum := ipv4Checksum(packetData[:ipHeaderLen])
-	binary.BigEndian.PutUint16(packetData[10:12], ipChecksum)
-
-	if len(d.decoded) > 1 {
-		switch d.decoded[1] {
-		case layers.LayerTypeTCP:
-			m.updateTCPChecksum(packetData, ipHeaderLen, oldSrc[:], newSrc[:])
-		case layers.LayerTypeUDP:
-			m.updateUDPChecksum(packetData, ipHeaderLen, oldSrc[:], newSrc[:])
-		case layers.LayerTypeICMPv4:
-			m.updateICMPChecksum(packetData, ipHeaderLen)
-		}
-	}
-
-	return nil
+	return m.rewritePacketIP(packetData, d, newIP, sourceIPOffset)
 }
 
 // updateTCPChecksum updates TCP checksum after IP address change per RFC 1624.