From ff49ece55d9d67f7043be5374711d8f267077091 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E4=B8=96=E7=95=8C?= <i@sekai.icu>
Date: Wed, 27 Aug 2025 15:45:18 +0800
Subject: [PATCH] Fix checksum changes

---
 internal/gtcpip/header/ipv4.go | 4 +++-
 internal/gtcpip/header/tcp.go  | 4 +++-
 internal/gtcpip/header/udp.go  | 5 +++--
 stack_gvisor_tcp.go            | 2 ++
 4 files changed, 11 insertions(+), 4 deletions(-)

diff --git a/internal/gtcpip/header/ipv4.go b/internal/gtcpip/header/ipv4.go
index d1cbf7c..ad06f38 100644
--- a/internal/gtcpip/header/ipv4.go
+++ b/internal/gtcpip/header/ipv4.go
@@ -479,6 +479,7 @@ func (b IPv4) SetDestinationAddress(addr tcpip.Address) {
 
 // CalculateChecksum calculates the checksum of the IPv4 header.
 func (b IPv4) CalculateChecksum() uint16 {
+	// return checksum.Checksum(b[:b.HeaderLength()], 0)
 	xsum0 := checksum.Checksum(b[:xsum], 0)
 	xsum0 = checksum.Checksum(b[xsum+2:b.HeaderLength()], xsum0)
 	return xsum0
@@ -573,7 +574,8 @@ func (b IPv4) IsChecksumValid() bool {
 	//        same set of octets, including the checksum field.  If the result
 	//        is all 1 bits (-0 in 1's complement arithmetic), the check
 	//        succeeds.
-	return b.CalculateChecksum() == 0xffff
+	//return b.CalculateChecksum() == 0xffff
+	return checksum.Checksum(b[:b.HeaderLength()], 0) == 0xffff
 }
 
 // IsV4MulticastAddress determines if the provided address is an IPv4 multicast
diff --git a/internal/gtcpip/header/tcp.go b/internal/gtcpip/header/tcp.go
index da5d3d8..1b58df8 100644
--- a/internal/gtcpip/header/tcp.go
+++ b/internal/gtcpip/header/tcp.go
@@ -351,6 +351,7 @@ func (b TCP) SetUrgentPointer(urgentPointer uint16) {
 // and the checksum of the segment data.
 func (b TCP) CalculateChecksum(partialChecksum uint16) uint16 {
 	// Calculate the rest of the checksum.
+	// return checksum.Checksum(b[:b.DataOffset()], partialChecksum)
 	xsum := checksum.Checksum(b[:TCPChecksumOffset], partialChecksum)
 	xsum = checksum.Checksum(b[TCPChecksumOffset+2:b.DataOffset()], xsum)
 	return xsum
@@ -360,7 +361,8 @@ func (b TCP) CalculateChecksum(partialChecksum uint16) uint16 {
 func (b TCP) IsChecksumValid(src, dst tcpip.Address, payloadChecksum, payloadLength uint16) bool {
 	xsum := PseudoHeaderChecksum(TCPProtocolNumber, src.AsSlice(), dst.AsSlice(), uint16(b.DataOffset())+payloadLength)
 	xsum = checksum.Combine(xsum, payloadChecksum)
-	return b.CalculateChecksum(xsum) == 0xffff
+	// return b.CalculateChecksum(xsum) == 0xffff
+	return checksum.Checksum(b[:b.DataOffset()], xsum) == 0xffff
 }
 
 // Options returns a slice that holds the unparsed TCP options in the segment.
diff --git a/internal/gtcpip/header/udp.go b/internal/gtcpip/header/udp.go
index eac9d63..a995a17 100644
--- a/internal/gtcpip/header/udp.go
+++ b/internal/gtcpip/header/udp.go
@@ -114,8 +114,9 @@ func (b UDP) SetLength(length uint16) {
 // checksum of the network-layer pseudo-header and the checksum of the payload.
 func (b UDP) CalculateChecksum(partialChecksum uint16) uint16 {
 	// Calculate the rest of the checksum.\
+	// return checksum.Checksum(b[:UDPMinimumSize], partialChecksum)
 	xsum := checksum.Checksum(b[:udpChecksum], partialChecksum)
-	xsum = checksum.Checksum(b[udpChecksum+2:], xsum)
+	xsum = checksum.Checksum(b[udpChecksum+2:UDPMinimumSize], xsum)
 	return xsum
 }
 
@@ -123,7 +124,7 @@ func (b UDP) CalculateChecksum(partialChecksum uint16) uint16 {
 func (b UDP) IsChecksumValid(src, dst tcpip.Address, payloadChecksum uint16) bool {
 	xsum := PseudoHeaderChecksum(UDPProtocolNumber, dst.AsSlice(), src.AsSlice(), b.Length())
 	xsum = checksum.Combine(xsum, payloadChecksum)
-	return b.CalculateChecksum(xsum) == 0xffff
+	return checksum.Checksum(b[:UDPMinimumSize], xsum) == 0xffff
 }
 
 // Encode encodes all the fields of the UDP header.
diff --git a/stack_gvisor_tcp.go b/stack_gvisor_tcp.go
index 024f4b4..1592799 100644
--- a/stack_gvisor_tcp.go
+++ b/stack_gvisor_tcp.go
@@ -51,6 +51,7 @@ func (f *TCPForwarder) HandlePacket(id stack.TransportEndpointID, pkt *stack.Pac
 			ipHdr.SetDestinationAddressWithChecksumUpdate(ipHdr.SourceAddress())
 			ipHdr.SetSourceAddressWithChecksumUpdate(inet4LoopbackAddress)
 			tcpHdr := header.TCP(pkt.TransportHeader().Slice())
+			tcpHdr.SetChecksum(0)
 			tcpHdr.SetChecksum(^checksum.Checksum(tcpHdr.Payload(), tcpHdr.CalculateChecksum(
 				header.PseudoHeaderChecksum(header.TCPProtocolNumber, ipHdr.SourceAddress(), ipHdr.DestinationAddress(), ipHdr.PayloadLength()),
 			)))
@@ -64,6 +65,7 @@ func (f *TCPForwarder) HandlePacket(id stack.TransportEndpointID, pkt *stack.Pac
 			ipHdr.SetDestinationAddress(ipHdr.SourceAddress())
 			ipHdr.SetSourceAddress(inet6LoopbackAddress)
 			tcpHdr := header.TCP(pkt.TransportHeader().Slice())
+			tcpHdr.SetChecksum(0)
 			tcpHdr.SetChecksum(^checksum.Checksum(tcpHdr.Payload(), tcpHdr.CalculateChecksum(
 				header.PseudoHeaderChecksum(header.TCPProtocolNumber, ipHdr.SourceAddress(), ipHdr.DestinationAddress(), ipHdr.PayloadLength()),
 			)))