使用原生 tcp 进行数据传输
Aidan Engineer
 2025-04-20 14:11:20      3.2k 字  15 分钟

用其做一些传输操作时替代原本 HTTP 的实现方式，可以用在命令行工具之类的操作上

先看下大体架构：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
// Conn 是你需要实现的一种连接类型，它支持下面描述的若干接口；
// 为了实现这些接口，你需要设计一个基于 TCP 的简单协议；
type Conn struct {
}

// Send 传入一个 key 表示发送者将要传输的数据对应的标识；
// 返回 writer 可供发送者分多次写入大量该 key 对应的数据；
// 当发送者已将该 key 对应的所有数据写入后，调用 writer.Close 告知接收者：该 key 的数据已经完全写入；
func (conn *Conn) Send(key string) (writer io.WriteCloser, err error) {
}

// Receive 返回一个 key 表示接收者将要接收到的数据对应的标识；
// 返回的 reader 可供接收者多次读取该 key 对应的数据；
// 当 reader 返回 io.EOF 错误时，表示接收者已经完整接收该 key 对应的数据；
func (conn *Conn) Receive() (key string, reader io.Reader, err error) {
}

// Close 关闭你实现的连接对象及其底层的 TCP 连接 
func (conn *Conn) Close() {
}

// NewConn 从一个 TCP 连接得到一个你实现的连接对象 
func NewConn(conn net.Conn) *Conn {
}

结构看作固定给出，以下就不做扩展了，在此基础上进行后续开发即可，其中各结构与函数的作用在注释中已做完整解释

  1. 首先对于 Conn 进行定义，golang 的 net 包下有完备的 tcp 实现，复用即可，根据函数的发送接收，考虑进行读写锁的添加防止调用 tcp 实际写入时产生的数据错乱
  2. Send 和 Receive 用来实现写入和接收，其通过接口进行参数返回，所以需要根据接口进行实际结构体的定义
  3. 写入会使用 writer.Close 用来告知完成，所以在结构体的实现时需要有标识位进行记录

实现的结构体如下：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
type Conn struct {
	tcp net.Conn
	wMu sync.Mutex // 写操作锁 
	rMu sync.Mutex // 读操作锁 
}

// 根据 io.WriteCloser 接口实现 
type sendWriter struct {
	conn   *Conn
	closed bool // 通过关闭标志告知数据传输结束 
}

func (w *sendWriter) Write(p []byte) (int, error) {
	if w.closed {
		return 0, io.ErrClosedPipe
	}

	// 写入数据长度，用于后续获取 
	lengthHeader := make([]byte, 4)
	binary.BigEndian.PutUint32(lengthHeader, uint32(len(p)))

	if _, err := w.conn.tcp.Write(lengthHeader); err != nil {
		return 0, err
	}
	n, err := w.conn.tcp.Write(p)
	if err != nil {
		return n, err
	}

	// fmt.Println("send data:", string(p[:n]))

	return n, nil
}

func (w *sendWriter) Close() error {
	if w.closed {
		return nil
	}
	w.closed = true

	// 写入长度为 0 的数据，读取判断结束 
	lengthHeader := make([]byte, 4)
	if _, err := w.conn.tcp.Write(lengthHeader); err != nil {
		return err
	}
	w.conn.wMu.Unlock()
	return nil
}

// 根据 io.Reader 接口实现 
type receiveReader struct {
	conn      *Conn
	remaining []byte
	closed    bool
}

func (r *receiveReader) Read(p []byte) (int, error) {
	if r.closed {
		return 0, io.EOF
	}

	if len(r.remaining) > 0 {
		n := copy(p, r.remaining)
		r.remaining = r.remaining[n:]
		return n, nil
	}

	// 读取固定位数保留的数据长度，按长度读取数据 
	lengthHeader := make([]byte, 4)
	if _, err := io.ReadFull(r.conn.tcp, lengthHeader); err != nil {
		return 0, err
	}
	dataLength := binary.BigEndian.Uint32(lengthHeader)

	if dataLength == 0 {
		r.closed = true
		r.conn.rMu.Unlock()
		return 0, io.EOF
	}

	data := make([]byte, dataLength)
	if _, err := io.ReadFull(r.conn.tcp, data); err != nil {
		return 0, err
	}

	n := copy(p, data)
	if n < len(data) {
		r.remaining = data[n:]
	}

	// fmt.Println("receive data:", string(data[:n]))

	return n, nil
}

以上代码所有具体操作都是借助 Conn 结构体实现，而 Conn 又通过 net.conn 进行直接引用，顾没有理解障碍

需要注意在接收时的 receiveReader.remaining 主要用来进行传输数据的保留，防止数据拆包导致的无法续读，当然也可以使用下标进行实现

