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.