forked from VKCOM/kittenhouse
/
udp.go
91 lines (72 loc) · 1.99 KB
/
udp.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 cmd
import (
"bytes"
"encoding/binary"
"errors"
"fmt"
"hash/crc32"
"log"
"net"
"syscall"
cmdconfig "github.com/vkcom/kittenhouse/core/cmdconfig"
"github.com/vkcom/kittenhouse/core/inmem"
)
var (
errPacketTooShort = errors.New("Packet too short")
errInvalidCRC = errors.New("Invalid CRC32")
errMissingNulByte = errors.New("Missing nul byte")
)
// parseUDPPacket parses UDP packet in the following format:
// | table_name | \0 | data | crc32 (4 byte) |
func parseUDPPacket(buf []byte) (table string, data []byte, flags byte, err error) {
if len(buf) < 5 {
return "", nil, 0, errPacketTooShort
}
crcBytes := buf[len(buf)-4:]
buf = buf[0 : len(buf)-4]
expectedCrc := binary.LittleEndian.Uint32(crcBytes)
if crc32.Update(0, crc32.IEEETable, buf) != expectedCrc {
return "", nil, 0, errInvalidCRC
}
nulPos := bytes.IndexByte(buf, 0)
if nulPos < 0 {
return "", nil, 0, errMissingNulByte
}
tableOffset := 0
flags = 0
if buf[0] < 32 {
tableOffset = 1
flags = buf[0]
}
table = string(buf[tableOffset:nulPos])
data = buf[nulPos+1:]
return table, data, flags, nil
}
func listenUDP() {
buf := make([]byte, maxUDPPacketSize)
addr, err := net.ResolveUDPAddr("udp", fmt.Sprintf("%s:%d", cmdconfig.Argv.Host, cmdconfig.Argv.Port))
if err != nil {
log.Fatalf("Could not resolve udp addr: %s", err.Error())
}
conn, err := net.ListenUDP("udp", addr)
if err != nil {
log.Fatalf("Could not listen udp addr: %s", err.Error())
}
defer conn.Close()
log.Printf("Listening %s:%d (UDP)", cmdconfig.Argv.Host, cmdconfig.Argv.Port)
for {
n, _, err := conn.ReadFromUDP(buf)
if err == syscall.EINVAL {
log.Fatalf("Could not read from UDP socket: %s", err.Error())
} else if err != nil {
log.Printf("Could not read from UDP: %s", err.Error())
continue
}
table, data, flags, err := parseUDPPacket(buf[0:n])
if err != nil {
log.Printf("Got bad packet from UDP: %s", err.Error())
continue
}
inmem.Write(table, data, flags&FlagRowBinary == FlagRowBinary)
}
}