根据自定义结构实现的代码逻辑如下：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
func (conn *Conn) Send(key string) (writer io.WriteCloser, err error) {
	conn.wMu.Lock()

	keyBytes := []byte(key)
	lengthHeader := make([]byte, 4)
	binary.BigEndian.PutUint32(lengthHeader, uint32(len(keyBytes)))

	if _, err := conn.tcp.Write(lengthHeader); err != nil {
		conn.wMu.Unlock()
		return nil, err
	}
	if _, err := conn.tcp.Write(keyBytes); err != nil {
		conn.wMu.Unlock()
		return nil, err
	}

	// fmt.Println("send key:", key)

	writer = &sendWriter{
		conn:   conn,
		closed: false,
	}
	return writer, nil
}

func (conn *Conn) Receive() (key string, reader io.Reader, err error) {
	conn.rMu.Lock()

	keyLengthHeader := make([]byte, 4)
	if _, err := io.ReadFull(conn.tcp, keyLengthHeader); err != nil {
		conn.rMu.Unlock()
		return "", nil, err
	}
	keyLength := binary.BigEndian.Uint32(keyLengthHeader)

	keyBytes := make([]byte, keyLength)
	if _, err := io.ReadFull(conn.tcp, keyBytes); err != nil {
		conn.rMu.Unlock()
		return "", nil, err
	}
	key = string(keyBytes)

	// fmt.Println("receive key:", key)

	reader = &receiveReader{
		conn: conn,
	}
	return key, reader, nil
}

func (conn *Conn) Close() {
	conn.tcp.Close()
}

func NewConn(conn net.Conn) *Conn {
	return &Conn{
		tcp: conn,
	}
}

以上可以看作最简单的基于 tcp 的应用层开发，后续可能在这个点子上做一些别的东西，比如客户端文件上传之类的东西，完整代码如下，因来源特殊性，此代码禁止外传 

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
package main

import (
	"crypto/rand"
	"crypto/sha256"
	"encoding/binary"
	"encoding/hex"
	"fmt"
	"hash"
	"io"
	"net"
	"sync"
)

// Conn 是你需要实现的一种连接类型，它支持下面描述的若干接口；
// 为了实现这些接口，你需要设计一个基于 TCP 的简单协议；
type Conn struct {
	tcp net.Conn
	wMu sync.Mutex // 写操作锁 
	rMu sync.Mutex // 读操作锁 
}

// Send 传入一个 key 表示发送者将要传输的数据对应的标识；
// 返回 writer 可供发送者分多次写入大量该 key 对应的数据；
// 当发送者已将该 key 对应的所有数据写入后，调用 writer.Close 告知接收者：该 key 的数据已经完全写入；
func (conn *Conn) Send(key string) (writer io.WriteCloser, err error) {
	conn.wMu.Lock()

	keyBytes := []byte(key)
	lengthHeader := make([]byte, 4)
	binary.BigEndian.PutUint32(lengthHeader, uint32(len(keyBytes)))

	if _, err := conn.tcp.Write(lengthHeader); err != nil {
		conn.wMu.Unlock()
		return nil, err
	}
	if _, err := conn.tcp.Write(keyBytes); err != nil {
		conn.wMu.Unlock()
		return nil, err
	}

	// fmt.Println("send key:", key)

	writer = &sendWriter{
		conn:   conn,
		closed: false,
	}
	return writer, nil
}

// Receive 返回一个 key 表示接收者将要接收到的数据对应的标识；
// 返回的 reader 可供接收者多次读取该 key 对应的数据；
// 当 reader 返回 io.EOF 错误时，表示接收者已经完整接收该 key 对应的数据；
func (conn *Conn) Receive() (key string, reader io.Reader, err error) {
	conn.rMu.Lock()

	keyLengthHeader := make([]byte, 4)
	if _, err := io.ReadFull(conn.tcp, keyLengthHeader); err != nil {
		conn.rMu.Unlock()
		return "", nil, err
	}
	keyLength := binary.BigEndian.Uint32(keyLengthHeader)

	keyBytes := make([]byte, keyLength)
	if _, err := io.ReadFull(conn.tcp, keyBytes); err != nil {
		conn.rMu.Unlock()
		return "", nil, err
	}
	key = string(keyBytes)

	// fmt.Println("receive key:", key)

	reader = &receiveReader{
		conn: conn,
	}
	return key, reader, nil
}

// Close 关闭你实现的连接对象及其底层的 TCP 连接 
func (conn *Conn) Close() {
	conn.tcp.Close()
}

// NewConn 从一个 TCP 连接得到一个你实现的连接对象 
func NewConn(conn net.Conn) *Conn {
	return &Conn{
		tcp: conn,
	}
}

// 除了上面规定的接口，你还可以自行定义新的类型，变量和函数以满足实现需求 

