1/*
2 * Licensed to the Apache Software Foundation (ASF) under one
3 * or more contributor license agreements. See the NOTICE file
4 * distributed with this work for additional information
5 * regarding copyright ownership. The ASF licenses this file
6 * to you under the Apache License, Version 2.0 (the
7 * "License"); you may not use this file except in compliance
8 * with the License. You may obtain a copy of the License at
9 *
10 *   http://www.apache.org/licenses/LICENSE-2.0
11 *
12 * Unless required by applicable law or agreed to in writing,
13 * software distributed under the License is distributed on an
14 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15 * KIND, either express or implied. See the License for the
16 * specific language governing permissions and limitations
17 * under the License.
18 */
19
20package thrift
21
22import (
23	"bytes"
24	"math"
25	"strings"
26	"testing"
27)
28
29func TestReadWriteBinaryProtocol(t *testing.T) {
30	ReadWriteProtocolTest(t, NewTBinaryProtocolFactoryDefault())
31}
32
33const (
34	safeReadBytesSource = `
35Lorem ipsum dolor sit amet, consectetur adipiscing elit. Integer sit amet
36tincidunt nibh. Phasellus vel convallis libero, sit amet posuere quam. Nullam
37blandit velit at nibh fringilla, sed egestas erat dapibus. Sed hendrerit
38tincidunt accumsan. Curabitur consectetur bibendum dui nec hendrerit. Fusce quis
39turpis nec magna efficitur volutpat a ut nibh. Vestibulum odio risus, tristique
40a nisi et, congue mattis mi. Vivamus a nunc justo. Mauris molestie sagittis
41magna, hendrerit auctor lectus egestas non. Phasellus pretium, odio sit amet
42bibendum feugiat, velit nunc luctus erat, ac bibendum mi dui molestie nulla.
43Nullam fermentum magna eu elit vehicula tincidunt. Etiam ornare laoreet
44dignissim. Ut sed nunc ac neque vulputate fermentum. Morbi volutpat dapibus
45magna, at porttitor quam facilisis a. Donec eget fermentum risus. Aliquam erat
46volutpat.
47
48Phasellus molestie id ante vel iaculis. Fusce eget quam nec quam viverra laoreet
49vitae a dui. Mauris blandit blandit dui, iaculis interdum diam mollis at. Morbi
50vel sem et.
51`
52	safeReadBytesSourceLen = len(safeReadBytesSource)
53)
54
55func TestSafeReadBytes(t *testing.T) {
56	srcData := []byte(safeReadBytesSource)
57
58	for _, c := range []struct {
59		label     string
60		askedSize int32
61		dataSize  int
62	}{
63		{
64			label:     "normal",
65			askedSize: 100,
66			dataSize:  100,
67		},
68		{
69			label:     "max-askedSize",
70			askedSize: math.MaxInt32,
71			dataSize:  safeReadBytesSourceLen,
72		},
73	} {
74		t.Run(c.label, func(t *testing.T) {
75			data := bytes.NewReader(srcData[:c.dataSize])
76			buf, err := safeReadBytes(c.askedSize, data)
77			if len(buf) != c.dataSize {
78				t.Errorf(
79					"Expected to read %d bytes, got %d",
80					c.dataSize,
81					len(buf),
82				)
83			}
84			if !strings.HasPrefix(safeReadBytesSource, string(buf)) {
85				t.Errorf("Unexpected read data: %q", buf)
86			}
87			if int32(c.dataSize) < c.askedSize {
88				// We expect error in this case
89				if err == nil {
90					t.Errorf(
91						"Expected error when dataSize %d < askedSize %d, got nil",
92						c.dataSize,
93						c.askedSize,
94					)
95				}
96			} else {
97				// We expect no error in this case
98				if err != nil {
99					t.Errorf(
100						"Expected no error when dataSize %d >= askedSize %d, got: %v",
101						c.dataSize,
102						c.askedSize,
103						err,
104					)
105				}
106			}
107		})
108	}
109}
110
111func generateSafeReadBytesBenchmark(askedSize int32, dataSize int) func(b *testing.B) {
112	return func(b *testing.B) {
113		data := make([]byte, dataSize)
114		b.ResetTimer()
115		for i := 0; i < b.N; i++ {
116			safeReadBytes(askedSize, bytes.NewReader(data))
117		}
118	}
119}
120
121func TestSafeReadBytesAlloc(t *testing.T) {
122	if testing.Short() {
123		// NOTE: Since this test runs a benchmark test, it takes at
124		// least 1 second.
125		//
126		// In general we try to avoid unit tests taking that long to run,
127		// but it's to verify a security issue so we made an exception
128		// here:
129		// https://issues.apache.org/jira/browse/THRIFT-5322
130		t.Skip("skipping test in short mode.")
131	}
132
133	const (
134		askedSize = int32(math.MaxInt32)
135		dataSize  = 4096
136	)
137
138	// The purpose of this test is that in the case a string header says
139	// that it has a string askedSize bytes long, the implementation should
140	// not just allocate askedSize bytes upfront. So when there're actually
141	// not enough data to be read (dataSize), the actual allocated bytes
142	// should be somewhere between dataSize and askedSize.
143	//
144	// Different approachs could have different memory overheads, so this
145	// target is arbitrary in nature. But when dataSize is small enough
146	// compare to askedSize, half the askedSize is a good and safe target.
147	const target = int64(askedSize) / 2
148
149	bm := testing.Benchmark(generateSafeReadBytesBenchmark(askedSize, dataSize))
150	actual := bm.AllocedBytesPerOp()
151	if actual > target {
152		t.Errorf(
153			"Expected allocated bytes per op to be <= %d, got %d",
154			target,
155			actual,
156		)
157	} else {
158		t.Logf("Allocated bytes: %d B/op", actual)
159	}
160}
161
162func BenchmarkSafeReadBytes(b *testing.B) {
163	for _, c := range []struct {
164		label     string
165		askedSize int32
166		dataSize  int
167	}{
168		{
169			label:     "normal",
170			askedSize: 100,
171			dataSize:  100,
172		},
173		{
174			label:     "max-askedSize",
175			askedSize: math.MaxInt32,
176			dataSize:  4096,
177		},
178	} {
179		b.Run(c.label, generateSafeReadBytesBenchmark(c.askedSize, c.dataSize))
180	}
181}
182