From 84b38f4c6e065d3f9314a8abbb3f5857ed4fa44e Mon Sep 17 00:00:00 2001
From: Lucas Manning <lucasmanning@google.com>
Date: Mon, 8 Nov 2021 13:26:02 -0800
Subject: Add reference counting to packet buffers.

PiperOrigin-RevId: 408426639
---
 pkg/tcpip/link/tun/device.go | 2 ++
 1 file changed, 2 insertions(+)

(limited to 'pkg/tcpip/link/tun')

diff --git a/pkg/tcpip/link/tun/device.go b/pkg/tcpip/link/tun/device.go
index fa2131c28..5230ac281 100644
--- a/pkg/tcpip/link/tun/device.go
+++ b/pkg/tcpip/link/tun/device.go
@@ -76,6 +76,7 @@ func (d *Device) Release(ctx context.Context) {
 
 	// Decrease refcount if there is an endpoint associated with this file.
 	if d.endpoint != nil {
+		d.endpoint.Drain()
 		d.endpoint.RemoveNotify(d.notifyHandle)
 		d.endpoint.DecRef(ctx)
 		d.endpoint = nil
@@ -231,6 +232,7 @@ func (d *Device) Write(data []byte) (int64, error) {
 		ReserveHeaderBytes: len(ethHdr),
 		Data:               buffer.View(data).ToVectorisedView(),
 	})
+	defer pkt.DecRef()
 	copy(pkt.LinkHeader().Push(len(ethHdr)), ethHdr)
 	endpoint.InjectLinkAddr(protocol, remote, pkt)
 	return dataLen, nil
-- 
cgit v1.2.3