// ============= 以下为自定义实现 =============
// 根据 io.WriteCloser 接口实现 
type sendWriter struct {
	conn   *Conn
	closed bool // 通过关闭标志告知数据传输结束 
}

func (w *sendWriter) Write(p []byte) (int, error) {
	if w.closed {
		return 0, io.ErrClosedPipe
	}

	// 写入数据长度，用于后续获取 
	lengthHeader := make([]byte, 4)
	binary.BigEndian.PutUint32(lengthHeader, uint32(len(p)))

	if _, err := w.conn.tcp.Write(lengthHeader); err != nil {
		return 0, err
	}
	n, err := w.conn.tcp.Write(p)
	if err != nil {
		return n, err
	}

	// fmt.Println("send data:", string(p[:n]))

	return n, nil
}

func (w *sendWriter) Close() error {
	if w.closed {
		return nil
	}
	w.closed = true

	// 写入长度为 0 的数据，读取判断结束 
	lengthHeader := make([]byte, 4)
	if _, err := w.conn.tcp.Write(lengthHeader); err != nil {
		return err
	}
	w.conn.wMu.Unlock()
	return nil
}

// 根据 io.Reader 接口实现 
type receiveReader struct {
	conn      *Conn
	remaining []byte
	closed    bool
}

func (r *receiveReader) Read(p []byte) (int, error) {
	if r.closed {
		return 0, io.EOF
	}

	if len(r.remaining) > 0 {
		n := copy(p, r.remaining)
		r.remaining = r.remaining[n:]
		return n, nil
	}

	// 读取固定位数保留的数据长度，按长度读取数据 
	lengthHeader := make([]byte, 4)
	if _, err := io.ReadFull(r.conn.tcp, lengthHeader); err != nil {
		return 0, err
	}
	dataLength := binary.BigEndian.Uint32(lengthHeader)

	if dataLength == 0 {
		r.closed = true
		r.conn.rMu.Unlock()
		return 0, io.EOF
	}

	data := make([]byte, dataLength)
	if _, err := io.ReadFull(r.conn.tcp, data); err != nil {
		return 0, err
	}

	n := copy(p, data)
	if n < len(data) {
		r.remaining = data[n:]
	}

	// fmt.Println("receive data:", string(data[:n]))

	return n, nil
}

//////////////////////////////////////////////
///////// 接下来的代码为测试代码，请勿修改 /////////
//////////////////////////////////////////////

// 连接到测试服务器，获得一个你实现的连接对象 
func dial(serverAddr string) *Conn {
	conn, err := net.Dial("tcp", serverAddr)
	if err != nil {
		panic(err)
	}
	return NewConn(conn)
}

// 启动测试服务器 
func startServer(handle func(*Conn)) net.Listener {
	ln, err := net.Listen("tcp", ":0")
	if err != nil {
		panic(err)
	}
	go func() {
		for {
			conn, err := ln.Accept()
			if err != nil {
				fmt.Println("[WARNING] ln.Accept", err)
				return
			}
			go handle(NewConn(conn))
		}
	}()
	return ln
}

// 简单断言 
func assertEqual[T comparable](actual T, expected T) {
	if actual != expected {
		panic(fmt.Sprintf("actual:%v expected:%v\n", actual, expected))
	}
}

// 简单 case：单连接，双向传输少量数据 
func testCase0() {
	const (
		key  = "Bible"
		data = `Then I heard the voice of the Lord saying, “Whom shall I send? And who will go for us?”
And I said, “Here am I. Send me!”
Isaiah 6:8`
	)
	ln := startServer(func(conn *Conn) {
		// 服务端等待客户端进行传输 
		_key, reader, err := conn.Receive()
		if err != nil {
			panic(err)
		}
		assertEqual(_key, key)
		dataB, err := io.ReadAll(reader)
		if err != nil {
			panic(err)
		}
		assertEqual(string(dataB), data)

		// 服务端向客户端进行传输 
		writer, err := conn.Send(key)
		if err != nil {
			panic(err)
		}
		n, err := writer.Write([]byte(data))
		if err != nil {
			panic(err)
		}
		if n != len(data) {
			panic(n)
		}
		conn.Close()
	})
	//goland:noinspection GoUnhandledErrorResult
	defer ln.Close()

	conn := dial(ln.Addr().String())
	// 客户端向服务端传输 
	writer, err := conn.Send(key)
	if err != nil {
		panic(err)
	}
	n, err := writer.Write([]byte(data))
	if n != len(data) {
		panic(n)
	}
	err = writer.Close()
	if err != nil {
		panic(err)
	}
	// 客户端等待服务端传输 
	_key, reader, err := conn.Receive()
	if err != nil {
		panic(err)
	}
	assertEqual(_key, key)
	dataB, err := io.ReadAll(reader)
	if err != nil {
		panic(err)
	}
	assertEqual(string(dataB), data)
	conn.Close()
}

