6 files changed, 104 insertions, 13 deletions
diff --git a/pkg/crypto/crypto_stdlib.go b/pkg/crypto/crypto_stdlib.go
index 74a55a123..514592b08 100644
--- a/pkg/crypto/crypto_stdlib.go
+++ b/pkg/crypto/crypto_stdlib.go
@@ -22,11 +22,11 @@ import (
 
 // EcdsaVerify verifies the signature in r, s of hash using ECDSA and the
 // public key, pub. Its return value records whether the signature is valid.
-func EcdsaVerify(pub *ecdsa.PublicKey, hash []byte, r, s *big.Int) bool {
-	return ecdsa.Verify(pub, hash, r, s)
+func EcdsaVerify(pub *ecdsa.PublicKey, hash []byte, r, s *big.Int) (bool, error) {
+	return ecdsa.Verify(pub, hash, r, s), nil
 }
 
 // SumSha384 returns the SHA384 checksum of the data.
-func SumSha384(data []byte) (sum384 [sha512.Size384]byte) {
-	return sha512.Sum384(data)
+func SumSha384(data []byte) ([sha512.Size384]byte, error) {
+	return sha512.Sum384(data), nil
 }
diff --git a/pkg/refsvfs2/refs_template.go b/pkg/refsvfs2/refs_template.go
index 3fbc91aa5..1102c8adc 100644
--- a/pkg/refsvfs2/refs_template.go
+++ b/pkg/refsvfs2/refs_template.go
@@ -13,7 +13,7 @@
 // limitations under the License.
 
 // Package refs_template defines a template that can be used by reference
-// counted objects.
+// counted objects. The template comes with leak checking capabilities.
 package refs_template
 
 import (
@@ -40,6 +40,14 @@ var obj *T
 // Refs implements refs.RefCounter. It keeps a reference count using atomic
 // operations and calls the destructor when the count reaches zero.
 //
+// NOTE: Do not introduce additional fields to the Refs struct. It is used by
+// many filesystem objects, and we want to keep it as small as possible (i.e.,
+// the same size as using an int64 directly) to avoid taking up extra cache
+// space. In general, this template should not be extended at the cost of
+// performance. If it does not offer enough flexibility for a particular object
+// (example: b/187877947), we should implement the RefCounter/CheckedObject
+// interfaces manually.
+//
 // +stateify savable
 type Refs struct {
 	// refCount is composed of two fields:
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.