Dragonfly2

1
/*
2
 *     Copyright 2020 The Dragonfly Authors
3
 *
4
 * Licensed under the Apache License, Version 2.0 (the "License");
5
 * you may not use this file except in compliance with the License.
6
 * You may obtain a copy of the License at
7
 *
8
 *      http://www.apache.org/licenses/LICENSE-2.0
9
 *
10
 * Unless required by applicable law or agreed to in writing, software
11
 * distributed under the License is distributed on an "AS IS" BASIS,
12
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13
 * See the License for the specific language governing permissions and
14
 * limitations under the License.
15
 */
16

17
package peer
18

19
import (
20
	"context"
21
	"crypto/md5"
22
	"encoding/hex"
23
	"errors"
24
	"fmt"
25
	"io"
26
	"math"
27
	"net/http"
28
	"net/http/httptest"
29
	"net/url"
30
	"os"
31
	"testing"
32
	"time"
33

34
	"github.com/go-http-utils/headers"
35
	testifyassert "github.com/stretchr/testify/assert"
36
	"github.com/stretchr/testify/require"
37
	"google.golang.org/grpc/status"
38

39
	commonv1 "d7y.io/api/v2/pkg/apis/common/v1"
40

41
	"d7y.io/dragonfly/v2/client/daemon/test"
42
	logger "d7y.io/dragonfly/v2/internal/dflog"
43
	nethttp "d7y.io/dragonfly/v2/pkg/net/http"
44
	"d7y.io/dragonfly/v2/pkg/source"
45
	"d7y.io/dragonfly/v2/pkg/source/clients/httpprotocol"
46
)
47

48
func TestPieceDownloader_isConnectionError(t *testing.T) {
49
	tests := []struct {
50
		name   string
51
		err    error
52
		expect bool
53
	}{
54
		{
55
			name: "connection error",
56
			err: &pieceDownloadError{
57
				connectionError: true,
58
			},
59
			expect: true,
60
		},
61
		{
62
			name: "connection ok",
63
			err: &pieceDownloadError{
64
				connectionError: false,
65
			},
66
			expect: false,
67
		},
68
		{
69
			name:   "no error",
70
			err:    nil,
71
			expect: false,
72
		},
73
	}
74

75
	for _, tc := range tests {
76
		t.Run(tc.name, func(t *testing.T) {
77
			assert := testifyassert.New(t)
78
			assert.Equal(tc.expect, isConnectionError(tc.err))
79
		})
80
	}
81
}
82

83
func TestPieceDownloader_isPieceNotFound(t *testing.T) {
84
	tests := []struct {
85
		name   string
86
		err    error
87
		expect bool
88
	}{
89
		{
90
			name: "piece not found",
91
			err: &pieceDownloadError{
92
				statusCode: http.StatusNotFound,
93
			},
94
			expect: true,
95
		},
96
		{
97
			name:   "piece found",
98
			err:    &pieceDownloadError{},
99
			expect: false,
100
		},
101
		{
102
			name:   "no error",
103
			err:    nil,
104
			expect: false,
105
		},
106
	}
107

108
	for _, tc := range tests {
109
		t.Run(tc.name, func(t *testing.T) {
110
			assert := testifyassert.New(t)
111
			assert.Equal(tc.expect, isPieceNotFound(tc.err))
112
		})
113
	}
114
}
115

116
func TestPieceDownloader_isBackSourceError(t *testing.T) {
117
	tests := []struct {
118
		name   string
119
		err    error
120
		expect bool
121
	}{
122
		{
123
			name:   "back source error",
124
			err:    &backSourceError{},
125
			expect: true,
126
		},
127
		{
128
			name:   "unexpected status code error",
129
			err:    &source.UnexpectedStatusCodeError{},
130
			expect: true,
131
		},
132
		{
133
			name:   "no error",
134
			err:    nil,
135
			expect: false,
136
		},
137
	}
138

139
	for _, tc := range tests {
140
		t.Run(tc.name, func(t *testing.T) {
141
			assert := testifyassert.New(t)
142
			assert.Equal(tc.expect, isBackSourceError(tc.err))
143
		})
144
	}
145
}
146

147
func TestPieceDownloader_Error(t *testing.T) {
148
	tests := []struct {
149
		name   string
150
		err    error
151
		expect func(t *testing.T, errMsg string)
152
	}{
153
		{
154
			name: "no request URL",
155
			err: &pieceDownloadError{
156
				connectionError: true,
157
				target:          "",
158
				err:             errors.New("http: nil Request.URL"),
159
			},
160
			expect: func(t *testing.T, errMsg string) {
161
				assert := testifyassert.New(t)
162
				assert.Equal("connect with  with error: http: nil Request.URL", errMsg)
163
			},
164
		},
165
		{
166
			name: "request URL not found",
167
			err: &pieceDownloadError{
168
				connectionError: false,
169
				target:          "http://www.x.yy",
170
				status:          "404 Not Found",
171
			},
172
			expect: func(t *testing.T, errMsg string) {
173
				assert := testifyassert.New(t)
174
				assert.Equal("download http://www.x.yy with error status: 404 Not Found", errMsg)
175
			},
176
		},
177
		{
178
			name: "resource not found",
179
			err: &backSourceError{
180
				st: status.New(5, "Resource not found."),
181
			},
182
			expect: func(t *testing.T, errMsg string) {
183
				assert := testifyassert.New(t)
184
				assert.Equal("rpc error: code = NotFound desc = Resource not found.", errMsg)
185
			},
186
		},
187
	}
188

189
	for _, tc := range tests {
190
		t.Run(tc.name, func(t *testing.T) {
191
			errMsg := tc.err.Error()
192
			tc.expect(t, errMsg)
193
		})
194
	}
195
}
196

