/
conn.go
91 lines (86 loc) · 2.17 KB
/
conn.go
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
package quickproto
import (
"bytes"
"net"
"strconv"
"github.com/Nigel2392/simplecrypto/aes"
)
// Convenience function to craft an address from an IP and a port.
// Port could be string or int, IP must be string.
func CraftAddr(ip string, port any) string {
switch port := port.(type) {
case int:
return ip + ":" + strconv.Itoa(port)
case string:
return ip + ":" + port
default:
panic("invalid port type provided")
}
}
// ReadConn reads a message from a connection.
func ReadConn(conn net.Conn, conf *Config, aes_key *[32]byte, compress bool) (*Message, error) {
msg := conf.NewMessage()
buf := make([]byte, conf.BufSize)
var data []byte
// read until ending delimiter is found.
for !bytes.Contains(data, msg.EndingDelimiter()) {
// read data from connection.
n, err := conn.Read(buf)
if err != nil {
return nil, err
}
data = append(data, buf[:n]...)
// flush buffer.
}
// decrypt data if needed.
if compress {
data = bytes.TrimSuffix(data, msg.EndingDelimiter())
var err error
data, err = GZIPdecompress(data)
if err != nil {
return nil, err
}
}
if aes_key != nil {
var err error
if !compress {
data = bytes.TrimSuffix(data, msg.EndingDelimiter())
}
if data, err = aes.Decrypt(data, aes_key); err != nil {
return nil, err
}
data = append(data, msg.EndingDelimiter()...)
}
msg.Data = data
return msg.Parse()
}
// WriteConn writes a message to a connection and encrypts it if needed.
func WriteConn(conn net.Conn, msg *Message, aes_key *[32]byte, compress bool) error {
// Write data to connection.
send, err := msg.Generate()
if err != nil {
return err
}
if aes_key != nil {
send.Data = bytes.TrimSuffix(send.Data, msg.EndingDelimiter())
send.Data, err = aes.Encrypt(send.Data, aes_key)
if err != nil {
return err
}
if !compress {
send.Data = append(send.Data, msg.EndingDelimiter()...)
}
}
if compress {
send.Data, err = GZIPcompress(send.Data)
if err != nil {
return err
}
send.Data = append(send.Data, msg.EndingDelimiter()...)
}
_, err = conn.Write(send.Data)
if err != nil {
return err
}
return nil
}