Merge pull request #3 from lethain/extrafields

Extrafields

README.md

@@ -68,8 +68,6 @@ server is reachable or not.
 To Do
 -----
 
-- Custom fields
-- Tests for chunked messages
 - WriteMessage example
 
 License

gelf/writer.go

@@ -45,20 +45,21 @@ const (
 
 // Message represents the contents of the GELF message. It is gzipped
 // before sending.
-//
-// TODO: support optional ('_'-prefixed) fields?
 type Message struct {
- Version string `json:"version"`
- Host string `json:"host"`
- Short string `json:"short_message"`
- Full string `json:"full_message"`
- TimeUnix int64 `json:"timestamp"`
- Level int32 `json:"level"`
- Facility string `json:"facility"`
- File string `json:"file"`
- Line int `json:"line"`
+ Version string `json:"version"`
+ Host string `json:"host"`
+ Short string `json:"short_message"`
+ Full string `json:"full_message"`
+ TimeUnix int64 `json:"timestamp"`
+ Level int32 `json:"level"`
+ Facility string `json:"facility"`
+ File string `json:"file"`
+ Line int `json:"line"`
+ Extra map[string]interface{} `json:"-"`
 }
 
+type innerMessage Message //against circular (Un)MarshalJSON
+
 // Used to control GELF chunking. Should be less than (MTU - len(UDP
 // header)).
 //
@@ -158,9 +159,9 @@ func (w *Writer) writeChunked(zBytes []byte) (err error) {
 
  bytesLeft -= chunkLen
  }
- // debugging
+
  if bytesLeft != 0 {
- panic("aw shit")
+ return fmt.Errorf("error: %d bytes left after sending", bytesLeft)
  }
  return nil
 }