197
func TestPieceDownloader_DownloadPiece(t *testing.T) {
198
	assert := testifyassert.New(t)
199
	source.UnRegister("http")
200
	require.Nil(t, source.Register("http", httpprotocol.NewHTTPSourceClient(), httpprotocol.Adapter))
201
	defer source.UnRegister("http")
202
	testData, err := os.ReadFile(test.File)
203
	assert.Nil(err, "load test file")
204
	pieceDownloadTimeout := 30 * time.Second
205

206
	tests := []struct {
207
		handleFunc      func(w http.ResponseWriter, r *http.Request)
208
		taskID          string
209
		pieceRange      string
210
		rangeStart      uint64
211
		rangeSize       uint32
212
		targetPieceData []byte
213
	}{
214
		{
215
			handleFunc: func(w http.ResponseWriter, r *http.Request) {
216
				assert.Equal("/download/tas/task-0", r.URL.Path)
217
				data := []byte("test test ")
218
				w.Header().Set(headers.ContentLength, fmt.Sprintf("%d", len(data)))
219
				if _, err := w.Write(data); err != nil {
220
					t.Error(err)
221
				}
222
			},
223
			taskID:          "task-0",
224
			pieceRange:      "bytes=0-9",
225
			rangeStart:      0,
226
			rangeSize:       10,
227
			targetPieceData: []byte("test test "),
228
		},
229
		{
230
			handleFunc: func(w http.ResponseWriter, r *http.Request) {
231
				assert.Equal("/download/tas/task-1", r.URL.Path)
232
				rg := nethttp.MustParseRange(r.Header.Get("Range"), math.MaxInt64)
233
				w.Header().Set(headers.ContentLength, fmt.Sprintf("%d", rg.Length))
234
				if _, err := w.Write(testData[rg.Start : rg.Start+rg.Length]); err != nil {
235
					t.Error(err)
236
				}
237
			},
238
			taskID:          "task-1",
239
			pieceRange:      "bytes=0-99",
240
			rangeStart:      0,
241
			rangeSize:       100,
242
			targetPieceData: testData[:100],
243
		},
244
		{
245
			handleFunc: func(w http.ResponseWriter, r *http.Request) {
246
				assert.Equal("/download/tas/task-2", r.URL.Path)
247
				rg := nethttp.MustParseRange(r.Header.Get("Range"), math.MaxInt64)
248
				w.Header().Set(headers.ContentLength, fmt.Sprintf("%d", rg.Length))
249
				if _, err := w.Write(testData[rg.Start : rg.Start+rg.Length]); err != nil {
250
					t.Error(err)
251
				}
252
			},
253
			taskID:          "task-2",
254
			pieceRange:      fmt.Sprintf("bytes=512-%d", len(testData)-1),
255
			rangeStart:      512,
256
			rangeSize:       uint32(len(testData) - 512),
257
			targetPieceData: testData[512:],
258
		},
259
		{
260
			handleFunc: func(w http.ResponseWriter, r *http.Request) {
261
				assert.Equal("/download/tas/task-3", r.URL.Path)
262
				rg := nethttp.MustParseRange(r.Header.Get("Range"), math.MaxInt64)
263
				w.Header().Set(headers.ContentLength, fmt.Sprintf("%d", rg.Length))
264
				if _, err := w.Write(testData[rg.Start : rg.Start+rg.Length]); err != nil {
265
					t.Error(err)
266
				}
267
			},
268
			taskID:          "task-3",
269
			pieceRange:      "bytes=512-1024",
270
			rangeStart:      512,
271
			rangeSize:       513,
272
			targetPieceData: testData[512:1025],
273
		},
274
	}
275

276
	for _, tt := range tests {
277
		server := httptest.NewServer(http.HandlerFunc(tt.handleFunc))
278
		addr, _ := url.Parse(server.URL)
279
		pd := NewPieceDownloader(pieceDownloadTimeout, nil)
280
		hash := md5.New()
281
		hash.Write(tt.targetPieceData)
282
		digest := hex.EncodeToString(hash.Sum(nil)[:16])
283
		r, c, err := pd.DownloadPiece(context.Background(), &DownloadPieceRequest{
284
			TaskID:     tt.taskID,
285
			DstPid:     "",
286
			DstAddr:    addr.Host,
287
			CalcDigest: true,
288
			piece: &commonv1.PieceInfo{
289
				PieceNum:    0,
290
				RangeStart:  tt.rangeStart,
291
				RangeSize:   tt.rangeSize,
292
				PieceMd5:    digest,
293
				PieceOffset: tt.rangeStart,
294
				PieceStyle:  commonv1.PieceStyle_PLAIN,
295
			},
296
			log: logger.With("test", "test"),
297
		})
298
		assert.Nil(err, "downloaded piece should success")
299

300
		data, err := io.ReadAll(r)
301
		assert.Nil(err, "read piece data should success")
302
		c.Close()
303

304
		assert.Equal(data, tt.targetPieceData, "downloaded piece data should match")
305
		server.Close()
306
	}
307
}
308