0
Fork 0
mirror of https://codeberg.org/forgejo/forgejo.git synced 2024-12-29 02:44:08 -05:00
forgejo/vendor/github.com/couchbase/gomemcached/mc_req.go

482 lines
14 KiB
Go

package gomemcached
import (
"encoding/binary"
"fmt"
"io"
)
// The maximum reasonable body length to expect.
// Anything larger than this will result in an error.
// The current limit, 20MB, is the size limit supported by ep-engine.
var MaxBodyLen = int(20 * 1024 * 1024)
// MCRequest is memcached Request
type MCRequest struct {
// The command being issued
Opcode CommandCode
// The CAS (if applicable, or 0)
Cas uint64
// An opaque value to be returned with this request
Opaque uint32
// The vbucket to which this command belongs
VBucket uint16
// Command extras, key, and body
Extras, Key, Body, ExtMeta []byte
// Datatype identifier
DataType uint8
// len() calls are expensive - cache this in case for collection
Keylen int
// Flexible Framing Extras
FramingExtras []FrameInfo
// Stored length of incoming framing extras
FramingElen int
}
// Size gives the number of bytes this request requires.
func (req *MCRequest) Size() int {
return HDR_LEN + len(req.Extras) + len(req.Key) + len(req.Body) + len(req.ExtMeta) + req.FramingElen
}
// A debugging string representation of this request
func (req MCRequest) String() string {
return fmt.Sprintf("{MCRequest opcode=%s, bodylen=%d, key='%s'}",
req.Opcode, len(req.Body), req.Key)
}
func (req *MCRequest) fillRegularHeaderBytes(data []byte) int {
// Byte/ 0 | 1 | 2 | 3 |
// / | | | |
// |0 1 2 3 4 5 6 7|0 1 2 3 4 5 6 7|0 1 2 3 4 5 6 7|0 1 2 3 4 5 6 7|
// +---------------+---------------+---------------+---------------+
// 0| Magic | Opcode | Key length |
// +---------------+---------------+---------------+---------------+
// 4| Extras length | Data type | vbucket id |
// +---------------+---------------+---------------+---------------+
// 8| Total body length |
// +---------------+---------------+---------------+---------------+
// 12| Opaque |
// +---------------+---------------+---------------+---------------+
// 16| CAS |
// | |
// +---------------+---------------+---------------+---------------+
// Total 24 bytes
pos := 0
data[pos] = REQ_MAGIC
pos++
data[pos] = byte(req.Opcode)
pos++
binary.BigEndian.PutUint16(data[pos:pos+2],
uint16(len(req.Key)))
pos += 2
// 4
data[pos] = byte(len(req.Extras))
pos++
// Data type
if req.DataType != 0 {
data[pos] = byte(req.DataType)
}
pos++
binary.BigEndian.PutUint16(data[pos:pos+2], req.VBucket)
pos += 2
// 8
binary.BigEndian.PutUint32(data[pos:pos+4],
uint32(len(req.Body)+len(req.Key)+len(req.Extras)+len(req.ExtMeta)))
pos += 4
// 12
binary.BigEndian.PutUint32(data[pos:pos+4], req.Opaque)
pos += 4
// 16
if req.Cas != 0 {
binary.BigEndian.PutUint64(data[pos:pos+8], req.Cas)
}
pos += 8
if len(req.Extras) > 0 {
copy(data[pos:pos+len(req.Extras)], req.Extras)
pos += len(req.Extras)
}
if len(req.Key) > 0 {
copy(data[pos:pos+len(req.Key)], req.Key)
pos += len(req.Key)
}
return pos
}
// Returns pos and if trailing by half byte
func (req *MCRequest) fillFlexHeaderBytes(data []byte) (int, bool) {
// Byte/ 0 | 1 | 2 | 3 |
// / | | | |
// |0 1 2 3 4 5 6 7|0 1 2 3 4 5 6 7|0 1 2 3 4 5 6 7|0 1 2 3 4 5 6 7|
// +---------------+---------------+---------------+---------------+
// 0| Magic (0x08) | Opcode | Framing extras| Key Length |
// +---------------+---------------+---------------+---------------+
// 4| Extras length | Data type | vbucket id |
// +---------------+---------------+---------------+---------------+
// 8| Total body length |
// +---------------+---------------+---------------+---------------+
// 12| Opaque |
// +---------------+---------------+---------------+---------------+
// 16| CAS |
// | |
// +---------------+---------------+---------------+---------------+
// Total 24 bytes
data[0] = FLEX_MAGIC
data[1] = byte(req.Opcode)
data[2] = byte(req.FramingElen)
data[3] = byte(req.Keylen)
elen := len(req.Extras)
data[4] = byte(elen)
if req.DataType != 0 {
data[5] = byte(req.DataType)
}
binary.BigEndian.PutUint16(data[6:8], req.VBucket)
binary.BigEndian.PutUint32(data[8:12],
uint32(len(req.Body)+req.Keylen+elen+len(req.ExtMeta)+req.FramingElen))
binary.BigEndian.PutUint32(data[12:16], req.Opaque)
if req.Cas != 0 {
binary.BigEndian.PutUint64(data[16:24], req.Cas)
}
pos := HDR_LEN
// Add framing infos
var framingExtras []byte
var outputBytes []byte
var mergeModeSrc []byte
var frameBytes int
var halfByteMode bool
var mergeMode bool
for _, frameInfo := range req.FramingExtras {
if !mergeMode {
outputBytes, halfByteMode = frameInfo.Bytes()
if !halfByteMode {
framingExtras = append(framingExtras, outputBytes...)
frameBytes += len(outputBytes)
} else {
mergeMode = true
mergeModeSrc = outputBytes
}
} else {
outputBytes, halfByteMode = frameInfo.Bytes()
outputBytes := ShiftByteSliceRight4Bits(outputBytes)
if halfByteMode {
// Previous halfbyte merge with this halfbyte will result in a complete byte
mergeMode = false
outputBytes = Merge2HalfByteSlices(mergeModeSrc, outputBytes)
framingExtras = append(framingExtras, outputBytes...)
frameBytes += len(outputBytes)
} else {
// Merge half byte with a non-half byte will result in a combined half-byte that will
// become the source for the next iteration
mergeModeSrc = Merge2HalfByteSlices(mergeModeSrc, outputBytes)
}
}
}
if mergeMode {
// Commit the temporary merge area into framingExtras
framingExtras = append(framingExtras, mergeModeSrc...)
frameBytes += len(mergeModeSrc)
}
copy(data[pos:pos+frameBytes], framingExtras)
pos += frameBytes
// Add Extras
if len(req.Extras) > 0 {
if mergeMode {
outputBytes = ShiftByteSliceRight4Bits(req.Extras)
data = Merge2HalfByteSlices(data, outputBytes)
} else {
copy(data[pos:pos+elen], req.Extras)
}
pos += elen
}
// Add keys
if req.Keylen > 0 {
if mergeMode {
outputBytes = ShiftByteSliceRight4Bits(req.Key)
data = Merge2HalfByteSlices(data, outputBytes)
} else {
copy(data[pos:pos+req.Keylen], req.Key)
}
pos += req.Keylen
}
return pos, mergeMode
}
func (req *MCRequest) FillHeaderBytes(data []byte) (int, bool) {
if req.FramingElen == 0 {
return req.fillRegularHeaderBytes(data), false
} else {
return req.fillFlexHeaderBytes(data)
}
}
// HeaderBytes will return the wire representation of the request header
// (with the extras and key).
func (req *MCRequest) HeaderBytes() []byte {
data := make([]byte, HDR_LEN+len(req.Extras)+len(req.Key)+req.FramingElen)
req.FillHeaderBytes(data)
return data
}
// Bytes will return the wire representation of this request.
func (req *MCRequest) Bytes() []byte {
data := make([]byte, req.Size())
pos, halfByteMode := req.FillHeaderBytes(data)
// TODO - the halfByteMode should be revisited for a more efficient
// way of doing things
if len(req.Body) > 0 {
if halfByteMode {
shifted := ShiftByteSliceRight4Bits(req.Body)
data = Merge2HalfByteSlices(data, shifted)
} else {
copy(data[pos:pos+len(req.Body)], req.Body)
}
}
if len(req.ExtMeta) > 0 {
if halfByteMode {
shifted := ShiftByteSliceRight4Bits(req.ExtMeta)
data = Merge2HalfByteSlices(data, shifted)
} else {
copy(data[pos+len(req.Body):pos+len(req.Body)+len(req.ExtMeta)], req.ExtMeta)
}
}
return data
}
// Transmit will send this request message across a writer.
func (req *MCRequest) Transmit(w io.Writer) (n int, err error) {
if len(req.Body) < 128 {
n, err = w.Write(req.Bytes())
} else {
n, err = w.Write(req.HeaderBytes())
if err == nil {
m := 0
m, err = w.Write(req.Body)
n += m
}
}
return
}
func (req *MCRequest) receiveHeaderCommon(hdrBytes []byte) (elen, totalBodyLen int) {
elen = int(hdrBytes[4])
// Data type at 5
req.DataType = uint8(hdrBytes[5])
req.Opcode = CommandCode(hdrBytes[1])
// Vbucket at 6:7
req.VBucket = binary.BigEndian.Uint16(hdrBytes[6:])
totalBodyLen = int(binary.BigEndian.Uint32(hdrBytes[8:]))
req.Opaque = binary.BigEndian.Uint32(hdrBytes[12:])
req.Cas = binary.BigEndian.Uint64(hdrBytes[16:])
return
}
func (req *MCRequest) receiveRegHeader(hdrBytes []byte) (elen, totalBodyLen int) {
elen, totalBodyLen = req.receiveHeaderCommon(hdrBytes)
req.Keylen = int(binary.BigEndian.Uint16(hdrBytes[2:]))
return
}
func (req *MCRequest) receiveFlexibleFramingHeader(hdrBytes []byte) (elen, totalBodyLen, framingElen int) {
elen, totalBodyLen = req.receiveHeaderCommon(hdrBytes)
// For flexible framing header, key length is a single byte at byte index 3
req.Keylen = int(binary.BigEndian.Uint16(hdrBytes[2:]) & 0x0ff)
// Flexible framing lengh is a single byte at index 2
framingElen = int(binary.BigEndian.Uint16(hdrBytes[2:]) >> 8)
req.FramingElen = framingElen
return
}
func (req *MCRequest) populateRegularBody(r io.Reader, totalBodyLen, elen int) (int, error) {
var m int
var err error
if totalBodyLen > 0 {
buf := make([]byte, totalBodyLen)
m, err = io.ReadFull(r, buf)
if err == nil {
if req.Opcode >= TAP_MUTATION &&
req.Opcode <= TAP_CHECKPOINT_END &&
len(buf) > 1 {
// In these commands there is "engine private"
// data at the end of the extras. The first 2
// bytes of extra data give its length.
elen += int(binary.BigEndian.Uint16(buf))
}
req.Extras = buf[0:elen]
req.Key = buf[elen : req.Keylen+elen]
// get the length of extended metadata
extMetaLen := 0
if elen > 29 {
extMetaLen = int(binary.BigEndian.Uint16(req.Extras[28:30]))
}
bodyLen := totalBodyLen - req.Keylen - elen - extMetaLen
if bodyLen > MaxBodyLen {
return m, fmt.Errorf("%d is too big (max %d)",
bodyLen, MaxBodyLen)
}
req.Body = buf[req.Keylen+elen : req.Keylen+elen+bodyLen]
req.ExtMeta = buf[req.Keylen+elen+bodyLen:]
}
}
return m, err
}
func (req *MCRequest) populateFlexBody(r io.Reader, totalBodyLen, elen, framingElen int) (int, error) {
var m int
var err error
if totalBodyLen > 0 {
buf := make([]byte, totalBodyLen)
m, err = io.ReadFull(r, buf)
if err != nil {
return m, err
}
err = req.populateFlexBodyInternal(buf, totalBodyLen, elen, framingElen)
}
return m, err
}
func (req *MCRequest) populateFlexBodyInternal(buf []byte, totalBodyLen, elen, framingElen int) error {
var halfByteOffset bool
var err error
if framingElen > 0 {
var objs []FrameInfo
objs, err, halfByteOffset = parseFrameInfoObjects(buf, framingElen)
if err != nil {
return err
}
req.FramingExtras = objs
}
err = req.populateFlexBodyAfterFrames(buf, totalBodyLen, elen, framingElen, halfByteOffset)
if err != nil {
return err
}
return nil
}
func (req *MCRequest) populateFlexBodyAfterFrames(buf []byte, totalBodyLen, elen, framingElen int, halfByteOffset bool) error {
var idxCursor int = framingElen
if req.Opcode >= TAP_MUTATION && req.Opcode <= TAP_CHECKPOINT_END && len(buf[idxCursor:]) > 1 {
// In these commands there is "engine private"
// data at the end of the extras. The first 2
// bytes of extra data give its length.
if !halfByteOffset {
elen += int(binary.BigEndian.Uint16(buf[idxCursor:]))
} else {
// 0 1 2 3 4 .... 19 20 21 22 ... 32
// ^-----^ ^-------^ ^------------^
// offset data do not care
var buffer uint32 = binary.BigEndian.Uint32(buf[idxCursor:])
buffer &= 0xffff000
elen += int(buffer >> 12)
}
}
// Get the extras
if !halfByteOffset {
req.Extras = buf[idxCursor : idxCursor+elen]
} else {
preShift := buf[idxCursor : idxCursor+elen+1]
req.Extras = ShiftByteSliceLeft4Bits(preShift)
}
idxCursor += elen
// Get the Key
if !halfByteOffset {
req.Key = buf[idxCursor : idxCursor+req.Keylen]
} else {
preShift := buf[idxCursor : idxCursor+req.Keylen+1]
req.Key = ShiftByteSliceLeft4Bits(preShift)
}
idxCursor += req.Keylen
// get the length of extended metadata
extMetaLen := 0
if elen > 29 {
extMetaLen = int(binary.BigEndian.Uint16(req.Extras[28:30]))
}
idxCursor += extMetaLen
bodyLen := totalBodyLen - req.Keylen - elen - extMetaLen - framingElen
if bodyLen > MaxBodyLen {
return fmt.Errorf("%d is too big (max %d)",
bodyLen, MaxBodyLen)
}
if !halfByteOffset {
req.Body = buf[idxCursor : idxCursor+bodyLen]
idxCursor += bodyLen
} else {
preShift := buf[idxCursor : idxCursor+bodyLen+1]
req.Body = ShiftByteSliceLeft4Bits(preShift)
idxCursor += bodyLen
}
if extMetaLen > 0 {
if !halfByteOffset {
req.ExtMeta = buf[idxCursor:]
} else {
preShift := buf[idxCursor:]
req.ExtMeta = ShiftByteSliceLeft4Bits(preShift)
}
}
return nil
}
// Receive will fill this MCRequest with the data from a reader.
func (req *MCRequest) Receive(r io.Reader, hdrBytes []byte) (int, error) {
if len(hdrBytes) < HDR_LEN {
hdrBytes = []byte{
0, 0, 0, 0, 0, 0, 0, 0,
0, 0, 0, 0, 0, 0, 0, 0,
0, 0, 0, 0, 0, 0, 0, 0}
}
n, err := io.ReadFull(r, hdrBytes)
if err != nil {
fmt.Printf("Err %v\n", err)
return n, err
}
switch hdrBytes[0] {
case RES_MAGIC:
fallthrough
case REQ_MAGIC:
elen, totalBodyLen := req.receiveRegHeader(hdrBytes)
bodyRead, err := req.populateRegularBody(r, totalBodyLen, elen)
return n + bodyRead, err
case FLEX_MAGIC:
elen, totalBodyLen, framingElen := req.receiveFlexibleFramingHeader(hdrBytes)
bodyRead, err := req.populateFlexBody(r, totalBodyLen, elen, framingElen)
return n + bodyRead, err
default:
return n, fmt.Errorf("bad magic: 0x%02x", hdrBytes[0])
}
}