/
mbserver.go
222 lines (176 loc) · 4.93 KB
/
mbserver.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
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
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
/* A basic binary memcache server
Run as: go run mbserver.go <ip>:<port>
Written by : Lynsey Haynes */
package main
import (
"encoding/binary"
"fmt"
"net"
"os"
"sync"
)
const (
HEADER_SIZE = 24
REQUEST = 0x80
RESPONSE = 0x81
GET = 0x00
SET = 0x01
NOOP = 0x0A
UNKNOWN = 0x81
SUCCESS = 0x0000
NOTFOUND = 0x0001
MAX_INT = 2147483647
)
type mbpacket struct {
}
type data struct {
val []byte
kind []byte
}
var rwmutex sync.RWMutex
var kvmap map[string]data
func checkError(err error) {
if err != nil {
fmt.Println("Error: ", err.Error())
os.Exit(1)
}
}
// --------------------------------------------------------------------------
func constructPacket(opcode uint8, extras []byte, status uint8, body []byte) []byte {
var magic []byte = []byte{0x81}
var opaque []byte = []byte{0x00, 0x00, 0x00, 0x00}
var CAS []byte = []byte{0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00}
var keylen []byte = []byte{0x00, 0x00}
var bodylen = make([]byte, 4)
var extralen = uint8(len(extras))
binary.BigEndian.PutUint32(bodylen, uint32(len(body)+len(extras)))
msg := append(magic, []byte{opcode}...)
msg = append(msg, keylen...)
msg = append(msg, []byte{extralen}...)
msg = append(msg, []byte{0x00}...)
msg = append(msg, []byte{0x00, status}...)
msg = append(msg, bodylen...)
msg = append(msg, opaque...)
msg = append(msg, CAS...)
msg = append(msg, extras...)
msg = append(msg, body...)
return msg
}
// --------------------------------------------------------------------------
// Send an error response to the connection and close it
func handleUnknown(conn *net.TCPConn) {
var empty []byte
msg := constructPacket(NOOP, empty, UNKNOWN, empty)
fmt.Printf("%X", msg)
fmt.Println(" ", len(msg))
conn.Write(msg)
}
// --------------------------------------------------------------------------
// Set the requested key to the requested value and respond with success
func handleSet(header []byte, conn *net.TCPConn) {
var keylen uint16
var extralen uint8
var bodylen uint32
keylen = binary.BigEndian.Uint16(header[2:4])
extralen = header[4]
bodylen = binary.BigEndian.Uint32(header[8:12])
if bodylen < uint32(keylen)+uint32(extralen) {
fmt.Printf("Error! Body length not long enough. Header: %X\n", header)
fmt.Println("BodyLength = ", bodylen)
return
}
body := make([]byte, bodylen)
if bodylen <= MAX_INT {
conn.SetReadBuffer(int(bodylen))
conn.Read(body)
} else { //size of body is larger than max ReadBuffer size
restlen := int(bodylen - MAX_INT)
rest := make([]byte, restlen)
conn.SetReadBuffer(MAX_INT)
conn.Read(body)
conn.SetReadBuffer(restlen)
conn.Read(rest)
body = append(body, rest...)
}
extras := body[:extralen-4]
key := string(body[extralen:(keylen + uint16(extralen))])
val := body[(uint16(extralen) + keylen):bodylen]
rwmutex.Lock()
kvmap[key] = data{val, extras}
var empty []byte
msg := constructPacket(SET, empty, SUCCESS, empty)
conn.Write(msg)
rwmutex.Unlock()
}
// --------------------------------------------------------------------------
// handle GET request. if value is there, return value, otherwise 'not found'
func handleGet(header []byte, conn *net.TCPConn) {
var keylen uint16
keylen = binary.BigEndian.Uint16(header[2:4])
keybuf := make([]byte, int(keylen))
conn.SetReadBuffer(int(keylen))
conn.Read(keybuf)
key := string(keybuf)
// ===============
rwmutex.RLock()
val := kvmap[key].val
kind := kvmap[key].kind
var msg []byte
if val != nil {
msg = constructPacket(GET, kind, SUCCESS, val)
} else {
var empty []byte
msg = constructPacket(GET, empty, NOTFOUND, []byte("Not found"))
}
conn.Write(msg)
rwmutex.RUnlock()
}
// --------------------------------------------------------------------------
// Handles incoming GET, SET, and UNKNOWN requests based on the packet header.
func handleRequest(conn *net.TCPConn) {
header := make([]byte, 24)
conn.SetReadBuffer(24)
n, err := conn.Read(header)
if err != nil || n != 24 {
return
}
magic := header[0]
if magic != REQUEST || n < HEADER_SIZE {
handleUnknown(conn)
return
}
opcode := header[1]
if opcode == GET {
handleGet(header, conn)
} else if opcode == SET {
handleSet(header, conn)
} else {
handleUnknown(conn)
}
}
// --------------------------------------------------------------------------
func handleRequests(conn *net.TCPConn) {
defer conn.Close()
for {
handleRequest(conn)
}
}
// --------------------------------------------------------------------------
//Listens for connections from possibly concurrent clients
func listenForConnections(addr string) {
tcpaddr, _ := net.ResolveTCPAddr("tcp", addr)
server, err := net.ListenTCP("tcp", tcpaddr)
checkError(err)
for {
conn, err := server.AcceptTCP()
checkError(err)
go handleRequests(conn)
}
}
// --------------------------------------------------------------------------
// Where the magic happens
func main() {
kvmap = make(map[string]data)
addr := os.Args[1]
listenForConnections(addr)
}