2
0

encode.go 7.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322
  1. package logfmt
  2. import (
  3. "bytes"
  4. "encoding"
  5. "errors"
  6. "fmt"
  7. "io"
  8. "reflect"
  9. "strings"
  10. "unicode/utf8"
  11. )
  12. // MarshalKeyvals returns the logfmt encoding of keyvals, a variadic sequence
  13. // of alternating keys and values.
  14. func MarshalKeyvals(keyvals ...interface{}) ([]byte, error) {
  15. buf := &bytes.Buffer{}
  16. if err := NewEncoder(buf).EncodeKeyvals(keyvals...); err != nil {
  17. return nil, err
  18. }
  19. return buf.Bytes(), nil
  20. }
  21. // An Encoder writes logfmt data to an output stream.
  22. type Encoder struct {
  23. w io.Writer
  24. scratch bytes.Buffer
  25. needSep bool
  26. }
  27. // NewEncoder returns a new encoder that writes to w.
  28. func NewEncoder(w io.Writer) *Encoder {
  29. return &Encoder{
  30. w: w,
  31. }
  32. }
  33. var (
  34. space = []byte(" ")
  35. equals = []byte("=")
  36. newline = []byte("\n")
  37. null = []byte("null")
  38. )
  39. // EncodeKeyval writes the logfmt encoding of key and value to the stream. A
  40. // single space is written before the second and subsequent keys in a record.
  41. // Nothing is written if a non-nil error is returned.
  42. func (enc *Encoder) EncodeKeyval(key, value interface{}) error {
  43. enc.scratch.Reset()
  44. if enc.needSep {
  45. if _, err := enc.scratch.Write(space); err != nil {
  46. return err
  47. }
  48. }
  49. if err := writeKey(&enc.scratch, key); err != nil {
  50. return err
  51. }
  52. if _, err := enc.scratch.Write(equals); err != nil {
  53. return err
  54. }
  55. if err := writeValue(&enc.scratch, value); err != nil {
  56. return err
  57. }
  58. _, err := enc.w.Write(enc.scratch.Bytes())
  59. enc.needSep = true
  60. return err
  61. }
  62. // EncodeKeyvals writes the logfmt encoding of keyvals to the stream. Keyvals
  63. // is a variadic sequence of alternating keys and values. Keys of unsupported
  64. // type are skipped along with their corresponding value. Values of
  65. // unsupported type or that cause a MarshalerError are replaced by their error
  66. // but do not cause EncodeKeyvals to return an error. If a non-nil error is
  67. // returned some key/value pairs may not have be written.
  68. func (enc *Encoder) EncodeKeyvals(keyvals ...interface{}) error {
  69. if len(keyvals) == 0 {
  70. return nil
  71. }
  72. if len(keyvals)%2 == 1 {
  73. keyvals = append(keyvals, nil)
  74. }
  75. for i := 0; i < len(keyvals); i += 2 {
  76. k, v := keyvals[i], keyvals[i+1]
  77. err := enc.EncodeKeyval(k, v)
  78. if err == ErrUnsupportedKeyType {
  79. continue
  80. }
  81. if _, ok := err.(*MarshalerError); ok || err == ErrUnsupportedValueType {
  82. v = err
  83. err = enc.EncodeKeyval(k, v)
  84. }
  85. if err != nil {
  86. return err
  87. }
  88. }
  89. return nil
  90. }
  91. // MarshalerError represents an error encountered while marshaling a value.
  92. type MarshalerError struct {
  93. Type reflect.Type
  94. Err error
  95. }
  96. func (e *MarshalerError) Error() string {
  97. return "error marshaling value of type " + e.Type.String() + ": " + e.Err.Error()
  98. }
  99. // ErrNilKey is returned by Marshal functions and Encoder methods if a key is
  100. // a nil interface or pointer value.
  101. var ErrNilKey = errors.New("nil key")
  102. // ErrInvalidKey is returned by Marshal functions and Encoder methods if, after
  103. // dropping invalid runes, a key is empty.
  104. var ErrInvalidKey = errors.New("invalid key")
  105. // ErrUnsupportedKeyType is returned by Encoder methods if a key has an
  106. // unsupported type.
  107. var ErrUnsupportedKeyType = errors.New("unsupported key type")
  108. // ErrUnsupportedValueType is returned by Encoder methods if a value has an
  109. // unsupported type.
  110. var ErrUnsupportedValueType = errors.New("unsupported value type")
  111. func writeKey(w io.Writer, key interface{}) error {
  112. if key == nil {
  113. return ErrNilKey
  114. }
  115. switch k := key.(type) {
  116. case string:
  117. return writeStringKey(w, k)
  118. case []byte:
  119. if k == nil {
  120. return ErrNilKey
  121. }
  122. return writeBytesKey(w, k)
  123. case encoding.TextMarshaler:
  124. kb, err := safeMarshal(k)
  125. if err != nil {
  126. return err
  127. }
  128. if kb == nil {
  129. return ErrNilKey
  130. }
  131. return writeBytesKey(w, kb)
  132. case fmt.Stringer:
  133. ks, ok := safeString(k)
  134. if !ok {
  135. return ErrNilKey
  136. }
  137. return writeStringKey(w, ks)
  138. default:
  139. rkey := reflect.ValueOf(key)
  140. switch rkey.Kind() {
  141. case reflect.Array, reflect.Chan, reflect.Func, reflect.Map, reflect.Slice, reflect.Struct:
  142. return ErrUnsupportedKeyType
  143. case reflect.Ptr:
  144. if rkey.IsNil() {
  145. return ErrNilKey
  146. }
  147. return writeKey(w, rkey.Elem().Interface())
  148. }
  149. return writeStringKey(w, fmt.Sprint(k))
  150. }
  151. }
  152. // keyRuneFilter returns r for all valid key runes, and -1 for all invalid key
  153. // runes. When used as the mapping function for strings.Map and bytes.Map
  154. // functions it causes them to remove invalid key runes from strings or byte
  155. // slices respectively.
  156. func keyRuneFilter(r rune) rune {
  157. if r <= ' ' || r == '=' || r == '"' || r == utf8.RuneError {
  158. return -1
  159. }
  160. return r
  161. }
  162. func writeStringKey(w io.Writer, key string) error {
  163. k := strings.Map(keyRuneFilter, key)
  164. if k == "" {
  165. return ErrInvalidKey
  166. }
  167. _, err := io.WriteString(w, k)
  168. return err
  169. }
  170. func writeBytesKey(w io.Writer, key []byte) error {
  171. k := bytes.Map(keyRuneFilter, key)
  172. if len(k) == 0 {
  173. return ErrInvalidKey
  174. }
  175. _, err := w.Write(k)
  176. return err
  177. }
  178. func writeValue(w io.Writer, value interface{}) error {
  179. switch v := value.(type) {
  180. case nil:
  181. return writeBytesValue(w, null)
  182. case string:
  183. return writeStringValue(w, v, true)
  184. case []byte:
  185. return writeBytesValue(w, v)
  186. case encoding.TextMarshaler:
  187. vb, err := safeMarshal(v)
  188. if err != nil {
  189. return err
  190. }
  191. if vb == nil {
  192. vb = null
  193. }
  194. return writeBytesValue(w, vb)
  195. case error:
  196. se, ok := safeError(v)
  197. return writeStringValue(w, se, ok)
  198. case fmt.Stringer:
  199. ss, ok := safeString(v)
  200. return writeStringValue(w, ss, ok)
  201. default:
  202. rvalue := reflect.ValueOf(value)
  203. switch rvalue.Kind() {
  204. case reflect.Array, reflect.Chan, reflect.Func, reflect.Map, reflect.Slice, reflect.Struct:
  205. return ErrUnsupportedValueType
  206. case reflect.Ptr:
  207. if rvalue.IsNil() {
  208. return writeBytesValue(w, null)
  209. }
  210. return writeValue(w, rvalue.Elem().Interface())
  211. }
  212. return writeStringValue(w, fmt.Sprint(v), true)
  213. }
  214. }
  215. func needsQuotedValueRune(r rune) bool {
  216. return r <= ' ' || r == '=' || r == '"' || r == utf8.RuneError
  217. }
  218. func writeStringValue(w io.Writer, value string, ok bool) error {
  219. var err error
  220. if ok && value == "null" {
  221. _, err = io.WriteString(w, `"null"`)
  222. } else if strings.IndexFunc(value, needsQuotedValueRune) != -1 {
  223. _, err = writeQuotedString(w, value)
  224. } else {
  225. _, err = io.WriteString(w, value)
  226. }
  227. return err
  228. }
  229. func writeBytesValue(w io.Writer, value []byte) error {
  230. var err error
  231. if bytes.IndexFunc(value, needsQuotedValueRune) != -1 {
  232. _, err = writeQuotedBytes(w, value)
  233. } else {
  234. _, err = w.Write(value)
  235. }
  236. return err
  237. }
  238. // EndRecord writes a newline character to the stream and resets the encoder
  239. // to the beginning of a new record.
  240. func (enc *Encoder) EndRecord() error {
  241. _, err := enc.w.Write(newline)
  242. if err == nil {
  243. enc.needSep = false
  244. }
  245. return err
  246. }
  247. // Reset resets the encoder to the beginning of a new record.
  248. func (enc *Encoder) Reset() {
  249. enc.needSep = false
  250. }
  251. func safeError(err error) (s string, ok bool) {
  252. defer func() {
  253. if panicVal := recover(); panicVal != nil {
  254. if v := reflect.ValueOf(err); v.Kind() == reflect.Ptr && v.IsNil() {
  255. s, ok = "null", false
  256. } else {
  257. s, ok = fmt.Sprintf("PANIC:%v", panicVal), false
  258. }
  259. }
  260. }()
  261. s, ok = err.Error(), true
  262. return
  263. }
  264. func safeString(str fmt.Stringer) (s string, ok bool) {
  265. defer func() {
  266. if panicVal := recover(); panicVal != nil {
  267. if v := reflect.ValueOf(str); v.Kind() == reflect.Ptr && v.IsNil() {
  268. s, ok = "null", false
  269. } else {
  270. s, ok = fmt.Sprintf("PANIC:%v", panicVal), true
  271. }
  272. }
  273. }()
  274. s, ok = str.String(), true
  275. return
  276. }
  277. func safeMarshal(tm encoding.TextMarshaler) (b []byte, err error) {
  278. defer func() {
  279. if panicVal := recover(); panicVal != nil {
  280. if v := reflect.ValueOf(tm); v.Kind() == reflect.Ptr && v.IsNil() {
  281. b, err = nil, nil
  282. } else {
  283. b, err = nil, fmt.Errorf("panic when marshalling: %s", panicVal)
  284. }
  285. }
  286. }()
  287. b, err = tm.MarshalText()
  288. if err != nil {
  289. return nil, &MarshalerError{
  290. Type: reflect.TypeOf(tm),
  291. Err: err,
  292. }
  293. }
  294. return
  295. }