7 files changed, 65 insertions, 0 deletions
diff --git a/pkg/tcpip/header/ipv4.go b/pkg/tcpip/header/ipv4.go
index 76839eb92..d8111a439 100644
--- a/pkg/tcpip/header/ipv4.go
+++ b/pkg/tcpip/header/ipv4.go
@@ -248,6 +248,18 @@ func (b IPv4) SetDestinationAddress(addr tcpip.Address) {
 	copy(b[dstAddr:dstAddr+IPv4AddressSize], addr)
 }
 
+func (b IPv4) UpdateTTL(n uint8) {
+	var sum uint32
+	var old uint16
+
+	old = binary.BigEndian.Uint16(b[ttl:])
+	b[ttl] -= n
+	sum = uint32(old) + uint32(^binary.BigEndian.Uint16(b[ttl:]) & 0xffff)
+	sum = sum + uint32(b.Checksum())
+	sum = (sum & 0xffff) + (sum>>16)
+	b.SetChecksum(uint16(sum + (sum>>16)))
+}
+
 // CalculateChecksum calculates the checksum of the ipv4 header.
 func (b IPv4) CalculateChecksum() uint16 {
 	return Checksum(b[:b.HeaderLength()], 0)
diff --git a/pkg/tcpip/header/ipv6.go b/pkg/tcpip/header/ipv6.go
index 76e88e9b3..ff1aa937e 100644
--- a/pkg/tcpip/header/ipv6.go
+++ b/pkg/tcpip/header/ipv6.go
@@ -226,6 +226,11 @@ func (b IPv6) SetNextHeader(v uint8) {
 	b[nextHdr] = v
 }
 
+// SetHopLimit sets the value of the "hop limit" field of the ipv6 header.
+func (b IPv6) SetHopLimit(v uint8) {
+	b[hopLimit] = v
+}
+
 // SetChecksum implements Network.SetChecksum. Given that IPv6 doesn't have a
 // checksum, it is empty.
 func (IPv6) SetChecksum(uint16) {
diff --git a/pkg/tcpip/network/arp/arp.go b/pkg/tcpip/network/arp/arp.go
index e9fcc89a8..9ffacd544 100644
--- a/pkg/tcpip/network/arp/arp.go
+++ b/pkg/tcpip/network/arp/arp.go
@@ -124,6 +124,10 @@ func (e *endpoint) HandlePacket(r *stack.Route, pkt tcpip.PacketBuffer) {
 	}
 }
 
+func (*endpoint) UpdateTTL(tcpip.PacketBuffer, uint8) *tcpip.Error {
+	return tcpip.ErrNotSupported
+}
+
 // protocol implements stack.NetworkProtocol and stack.LinkAddressResolver.
 type protocol struct {
 }
diff --git a/pkg/tcpip/network/ipv4/ipv4.go b/pkg/tcpip/network/ipv4/ipv4.go
index 4f1742938..9953f83cd 100644
--- a/pkg/tcpip/network/ipv4/ipv4.go
+++ b/pkg/tcpip/network/ipv4/ipv4.go
@@ -21,6 +21,7 @@
 package ipv4
 
 import (
+	"fmt"
 	"sync/atomic"
 
 	"gvisor.dev/gvisor/pkg/tcpip"
@@ -407,6 +408,24 @@ func (e *endpoint) HandlePacket(r *stack.Route, pkt tcpip.PacketBuffer) {
 	e.dispatcher.DeliverTransportPacket(r, p, pkt)
 }
 
+func (*endpoint) UpdateTTL(pkt tcpip.PacketBuffer, n uint8) *tcpip.Error {
+	headerView := pkt.Header.View()
+	h := header.IPv4(headerView)
+	if !h.IsValid(pkt.Data.Size()) {
+		fmt.Println("Malformed packet", len(h), int(h.HeaderLength()), int(h.TotalLength()), header.IPVersion(h), pkt.Data.Size())
+		//r.Stats().IP.MalformedPacketsReceived.Increment()
+		return tcpip.ErrInvalidOptionValue
+	}
+	//pkt.NetworkHeader = headerView[:h.HeaderLength()]
+
+	ttl := h.TTL()
+	if ttl <= n {
+		return tcpip.ErrTimeout
+	}
+	h.UpdateTTL(n)
+	return nil
+}
+
 // Close cleans up resources associated with the endpoint.
 func (e *endpoint) Close() {}
 
diff --git a/pkg/tcpip/network/ipv6/ipv6.go b/pkg/tcpip/network/ipv6/ipv6.go
index 9aef5234b..df86dd3b2 100644
--- a/pkg/tcpip/network/ipv6/ipv6.go
+++ b/pkg/tcpip/network/ipv6/ipv6.go
@@ -21,6 +21,7 @@
 package ipv6
 
 import (
+	"fmt"
 	"sync/atomic"
 
 	"gvisor.dev/gvisor/pkg/tcpip"
@@ -189,6 +190,23 @@ func (e *endpoint) HandlePacket(r *stack.Route, pkt tcpip.PacketBuffer) {
 	e.dispatcher.DeliverTransportPacket(r, p, pkt)
 }
 
+func (*endpoint) UpdateTTL(pkt tcpip.PacketBuffer, n uint8) *tcpip.Error {
+	headerView := pkt.Header.View()
+	h := header.IPv6(headerView)
+	if !h.IsValid(pkt.Data.Size()) {
+		fmt.Println("Malformed packet", len(h), header.IPVersion(h), pkt.Data.Size())
+		//r.Stats().IP.MalformedPacketsReceived.Increment()
+		return tcpip.ErrInvalidOptionValue
+	}
+
+	ttl := h.HopLimit()
+	if ttl <= n {
+		return tcpip.ErrTimeout
+	}
+	h.SetHopLimit(ttl - n)
+	return nil
+}
+
 // Close cleans up resources associated with the endpoint.
 func (*endpoint) Close() {}
 
diff --git a/pkg/tcpip/stack/nic.go b/pkg/tcpip/stack/nic.go
index 9dcb1d52c..af211f953 100644
--- a/pkg/tcpip/stack/nic.go
+++ b/pkg/tcpip/stack/nic.go
@@ -1261,6 +1261,11 @@ func (n *NIC) DeliverNetworkPacket(linkEP LinkEndpoint, remote, local tcpip.Link
 			return
 		}
 
+		if err := r.ref.ep.UpdateTTL(pkt, 1); err != nil {
+			fmt.Println("Drop")
+			return
+		}
+		
 		// n doesn't have a destination endpoint.
 		// Send the packet out of n.
 		// TODO(b/128629022): move this logic to route.WritePacket.
diff --git a/pkg/tcpip/stack/registration.go b/pkg/tcpip/stack/registration.go
index fa28b46b1..1231dee63 100644
--- a/pkg/tcpip/stack/registration.go
+++ b/pkg/tcpip/stack/registration.go
@@ -267,6 +267,8 @@ type NetworkEndpoint interface {
 	// HandlePacket takes ownership of pkt.
 	HandlePacket(r *Route, pkt tcpip.PacketBuffer)
 
+	UpdateTTL(pkt tcpip.PacketBuffer, n uint8) *tcpip.Error
+
 	// Close is called when the endpoint is reomved from a stack.
 	Close()
 }