func handleSpecCommand(cmd string, keys [][]byte, timeout int) ([]byte, bool, bool, error) {
var b []byte
shouldClose := false
switch cmd {
case "PING":
b = []byte("+PONG\r\n")
case "QUIT":
b = OK_BYTES
shouldClose = true
case "SELECT":
b = OK_BYTES
case "ECHO":
if len(keys) > 0 {
var err error
b, err = respcoding.Marshal(string(keys[0]))
if err != nil {
return nil, true, false, errors.Trace(err)
}
} else {
return nil, true, false, nil
}
}
if len(b) > 0 {
return b, shouldClose, true, nil
}
return b, shouldClose, false, nil
}
func (oper *MultiOperator) mgetResults(mop *MulOp) ([]byte, error) {
slotmap := getSlotMap(mop.keys)
results := make([]interface{}, len(mop.keys))
conn := oper.pool.Get()
defer conn.Close()
for _, vec := range slotmap {
req := make([]interface{}, 0, len(vec))
for _, p := range vec {
req = append(req, p.key)
}
replys, err := redis.Values(conn.Do("mget", req...))
if err != nil {
return nil, errors.Trace(err)
}
for i, reply := range replys {
if reply != nil {
results[vec[i].pos] = reply
} else {
results[vec[i].pos] = nil
}
}
}
b, err := respcoding.Marshal(results)
return b, errors.Trace(err)
}
func (oper *MultiOperator) delResults(mop *MulOp) ([]byte, error) {
var results int64
conn := oper.pool.Get()
defer conn.Close()
for _, k := range mop.keys {
n, err := conn.Do("del", k)
if err != nil {
return nil, errors.Trace(err)
}
results += n.(int64)
}
b, err := respcoding.Marshal(int(results))
return b, errors.Trace(err)
}
func (oper *MultiOperator) mgetResults(mop *MulOp) ([]byte, error) {
results := make([]interface{}, len(mop.keys))
conn := oper.pool.Get()
defer conn.Close()
for i, key := range mop.keys {
replys, err := redis.Values(conn.Do("mget", key))
if err != nil {
return nil, errors.Trace(err)
}
for _, reply := range replys {
if reply != nil {
results[i] = reply
} else {
results[i] = nil
}
}
}
b, err := respcoding.Marshal(results)
return b, errors.Trace(err)
}
func filter(opstr string, keys [][]byte, c *session, timeoutSec int) (rawresp []byte, next bool, err error) {
if !allowOp(opstr) {
errmsg, err := respcoding.Marshal(fmt.Errorf("%s not allowed", opstr))
if err != nil {
log.Fatal("should never happend", opstr)
}
return errmsg, false, errors.New(string(errmsg))
}
buf, shouldClose, handled, err := handleSpecCommand(opstr, keys, timeoutSec)
if shouldClose { //quit command
return buf, false, errors.Trace(io.EOF)
}
if err != nil {
return nil, false, errors.Trace(err)
}
if handled {
return buf, false, nil
}
return nil, true, nil
}
func handleSpecCommand(cmd string, clientWriter DeadlineReadWriter, keys [][]byte) (bool, bool, error) {
var b []byte
shouldClose := false
switch cmd {
case "PING":
b = []byte("+PONG\r\n")
case "QUIT":
b = OK_BYTES
shouldClose = true
case "SELECT":
b = OK_BYTES
case "AUTH":
b = OK_BYTES
case "ECHO":
if len(keys) > 0 {
var err error
b, err = respcoding.Marshal(string(keys[0]))
if err != nil {
return true, false, errors.Trace(err)
}
} else {
return true, false, nil
}
}
if len(b) > 0 {
clientWriter.SetWriteDeadline(time.Now().Add(5 * time.Second))
_, err := clientWriter.Write(b)
if err != nil {
return shouldClose, true, errors.Errorf("%s, cmd:%s", err.Error(), cmd)
}
return shouldClose, true, nil
}
return shouldClose, false, nil
}
func Parse(r *bufio.Reader) (*Resp, error) {
line, err := readLine(r)
if err != nil {
return nil, errors.Trace(err)
}
resp := &Resp{Raw: line}
switch line[0] {
case '-':
resp.Type = ErrorResp
return resp, nil
case '+':
resp.Type = SimpleString
return resp, nil
case ':':
resp.Type = IntegerResp
return resp, nil
case '$':
resp.Type = BulkResp
size, err := Btoi(line[1 : len(line)-2])
if err != nil {
return nil, errors.Trace(err)
}
err = ReadBulk(r, size, &resp.Raw)
if err != nil {
return nil, errors.Trace(err)
}
return resp, nil
case '*':
i, err := Btoi(line[1 : len(line)-2]) //strip \r\n
if err != nil {
return nil, errors.Trace(err)
}
resp.Type = MultiResp
if i >= 0 {
multi := make([]*Resp, i)
for j := 0; j < i; j++ {
rp, err := Parse(r)
if err != nil {
return nil, errors.Trace(err)
}
multi[j] = rp
}
resp.Multi = multi
}
return resp, nil
default:
if !IsLetter(line[0]) { //handle telnet text command
return nil, errors.New("redis protocol error, " + string(line))
}
resp.Type = MultiResp
strs := strings.Split(string(line), " ")
resp.Raw = make([]byte, 0, 20)
resp.Raw = append(resp.Raw, '*')
resp.Raw = append(resp.Raw, []byte(strconv.Itoa(len(strs)))...)
resp.Raw = append(resp.Raw, NEW_LINE...)
for i := 0; i < len(strs); i++ { //last element is \r\n
if str := strings.TrimSpace(strs[i]); len(str) > 0 {
b, err := respcoding.Marshal(str)
if err != nil {
return nil, errors.New("redis protocol error, " + string(line))
}
resp.Multi = append(resp.Multi, &Resp{Type: BulkResp, Raw: b})
}
}
return resp, nil
}
}