forked from apple/swift-nio-extras
-
Notifications
You must be signed in to change notification settings - Fork 0
/
LengthFieldPrepender.swift
135 lines (117 loc) · 5.18 KB
/
LengthFieldPrepender.swift
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
//===----------------------------------------------------------------------===//
//
// This source file is part of the SwiftNIO open source project
//
// Copyright (c) 2017-2018 Apple Inc. and the SwiftNIO project authors
// Licensed under Apache License v2.0
//
// See LICENSE.txt for license information
// See CONTRIBUTORS.txt for the list of SwiftNIO project authors
//
// SPDX-License-Identifier: Apache-2.0
//
//===----------------------------------------------------------------------===//
import NIO
extension ByteBuffer {
@discardableResult
@inlinable
mutating func write24UInt(
_ integer: UInt32,
endianness: Endianness = .big
) -> Int {
precondition(integer & 0xFF_FF_FF == integer, "integer value does not fit into 24 bit integer")
switch endianness {
case .little:
return writeInteger(UInt8(integer & 0xFF), endianness: .little) +
writeInteger(UInt16((integer >> 8) & 0xFF_FF), endianness: .little)
case .big:
return writeInteger(UInt16((integer >> 8) & 0xFF_FF), endianness: .big) +
writeInteger(UInt8(integer & 0xFF), endianness: .big)
}
}
}
public enum LengthFieldPrependerError: Error {
case messageDataTooLongForLengthField
}
///
/// An encoder that takes a `ByteBuffer` message and prepends the number of bytes in the message.
/// The length field is always the same fixed length specified on construction.
/// These bytes contain a binary specification of the message size.
///
/// For example, if you received a packet with the 3 byte length (BCD)...
/// Given that the specified header length is 1 byte, there would be a single byte prepended which contains the number 3
/// +---+-----+
/// | A | BCD | ('A' contains 0x03)
/// +---+-----+
/// This initial prepended byte is called the 'length field'.
///
public final class LengthFieldPrepender: ChannelOutboundHandler {
///
/// An enumeration to describe the length of a piece of data in bytes.
///
public enum ByteLength {
case one
case two
case four
case eight
fileprivate var bitLength: NIOLengthFieldBitLength {
switch self {
case .one: return .oneByte
case .two: return .twoBytes
case .four: return .fourBytes
case .eight: return .eightBytes
}
}
}
public typealias OutboundIn = ByteBuffer
public typealias OutboundOut = ByteBuffer
private let lengthFieldLength: NIOLengthFieldBitLength
private let lengthFieldEndianness: Endianness
private var lengthBuffer: ByteBuffer?
/// Create `LengthFieldPrepender` with a given length field length.
///
/// - parameters:
/// - lengthFieldLength: The length of the field specifying the remaining length of the frame.
/// - lengthFieldEndianness: The endianness of the field specifying the remaining length of the frame.
///
public convenience init(lengthFieldLength: ByteLength, lengthFieldEndianness: Endianness = .big) {
self.init(lengthFieldBitLength: lengthFieldLength.bitLength, lengthFieldEndianness: lengthFieldEndianness)
}
public init(lengthFieldBitLength: NIOLengthFieldBitLength, lengthFieldEndianness: Endianness = .big) {
// The value contained in the length field must be able to be represented by an integer type on the platform.
// ie. .eight == 64bit which would not fit into the Int type on a 32bit platform.
precondition(lengthFieldBitLength.length <= Int.bitWidth/8)
self.lengthFieldLength = lengthFieldBitLength
self.lengthFieldEndianness = lengthFieldEndianness
}
public func write(context: ChannelHandlerContext, data: NIOAny, promise: EventLoopPromise<Void>?) {
let dataBuffer = self.unwrapOutboundIn(data)
let dataLength = dataBuffer.readableBytes
guard dataLength <= self.lengthFieldLength.max else {
promise?.fail(LengthFieldPrependerError.messageDataTooLongForLengthField)
return
}
var dataLengthBuffer: ByteBuffer
if let existingBuffer = self.lengthBuffer {
dataLengthBuffer = existingBuffer
dataLengthBuffer.clear()
} else {
dataLengthBuffer = context.channel.allocator.buffer(capacity: self.lengthFieldLength.length)
self.lengthBuffer = dataLengthBuffer
}
switch self.lengthFieldLength.bitLength {
case .bits8:
dataLengthBuffer.writeInteger(UInt8(dataLength), endianness: self.lengthFieldEndianness)
case .bits16:
dataLengthBuffer.writeInteger(UInt16(dataLength), endianness: self.lengthFieldEndianness)
case .bits24:
dataLengthBuffer.write24UInt(UInt32(dataLength), endianness: self.lengthFieldEndianness)
case .bits32:
dataLengthBuffer.writeInteger(UInt32(dataLength), endianness: self.lengthFieldEndianness)
case .bits64:
dataLengthBuffer.writeInteger(UInt64(dataLength), endianness: self.lengthFieldEndianness)
}
context.write(self.wrapOutboundOut(dataLengthBuffer), promise: nil)
context.write(data, promise: promise)
}
}