rdp/Segment.cpp at master · guangqianpeng/rdp · GitHub

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
//
// Created by frank on 17-8-14.
//

#include "Segment.h"

#include <netinet/in.h>
#include <cstring>
#include <cassert>

namespace
{
struct Header
{
	uint32_t ident;
	uint8_t cmd;
	uint8_t frg;
	uint16_t wnd;
	uint32_t ts;
	uint32_t seq;
	uint32_t una;
	uint32_t len;
};

Header createHeader(const Segment& seg)
{
	Header hdr;
	hdr.ident = htonl(seg.ident);
	hdr.cmd = static_cast<uint8_t>(seg.cmd);
	hdr.frg = static_cast<uint8_t >(seg.frg);
	hdr.wnd = htons(static_cast<uint16_t>(seg.wnd));
	hdr.ts = htonl(seg.ts);
	hdr.seq = htonl(seg.seq);
	hdr.una = htonl(seg.una);
	hdr.len = htonl(seg.len);
	return hdr;
}

}

SegmentPtr createSegment(const char *data, uint32_t size)
{
	uint32_t segTotalLen = sizeof(Segment) + size;

	auto seg = static_cast<Segment*>(malloc(segTotalLen));

	/* other data members cannot initialize here */
	seg->len = size;
	if (size > 0)
		memcpy(seg->data, data, size);

	return SegmentPtr(seg);
}

uint32_t encodeSegment(const Segment& seg, char*& buffer, uint32_t& size)
{
	uint32_t totalLen = sizeof(Header) + seg.len;

	assert(size >= totalLen);

	Header hdr = createHeader(seg);

	memcpy(buffer, &hdr, sizeof(hdr));
	if (seg.len > 0)
		memcpy(buffer + sizeof(hdr), seg.data, seg.len);

	buffer += totalLen;
	size -= totalLen;

	return totalLen;
}

uint32_t encodeSegment(const Segment& seg, std::vector<char>& buffer)
{
	Header hdr = createHeader(seg);
	char *ptr = reinterpret_cast<char*>(&hdr);
	buffer.insert(buffer.end(), ptr, ptr + sizeof(hdr));
	if (seg.len > 0)
		buffer.insert(buffer.end(), seg.data, seg.data + seg.len);
	return sizeof(hdr) + seg.len;
}

SegmentPtr decodeSegment(const char *&data, uint32_t &size)
{
	auto *hdr = reinterpret_cast<const Header*>(data);

	uint32_t ident = ntohl(hdr->ident);
	uint32_t cmd = hdr->cmd;
	uint32_t frg = hdr->frg;
	uint32_t wnd = ntohs(hdr->wnd);
	uint32_t ts = ntohl(hdr->ts);
	uint32_t seq = ntohl(hdr->seq);
	uint32_t una = ntohl(hdr->una);
	uint32_t len = ntohl(hdr->len);

	uint32_t segTotalLen = sizeof(Segment) + len;
	uint32_t hdrTotalLen = ENCODE_OVERHEAD + len;

	if (size < hdrTotalLen)
		return nullptr;

	auto seg = static_cast<Segment*>(malloc(segTotalLen));

	seg->ident = ident;
	seg->cmd = cmd;
	seg->frg = frg;
	seg->wnd = wnd;
	seg->ts = ts;
	seg->seq = seq;
	seg->una = una;
	seg->len = len;

	if (len > 0)
		memcpy(seg->data, data + ENCODE_OVERHEAD, len);

	data += hdrTotalLen;
	size -= hdrTotalLen;

	return SegmentPtr(seg);
}