encode.go

Documentation: github.com/go-logfmt/logfmt

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