2020-04-08 23:39:40 +00:00
|
|
|
package common
|
2020-04-08 19:53:09 +00:00
|
|
|
|
|
|
|
import (
|
|
|
|
"encoding/binary"
|
|
|
|
"io"
|
|
|
|
"net"
|
2020-10-17 12:46:22 +00:00
|
|
|
"sync"
|
2020-04-08 19:53:09 +00:00
|
|
|
"time"
|
|
|
|
)
|
|
|
|
|
|
|
|
const (
|
|
|
|
VersionTLS11 = 0x0301
|
|
|
|
VersionTLS13 = 0x0303
|
|
|
|
|
2020-04-10 15:15:23 +00:00
|
|
|
recordLayerLength = 5
|
|
|
|
|
2020-04-08 19:53:09 +00:00
|
|
|
Handshake = 22
|
|
|
|
ApplicationData = 23
|
|
|
|
)
|
|
|
|
|
|
|
|
func AddRecordLayer(input []byte, typ byte, ver uint16) []byte {
|
2020-04-10 18:09:17 +00:00
|
|
|
msgLen := len(input)
|
2020-04-17 00:03:36 +00:00
|
|
|
retLen := msgLen + recordLayerLength
|
2020-04-10 18:09:17 +00:00
|
|
|
var ret []byte
|
2020-04-17 00:03:36 +00:00
|
|
|
if cap(input) >= retLen {
|
|
|
|
ret = input[:retLen]
|
2020-04-10 18:09:17 +00:00
|
|
|
} else {
|
2020-04-17 00:03:36 +00:00
|
|
|
ret = make([]byte, retLen)
|
2020-04-10 18:09:17 +00:00
|
|
|
}
|
|
|
|
copy(ret[recordLayerLength:], input)
|
2020-04-08 19:53:09 +00:00
|
|
|
ret[0] = typ
|
2020-04-10 18:09:17 +00:00
|
|
|
ret[1] = byte(ver >> 8)
|
|
|
|
ret[2] = byte(ver)
|
|
|
|
ret[3] = byte(msgLen >> 8)
|
|
|
|
ret[4] = byte(msgLen)
|
2020-04-08 19:53:09 +00:00
|
|
|
return ret
|
|
|
|
}
|
|
|
|
|
|
|
|
type TLSConn struct {
|
|
|
|
net.Conn
|
2020-10-17 12:46:22 +00:00
|
|
|
writeM sync.Mutex
|
|
|
|
writeBuf []byte
|
|
|
|
}
|
|
|
|
|
|
|
|
func NewTLSConn(conn net.Conn) *TLSConn {
|
|
|
|
return &TLSConn{
|
|
|
|
Conn: conn,
|
|
|
|
writeBuf: make([]byte, 15000),
|
|
|
|
}
|
2020-04-08 19:53:09 +00:00
|
|
|
}
|
|
|
|
|
|
|
|
func (tls *TLSConn) LocalAddr() net.Addr {
|
|
|
|
return tls.Conn.LocalAddr()
|
|
|
|
}
|
|
|
|
|
|
|
|
func (tls *TLSConn) RemoteAddr() net.Addr {
|
|
|
|
return tls.Conn.RemoteAddr()
|
|
|
|
}
|
|
|
|
|
|
|
|
func (tls *TLSConn) SetDeadline(t time.Time) error {
|
|
|
|
return tls.Conn.SetDeadline(t)
|
|
|
|
}
|
|
|
|
|
|
|
|
func (tls *TLSConn) SetReadDeadline(t time.Time) error {
|
|
|
|
return tls.Conn.SetReadDeadline(t)
|
|
|
|
}
|
|
|
|
|
|
|
|
func (tls *TLSConn) SetWriteDeadline(t time.Time) error {
|
|
|
|
return tls.Conn.SetWriteDeadline(t)
|
|
|
|
}
|
|
|
|
|
|
|
|
func (tls *TLSConn) Read(buffer []byte) (n int, err error) {
|
|
|
|
// TCP is a stream. Multiple TLS messages can arrive at the same time,
|
|
|
|
// a single message can also be segmented due to MTU of the IP layer.
|
|
|
|
// This function guareentees a single TLS message to be read and everything
|
|
|
|
// else is left in the buffer.
|
2020-04-10 15:15:23 +00:00
|
|
|
_, err = io.ReadFull(tls.Conn, buffer[:recordLayerLength])
|
2020-04-08 19:53:09 +00:00
|
|
|
if err != nil {
|
|
|
|
return
|
|
|
|
}
|
|
|
|
|
|
|
|
dataLength := int(binary.BigEndian.Uint16(buffer[3:5]))
|
|
|
|
if dataLength > len(buffer) {
|
|
|
|
err = io.ErrShortBuffer
|
|
|
|
return
|
|
|
|
}
|
2020-08-12 10:32:39 +00:00
|
|
|
// we overwrite the record layer here
|
2020-04-08 19:53:09 +00:00
|
|
|
return io.ReadFull(tls.Conn, buffer[:dataLength])
|
|
|
|
}
|
|
|
|
|
|
|
|
func (tls *TLSConn) Write(in []byte) (n int, err error) {
|
2020-10-17 12:46:22 +00:00
|
|
|
msgLen := len(in)
|
|
|
|
tls.writeM.Lock()
|
|
|
|
tls.writeBuf = append(tls.writeBuf[:5], in...)
|
|
|
|
tls.writeBuf[0] = ApplicationData
|
|
|
|
tls.writeBuf[1] = byte(VersionTLS13 >> 8)
|
|
|
|
tls.writeBuf[2] = byte(VersionTLS13 & 0xFF)
|
|
|
|
tls.writeBuf[3] = byte(msgLen >> 8)
|
|
|
|
tls.writeBuf[4] = byte(msgLen & 0xFF)
|
|
|
|
n, err = tls.Conn.Write(tls.writeBuf[:recordLayerLength+msgLen])
|
|
|
|
tls.writeM.Unlock()
|
2020-04-10 15:15:23 +00:00
|
|
|
return n - recordLayerLength, err
|
2020-04-08 19:53:09 +00:00
|
|
|
}
|
|
|
|
|
|
|
|
func (tls *TLSConn) Close() error {
|
|
|
|
return tls.Conn.Close()
|
|
|
|
}
|