@@ -288,6 +289,7 @@ func (w *Writer) Write(p []byte) (n int, err error) {
  Facility: w.Facility,
  File: file,
  Line: line,
+ Extra: map[string]interface{}{},
  }
 
  if err = w.WriteMessage(&m); err != nil {
@@ -296,3 +298,64 @@ func (w *Writer) Write(p []byte) (n int, err error) {
 
  return len(p), nil
 }
+
+func (m *Message) MarshalJSON() ([]byte, error) {
+ var err error
+ var b, eb []byte
+
+ extra := m.Extra
+ b, err = json.Marshal((*innerMessage)(m))
+ m.Extra = extra
+ if err != nil {
+ return nil, err
+ }
+
+ if len(extra) == 0 {
+ return b, nil
+ }
+
+ if eb, err = json.Marshal(extra); err != nil {
+ return nil, err
+ }
+
+ // merge serialized message + serialized extra map
+ b[len(b)-1] = ','
+ return append(b, eb[1:len(eb)]...), nil
+}
+
+func (m *Message) UnmarshalJSON(data []byte) error {
+ i := make(map[string]interface{}, 16)
+ if err := json.Unmarshal(data, &i); err != nil {
+ return err
+ }
+ for k, v := range i {
+ if k[0] == '_' {
+ if m.Extra == nil {
+ m.Extra = make(map[string]interface{}, 1)
+ }
+ m.Extra[k] = v
+ continue
+ }
+ switch k {
+ case "version":
+ m.Version = v.(string)
+ case "host":
+ m.Host = v.(string)
+ case "short_message":
+ m.Short = v.(string)
+ case "full_message":
+ m.Full = v.(string)
+ case "timestamp":
+ m.TimeUnix = int64(v.(float64))
+ case "level":
+ m.Level = int32(v.(float64))
+ case "facility":
+ m.Facility = v.(string)
+ case "file":
+ m.File = v.(string)
+ case "line":
+ m.Line = int(v.(float64))
+ }
+ }
+ return nil
+}

gelf/writer_test.go

@@ -10,6 +10,7 @@ import (
  "fmt"
  "strings"
  "testing"
+ "time"
 )
 
 func TestNewWriter(t *testing.T) {
@@ -39,6 +40,25 @@ func sendAndRecv(msgData string, compress CompressType) (*Message, error) {
  return r.ReadMessage()
 }
 
+func sendAndRecvMsg(msg *Message, compress CompressType) (*Message, error) {
+ r, err := NewReader("127.0.0.1:0")
+ if err != nil {
+ return nil, fmt.Errorf("NewReader: %s", err)
+ }
+
+ w, err := NewWriter(r.Addr())
+ if err != nil {
+ return nil, fmt.Errorf("NewWriter: %s", err)
+ }
+ w.CompressionType = compress
+
+ if err = w.WriteMessage(msg); err != nil {
+ return nil, fmt.Errorf("w.Write: %s", err)
+ }
+
+ return r.ReadMessage()
+}
+
 // tests single-message (non-chunked) messages that are split over
 // multiple lines
 func TestWriteSmallMultiLine(t *testing.T) {
@@ -96,6 +116,12 @@ func TestWriteSmallOneLine(t *testing.T) {
  if !strings.HasSuffix(msg.File, fileExpected) {
  t.Errorf("msg.File: expected %s, got %s", fileExpected,
  msg.File)
+ return
+ }
+
+ if len(msg.Extra) != 0 {
+ t.Errorf("extra extra fields in %v (expect empty)", msg.Extra)
+ return
  }
 }
 
@@ -106,9 +132,14 @@ func TestGetCaller(t *testing.T) {
  return
  }
 
+ file, _ = getCaller(0)
+ if !strings.HasSuffix(file, "/gelf/writer_test.go") {
+ t.Errorf("not writer_test.go 1? %s", file)
+ }
+
  file, _ = getCallerIgnoringLogMulti(0)
  if !strings.HasSuffix(file, "/gelf/writer_test.go") {
- t.Errorf("not writer_test.go? %s", file)
+ t.Errorf("not writer_test.go 2? %s", file)
  }
 }
 
@@ -146,4 +177,58 @@ func TestWriteBigChunked(t *testing.T) {
  }
 }
 
-// tests single-message (non-chunked) messages that are a single line long
+// tests messages with extra data
+func TestExtraData(t *testing.T) {
+
+ // time.Now().Unix() seems fine, UnixNano() won't roundtrip
+ // through string -> float64 -> int64
+ extra := map[string]interface{}{"_a": 10 * time.Now().Unix(), "C": 9}
+
+ short := "quick"
+ full := short + "\nwith more detail"
+ m := Message{
+ Version: "1.0",
+ Host: "fake-host",
+ Short: string(short),
+ Full: string(full),
+ TimeUnix: time.Now().Unix(),
+ Level: 6, // info
+ Facility: "writer_test",
+ File: "writer_test.go",
+ Line: 186,
+ Extra: extra,
+ }
+
+ for _, i := range []CompressType{CompressGzip, CompressZlib} {
+ msg, err := sendAndRecvMsg(&m, i)
+ if err != nil {
+ t.Errorf("sendAndRecv: %s", err)
+ return
+ }
+
+ if msg.Short != short {
+ t.Errorf("msg.Short: expected %s, got %s", short, msg.Full)
+ return
+ }
+
+ if msg.Full != full {
+ t.Errorf("msg.Full: expected %s, got %s", full, msg.Full)
+ return
+ }
+
+ if msg.File != "writer_test.go" {
+ t.Errorf("msg.File: expected writer_test.go, got %s", msg.File)
+ return
+ }
+
+ if len(msg.Extra) != 1 {
+ t.Errorf("extra extra fields in %v", msg.Extra)
+ return
+ }
+
+ if int64(msg.Extra["_a"].(float64)) != extra["_a"].(int64) {
+ t.Errorf("_a didn't roundtrip (%v != %v)", int64(msg.Extra["_a"].(float64)), extra["_a"].(int64))
+ return
+ }
+ }
+}