Fix not calling decRef on merged segments

This code path is for outgoing packets, and we don't currently do memory
accounting on this path. So it wasn't breaking anything.

This change did not add a test for ref-counting issue fixed, but will switch to
the leak-checking ref-counter later when all ref-counting issues are fixed.

PiperOrigin-RevId: 373447913

4 files changed, 91 insertions, 8 deletions

diff --git a/pkg/tcpip/transport/tcp/BUILD b/pkg/tcpip/transport/tcp/BUILD
index 48417f192..0f20d3856 100644
--- a/pkg/tcpip/transport/tcp/BUILD
+++ b/pkg/tcpip/transport/tcp/BUILD
@@ -126,7 +126,15 @@ go_test(
 go_test(
     name = "tcp_test",
     size = "small",
-    srcs = ["timer_test.go"],
+    srcs = [
+        "segment_test.go",
+        "timer_test.go",
+    ],
     library = ":tcp",
-    deps = ["//pkg/sleep"],
+    deps = [
+        "//pkg/sleep",
+        "//pkg/tcpip/buffer",
+        "//pkg/tcpip/stack",
+        "@com_github_google_go_cmp//cmp:go_default_library",
+    ],
 )
diff --git a/pkg/tcpip/transport/tcp/segment.go b/pkg/tcpip/transport/tcp/segment.go
index c28641be3..7e5ba6ef7 100644
--- a/pkg/tcpip/transport/tcp/segment.go
+++ b/pkg/tcpip/transport/tcp/segment.go
@@ -140,6 +140,15 @@ func (s *segment) clone() *segment {
 	return t
 }
 
+// merge merges data in oth and clears oth.
+func (s *segment) merge(oth *segment) {
+	s.data.Append(oth.data)
+	s.dataMemSize = s.data.Size()
+
+	oth.data = buffer.VectorisedView{}
+	oth.dataMemSize = oth.data.Size()
+}
+
 // flagIsSet checks if at least one flag in flags is set in s.flags.
 func (s *segment) flagIsSet(flags header.TCPFlags) bool {
 	return s.flags&flags != 0
diff --git a/pkg/tcpip/transport/tcp/segment_test.go b/pkg/tcpip/transport/tcp/segment_test.go
new file mode 100644
index 000000000..486016fc0
--- /dev/null
+++ b/pkg/tcpip/transport/tcp/segment_test.go
@@ -0,0 +1,67 @@
+// Copyright 2021 The gVisor Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package tcp
+
+import (
+	"testing"
+
+	"github.com/google/go-cmp/cmp"
+	"gvisor.dev/gvisor/pkg/tcpip/buffer"
+	"gvisor.dev/gvisor/pkg/tcpip/stack"
+)
+
+type segmentSizeWants struct {
+	DataSize   int
+	SegMemSize int
+}
+
+func checkSegmentSize(t *testing.T, name string, seg *segment, want segmentSizeWants) {
+	t.Helper()
+	got := segmentSizeWants{
+		DataSize:   seg.data.Size(),
+		SegMemSize: seg.segMemSize(),
+	}
+	if diff := cmp.Diff(got, want); diff != "" {
+		t.Errorf("%s differs (-want +got):\n%s", name, diff)
+	}
+}
+
+func TestSegmentMerge(t *testing.T) {
+	id := stack.TransportEndpointID{}
+	seg1 := newOutgoingSegment(id, buffer.NewView(10))
+	defer seg1.decRef()
+	seg2 := newOutgoingSegment(id, buffer.NewView(20))
+	defer seg2.decRef()
+
+	checkSegmentSize(t, "seg1", seg1, segmentSizeWants{
+		DataSize:   10,
+		SegMemSize: SegSize + 10,
+	})
+	checkSegmentSize(t, "seg2", seg2, segmentSizeWants{
+		DataSize:   20,
+		SegMemSize: SegSize + 20,
+	})
+
+	seg1.merge(seg2)
+
+	checkSegmentSize(t, "seg1", seg1, segmentSizeWants{
+		DataSize:   30,
+		SegMemSize: SegSize + 30,
+	})
+	checkSegmentSize(t, "seg2", seg2, segmentSizeWants{
+		DataSize:   0,
+		SegMemSize: SegSize,
+	})
+}
diff --git a/pkg/tcpip/transport/tcp/snd.go b/pkg/tcpip/transport/tcp/snd.go
index 2b32cb7b2..f43e86677 100644
--- a/pkg/tcpip/transport/tcp/snd.go
+++ b/pkg/tcpip/transport/tcp/snd.go
@@ -716,15 +716,14 @@ func (s *sender) maybeSendSegment(seg *segment, limit int, end seqnum.Value) (se
 			// triggering bugs in poorly written DNS
 			// implementations.
 			var nextTooBig bool
-			for seg.Next() != nil && seg.Next().data.Size() != 0 {
-				if seg.data.Size()+seg.Next().data.Size() > available {
+			for nSeg := seg.Next(); nSeg != nil && nSeg.data.Size() != 0; nSeg = seg.Next() {
+				if seg.data.Size()+nSeg.data.Size() > available {
 					nextTooBig = true
 					break
 				}
-				seg.data.Append(seg.Next().data)
-
-				// Consume the segment that we just merged in.
-				s.writeList.Remove(seg.Next())
+				seg.merge(nSeg)
+				s.writeList.Remove(nSeg)
+				nSeg.decRef()
 			}
 			if !nextTooBig && seg.data.Size() < available {
 				// Segment is not full.