...
1 package proto
2
3 import (
4 "bufio"
5 "encoding"
6 "fmt"
7 "io"
8 "strconv"
9
10 "github.com/go-redis/redis/internal/util"
11 )
12
13 type Writer struct {
14 wr *bufio.Writer
15
16 lenBuf []byte
17 numBuf []byte
18 }
19
20 func NewWriter(wr io.Writer) *Writer {
21 return &Writer{
22 wr: bufio.NewWriter(wr),
23
24 lenBuf: make([]byte, 64),
25 numBuf: make([]byte, 64),
26 }
27 }
28
29 func (w *Writer) WriteArgs(args []interface{}) error {
30 err := w.wr.WriteByte(ArrayReply)
31 if err != nil {
32 return err
33 }
34
35 err = w.writeLen(len(args))
36 if err != nil {
37 return err
38 }
39
40 for _, arg := range args {
41 err := w.writeArg(arg)
42 if err != nil {
43 return err
44 }
45 }
46
47 return nil
48 }
49
50 func (w *Writer) writeLen(n int) error {
51 w.lenBuf = strconv.AppendUint(w.lenBuf[:0], uint64(n), 10)
52 w.lenBuf = append(w.lenBuf, '\r', '\n')
53 _, err := w.wr.Write(w.lenBuf)
54 return err
55 }
56
57 func (w *Writer) writeArg(v interface{}) error {
58 switch v := v.(type) {
59 case nil:
60 return w.string("")
61 case string:
62 return w.string(v)
63 case []byte:
64 return w.bytes(v)
65 case int:
66 return w.int(int64(v))
67 case int8:
68 return w.int(int64(v))
69 case int16:
70 return w.int(int64(v))
71 case int32:
72 return w.int(int64(v))
73 case int64:
74 return w.int(v)
75 case uint:
76 return w.uint(uint64(v))
77 case uint8:
78 return w.uint(uint64(v))
79 case uint16:
80 return w.uint(uint64(v))
81 case uint32:
82 return w.uint(uint64(v))
83 case uint64:
84 return w.uint(v)
85 case float32:
86 return w.float(float64(v))
87 case float64:
88 return w.float(v)
89 case bool:
90 if v {
91 return w.int(1)
92 } else {
93 return w.int(0)
94 }
95 case encoding.BinaryMarshaler:
96 b, err := v.MarshalBinary()
97 if err != nil {
98 return err
99 }
100 return w.bytes(b)
101 default:
102 return fmt.Errorf(
103 "redis: can't marshal %T (implement encoding.BinaryMarshaler)", v)
104 }
105 }
106
107 func (w *Writer) bytes(b []byte) error {
108 err := w.wr.WriteByte(StringReply)
109 if err != nil {
110 return err
111 }
112
113 err = w.writeLen(len(b))
114 if err != nil {
115 return err
116 }
117
118 _, err = w.wr.Write(b)
119 if err != nil {
120 return err
121 }
122
123 return w.crlf()
124 }
125
126 func (w *Writer) string(s string) error {
127 return w.bytes(util.StringToBytes(s))
128 }
129
130 func (w *Writer) uint(n uint64) error {
131 w.numBuf = strconv.AppendUint(w.numBuf[:0], n, 10)
132 return w.bytes(w.numBuf)
133 }
134
135 func (w *Writer) int(n int64) error {
136 w.numBuf = strconv.AppendInt(w.numBuf[:0], n, 10)
137 return w.bytes(w.numBuf)
138 }
139
140 func (w *Writer) float(f float64) error {
141 w.numBuf = strconv.AppendFloat(w.numBuf[:0], f, 'f', -1, 64)
142 return w.bytes(w.numBuf)
143 }
144
145 func (w *Writer) crlf() error {
146 err := w.wr.WriteByte('\r')
147 if err != nil {
148 return err
149 }
150 return w.wr.WriteByte('\n')
151 }
152
153 func (w *Writer) Reset(wr io.Writer) {
154 w.wr.Reset(wr)
155 }
156
157 func (w *Writer) Flush() error {
158 return w.wr.Flush()
159 }
160
View as plain text