1package gomemcached
2
3import (
4	"encoding/binary"
5	"fmt"
6	"io"
7)
8
9// The maximum reasonable body length to expect.
10// Anything larger than this will result in an error.
11// The current limit, 20MB, is the size limit supported by ep-engine.
12var MaxBodyLen = int(20 * 1024 * 1024)
13
14// MCRequest is memcached Request
15type MCRequest struct {
16	// The command being issued
17	Opcode CommandCode
18	// The CAS (if applicable, or 0)
19	Cas uint64
20	// An opaque value to be returned with this request
21	Opaque uint32
22	// The vbucket to which this command belongs
23	VBucket uint16
24	// Command extras, key, and body
25	Extras, Key, Body, ExtMeta []byte
26	// Datatype identifier
27	DataType uint8
28}
29
30// Size gives the number of bytes this request requires.
31func (req *MCRequest) Size() int {
32	return HDR_LEN + len(req.Extras) + len(req.Key) + len(req.Body) + len(req.ExtMeta)
33}
34
35// A debugging string representation of this request
36func (req MCRequest) String() string {
37	return fmt.Sprintf("{MCRequest opcode=%s, bodylen=%d, key='%s'}",
38		req.Opcode, len(req.Body), req.Key)
39}
40
41func (req *MCRequest) fillHeaderBytes(data []byte) int {
42
43	pos := 0
44	data[pos] = REQ_MAGIC
45	pos++
46	data[pos] = byte(req.Opcode)
47	pos++
48	binary.BigEndian.PutUint16(data[pos:pos+2],
49		uint16(len(req.Key)))
50	pos += 2
51
52	// 4
53	data[pos] = byte(len(req.Extras))
54	pos++
55	// Data type
56	if req.DataType != 0 {
57		data[pos] = byte(req.DataType)
58	}
59	pos++
60	binary.BigEndian.PutUint16(data[pos:pos+2], req.VBucket)
61	pos += 2
62
63	// 8
64	binary.BigEndian.PutUint32(data[pos:pos+4],
65		uint32(len(req.Body)+len(req.Key)+len(req.Extras)+len(req.ExtMeta)))
66	pos += 4
67
68	// 12
69	binary.BigEndian.PutUint32(data[pos:pos+4], req.Opaque)
70	pos += 4
71
72	// 16
73	if req.Cas != 0 {
74		binary.BigEndian.PutUint64(data[pos:pos+8], req.Cas)
75	}
76	pos += 8
77
78	if len(req.Extras) > 0 {
79		copy(data[pos:pos+len(req.Extras)], req.Extras)
80		pos += len(req.Extras)
81	}
82
83	if len(req.Key) > 0 {
84		copy(data[pos:pos+len(req.Key)], req.Key)
85		pos += len(req.Key)
86	}
87
88	return pos
89}
90
91// HeaderBytes will return the wire representation of the request header
92// (with the extras and key).
93func (req *MCRequest) HeaderBytes() []byte {
94	data := make([]byte, HDR_LEN+len(req.Extras)+len(req.Key))
95
96	req.fillHeaderBytes(data)
97
98	return data
99}
100
101// Bytes will return the wire representation of this request.
102func (req *MCRequest) Bytes() []byte {
103	data := make([]byte, req.Size())
104
105	pos := req.fillHeaderBytes(data)
106
107	if len(req.Body) > 0 {
108		copy(data[pos:pos+len(req.Body)], req.Body)
109	}
110
111	if len(req.ExtMeta) > 0 {
112		copy(data[pos+len(req.Body):pos+len(req.Body)+len(req.ExtMeta)], req.ExtMeta)
113	}
114
115	return data
116}
117
118// Transmit will send this request message across a writer.
119func (req *MCRequest) Transmit(w io.Writer) (n int, err error) {
120	if len(req.Body) < 128 {
121		n, err = w.Write(req.Bytes())
122	} else {
123		n, err = w.Write(req.HeaderBytes())
124		if err == nil {
125			m := 0
126			m, err = w.Write(req.Body)
127			n += m
128		}
129	}
130	return
131}
132
133// Receive will fill this MCRequest with the data from a reader.
134func (req *MCRequest) Receive(r io.Reader, hdrBytes []byte) (int, error) {
135	if len(hdrBytes) < HDR_LEN {
136		hdrBytes = []byte{
137			0, 0, 0, 0, 0, 0, 0, 0,
138			0, 0, 0, 0, 0, 0, 0, 0,
139			0, 0, 0, 0, 0, 0, 0, 0}
140	}
141	n, err := io.ReadFull(r, hdrBytes)
142	if err != nil {
143		return n, err
144	}
145
146	if hdrBytes[0] != RES_MAGIC && hdrBytes[0] != REQ_MAGIC {
147		return n, fmt.Errorf("bad magic: 0x%02x", hdrBytes[0])
148	}
149
150	klen := int(binary.BigEndian.Uint16(hdrBytes[2:]))
151	elen := int(hdrBytes[4])
152	// Data type at 5
153	req.DataType = uint8(hdrBytes[5])
154
155	req.Opcode = CommandCode(hdrBytes[1])
156	// Vbucket at 6:7
157	req.VBucket = binary.BigEndian.Uint16(hdrBytes[6:])
158	totalBodyLen := int(binary.BigEndian.Uint32(hdrBytes[8:]))
159
160	req.Opaque = binary.BigEndian.Uint32(hdrBytes[12:])
161	req.Cas = binary.BigEndian.Uint64(hdrBytes[16:])
162
163	if totalBodyLen > 0 {
164		buf := make([]byte, totalBodyLen)
165		m, err := io.ReadFull(r, buf)
166		n += m
167		if err == nil {
168			if req.Opcode >= TAP_MUTATION &&
169				req.Opcode <= TAP_CHECKPOINT_END &&
170				len(buf) > 1 {
171				// In these commands there is "engine private"
172				// data at the end of the extras.  The first 2
173				// bytes of extra data give its length.
174				elen += int(binary.BigEndian.Uint16(buf))
175			}
176
177			req.Extras = buf[0:elen]
178			req.Key = buf[elen : klen+elen]
179
180			// get the length of extended metadata
181			extMetaLen := 0
182			if elen > 29 {
183				extMetaLen = int(binary.BigEndian.Uint16(req.Extras[28:30]))
184			}
185
186			bodyLen := totalBodyLen - klen - elen - extMetaLen
187			if bodyLen > MaxBodyLen {
188				return n, fmt.Errorf("%d is too big (max %d)",
189					bodyLen, MaxBodyLen)
190			}
191
192			req.Body = buf[klen+elen : klen+elen+bodyLen]
193			req.ExtMeta = buf[klen+elen+bodyLen:]
194		}
195	}
196	return n, err
197}
198