1 /*
2  * SRT - Secure, Reliable, Transport
3  * Copyright (c) 2020 Haivision Systems Inc.
4  *
5  * This Source Code Form is subject to the terms of the Mozilla Public
6  * License, v. 2.0. If a copy of the MPL was not distributed with this
7  * file, You can obtain one at http://mozilla.org/MPL/2.0/.
8  *
9  * Based on the proposal by Russell Greene (Issue #440)
10  *
11  */
12 
13 #include <gtest/gtest.h>
14 
15 #ifdef _WIN32
16 #define INC_SRT_WIN_WINTIME // exclude gettimeofday from srt headers
17 #endif
18 
19 #include "srt.h"
20 
21 #include <thread>
22 #include <fstream>
23 #include <ctime>
24 #include <vector>
25 
26 //#pragma comment (lib, "ws2_32.lib")
27 
TEST(Transmission,FileUpload)28 TEST(Transmission, FileUpload)
29 {
30     srt_startup();
31 
32     // Generate the source file
33     // We need a file that will contain more data
34     // than can be contained in one sender buffer.
35 
36     SRTSOCKET sock_lsn = srt_create_socket(), sock_clr = srt_create_socket();
37 
38     int tt = SRTT_FILE;
39     srt_setsockflag(sock_lsn, SRTO_TRANSTYPE, &tt, sizeof tt);
40     srt_setsockflag(sock_clr, SRTO_TRANSTYPE, &tt, sizeof tt);
41 
42     // Configure listener
43     sockaddr_in sa_lsn = sockaddr_in();
44     sa_lsn.sin_family = AF_INET;
45     sa_lsn.sin_addr.s_addr = INADDR_ANY;
46     sa_lsn.sin_port = htons(5555);
47 
48     // Find unused a port not used by any other service.
49     // Otherwise srt_connect may actually connect.
50     int bind_res = -1;
51     for (int port = 5000; port <= 5555; ++port)
52     {
53         sa_lsn.sin_port = htons(port);
54         bind_res = srt_bind(sock_lsn, (sockaddr*)&sa_lsn, sizeof sa_lsn);
55         if (bind_res == 0)
56         {
57             std::cout << "Running test on port " << port << "\n";
58             break;
59         }
60 
61         ASSERT_TRUE(bind_res == SRT_EINVOP) << "Bind failed not due to an occupied port. Result " << bind_res;
62     }
63 
64     ASSERT_GE(bind_res, 0);
65 
66     srt_bind(sock_lsn, (sockaddr*)&sa_lsn, sizeof sa_lsn);
67 
68     int optval = 0;
69     int optlen = sizeof optval;
70     ASSERT_EQ(srt_getsockflag(sock_lsn, SRTO_SNDBUF, &optval, &optlen), 0);
71     const size_t filesize = 7 * optval;
72 
73     {
74         std::cout << "WILL CREATE source file with size=" << filesize << " (= 7 * " << optval << "[sndbuf])\n";
75         std::ofstream outfile("file.source", std::ios::out | std::ios::binary);
76         ASSERT_EQ(!!outfile, true) << srt_getlasterror_str();
77 
78         srand(time(0));
79 
80         for (size_t i = 0; i < filesize; ++i)
81         {
82             char outbyte = rand() % 255;
83             outfile.write(&outbyte, 1);
84         }
85     }
86 
87     srt_listen(sock_lsn, 1);
88 
89     // Start listener-receiver thread
90 
91     bool thread_exit = false;
92 
93     auto client = std::thread([&]
94     {
95         sockaddr_in remote;
96         int len = sizeof remote;
97         const SRTSOCKET accepted_sock = srt_accept(sock_lsn, (sockaddr*)&remote, &len);
98         ASSERT_GT(accepted_sock, 0);
99 
100         if (accepted_sock == SRT_INVALID_SOCK)
101         {
102             std::cerr << srt_getlasterror_str() << std::endl;
103             EXPECT_NE(srt_close(sock_lsn), SRT_ERROR);
104             return;
105         }
106 
107         std::ofstream copyfile("file.target", std::ios::out | std::ios::trunc | std::ios::binary);
108 
109         std::vector<char> buf(1456);
110 
111         for (;;)
112         {
113             int n = srt_recv(accepted_sock, buf.data(), 1456);
114             ASSERT_NE(n, SRT_ERROR);
115             if (n == 0)
116             {
117                 std::cerr << "Received 0 bytes, breaking.\n";
118                 break;
119             }
120 
121             // Write to file any amount of data received
122             copyfile.write(buf.data(), n);
123         }
124 
125         EXPECT_NE(srt_close(accepted_sock), SRT_ERROR);
126 
127         thread_exit = true;
128     });
129 
130     sockaddr_in sa = sockaddr_in();
131     sa.sin_family = AF_INET;
132     sa.sin_port = sa_lsn.sin_port;
133     ASSERT_EQ(inet_pton(AF_INET, "127.0.0.1", &sa.sin_addr), 1);
134 
135     srt_connect(sock_clr, (sockaddr*)&sa, sizeof(sa));
136 
137     std::cout << "Connection initialized" << std::endl;
138 
139     std::ifstream ifile("file.source", std::ios::in | std::ios::binary);
140     std::vector<char> buf(1456);
141 
142     for (;;)
143     {
144         size_t n = ifile.read(buf.data(), 1456).gcount();
145         size_t shift = 0;
146         while (n > 0)
147         {
148             const int st = srt_send(sock_clr, buf.data()+shift, n);
149             ASSERT_GT(st, 0) << srt_getlasterror_str();
150 
151             n -= st;
152             shift += st;
153         }
154 
155         if (ifile.eof())
156         {
157             break;
158         }
159 
160         ASSERT_EQ(ifile.good(), true);
161     }
162 
163     // Finished sending, close the socket
164     std::cout << "Finished sending, closing sockets:\n";
165     srt_close(sock_clr);
166     srt_close(sock_lsn);
167 
168     std::cout << "Sockets closed, joining receiver thread\n";
169     client.join();
170 
171     std::ifstream tarfile("file.target");
172     EXPECT_EQ(!!tarfile, true);
173 
174     tarfile.seekg(0, std::ios::end);
175     size_t tar_size = tarfile.tellg();
176     EXPECT_EQ(tar_size, filesize);
177 
178     std::cout << "Comparing files\n";
179     // Compare files
180     tarfile.seekg(0, std::ios::end);
181     ifile.seekg(0, std::ios::beg);
182 
183     for (size_t i = 0; i < tar_size; ++i)
184     {
185         EXPECT_EQ(ifile.get(), tarfile.get());
186     }
187 
188     EXPECT_EQ(ifile.get(), EOF);
189     EXPECT_EQ(tarfile.get(), EOF);
190 
191     remove("file.source");
192     remove("file.target");
193 
194     (void)srt_cleanup();
195 }
196