// 生成一个随机 key
func newRandomKey() string {
	buf := make([]byte, 8)
	_, err := rand.Read(buf)
	if err != nil {
		panic(err)
	}
	return hex.EncodeToString(buf)
}

// 读取随机数据，并返回随机数据的校验和：用于验证数据是否完整传输 
func readRandomData(reader io.Reader, hash hash.Hash) (checksum string) {
	hash.Reset()
	var buf = make([]byte, 23<<20) // 调用者读取时的 buf 大小不是固定的，你的实现中不可假定 buf 为固定值 
	for {
		n, err := reader.Read(buf)
		if err == io.EOF {
			break
		}
		if err != nil {
			panic(err)
		}
		_, err = hash.Write(buf[:n])
		if err != nil {
			panic(err)
		}
	}
	checksum = hex.EncodeToString(hash.Sum(nil))
	return checksum
}

// 写入随机数据，并返回随机数据的校验和：用于验证数据是否完整传输 
func writeRandomData(writer io.Writer, hash hash.Hash) (checksum string) {
	hash.Reset()
	const (
		dataSize = 500 << 20 // 一个 key 对应 500MB 随机二进制数据，dataSize 也可以是其他值，你的实现中不可假定 dataSize 为固定值 
		bufSize  = 1 << 20   // 调用者写入时的 buf 大小不是固定的，你的实现中不可假定 buf 为固定值 
	)
	var (
		buf  = make([]byte, bufSize)
		size = 0
	)
	for i := 0; i < dataSize/bufSize; i++ {
		_, err := rand.Read(buf)
		if err != nil {
			panic(err)
		}
		_, err = hash.Write(buf)
		if err != nil {
			panic(err)
		}
		n, err := writer.Write(buf)
		if err != nil {
			panic(err)
		}
		size += n
	}
	if size != dataSize {
		panic(size)
	}
	checksum = hex.EncodeToString(hash.Sum(nil))
	return checksum
}

// 复杂 case：多连接，双向传输，大量数据，多个不同的 key
func testCase1() {
	var (
		mapKeyToChecksum = map[string]string{}
		lock             sync.Mutex
	)
	ln := startServer(func(conn *Conn) {
		// 服务端等待客户端进行传输 
		key, reader, err := conn.Receive()
		if err != nil {
			panic(err)
		}
		var (
			h         = sha256.New()
			_checksum = readRandomData(reader, h)
		)
		lock.Lock()
		checksum, keyExist := mapKeyToChecksum[key]
		lock.Unlock()
		if !keyExist {
			panic(fmt.Sprintln(key, "not exist"))
		}
		assertEqual(_checksum, checksum)

		// 服务端向客户端连续进行 2 次传输 
		for _, key := range []string{newRandomKey(), newRandomKey()} {
			writer, err := conn.Send(key)
			if err != nil {
				panic(err)
			}
			checksum := writeRandomData(writer, h)
			lock.Lock()
			mapKeyToChecksum[key] = checksum
			lock.Unlock()
			err = writer.Close() // 表明该 key 的所有数据已传输完毕 
			if err != nil {
				panic(err)
			}
		}
		conn.Close()
	})
	//goland:noinspection GoUnhandledErrorResult
	defer ln.Close()

	conn := dial(ln.Addr().String())
	// 客户端向服务端传输 
	var (
		key = newRandomKey()
		h   = sha256.New()
	)
	writer, err := conn.Send(key)
	if err != nil {
		panic(err)
	}
	checksum := writeRandomData(writer, h)
	lock.Lock()
	mapKeyToChecksum[key] = checksum
	lock.Unlock()
	err = writer.Close()
	if err != nil {
		panic(err)
	}

	// 客户端等待服务端的多次传输 
	keyCount := 0
	for {
		key, reader, err := conn.Receive()
		if err == io.EOF {
			// 服务端所有的数据均传输完毕，关闭连接 
			break
		}
		if err != nil {
			panic(err)
		}
		_checksum := readRandomData(reader, h)
		lock.Lock()
		checksum, keyExist := mapKeyToChecksum[key]
		lock.Unlock()
		if !keyExist {
			panic(fmt.Sprintln(key, "not exist"))
		}
		assertEqual(_checksum, checksum)
		keyCount++
	}
	assertEqual(keyCount, 2)
	conn.Close()
}

func main() {
	testCase0()
	testCase1()
}
  • 本文标题：使用原生 tcp 进行数据传输
  • 本文作者：Aidan
  • 创建时间：2025-04-20 14:11:20
  • 本文链接：https://aidanblog.top/data_transfer-tcp/
  • 版权声明：本博客所有文章除特别声明外，均采用 BY-NC-SA 许可协议。转载请注明出处！
 评论