1package trace
2
3import (
4	"encoding/json"
5	"net"
6	"time"
7
8	"github.com/jonboulle/clockwork"
9	log "github.com/sirupsen/logrus"
10)
11
12const (
13	// UDPDefaultAddr is a default address to emit logs to
14	UDPDefaultAddr = "127.0.0.1:5000"
15	// UDPDefaultNet is a default network
16	UDPDefaultNet = "udp"
17)
18
19// UDPOptionSetter represents functional arguments passed to ELKHook
20type UDPOptionSetter func(f *UDPHook)
21
22// NewUDPHook returns logrus-compatible hook that sends data to UDP socket
23func NewUDPHook(opts ...UDPOptionSetter) (*UDPHook, error) {
24	f := &UDPHook{}
25	for _, o := range opts {
26		o(f)
27	}
28	if f.Clock == nil {
29		f.Clock = clockwork.NewRealClock()
30	}
31	if f.clientNet == "" {
32		f.clientNet = UDPDefaultNet
33	}
34	if f.clientAddr == "" {
35		f.clientAddr = UDPDefaultAddr
36	}
37	addr, err := net.ResolveUDPAddr(f.clientNet, f.clientAddr)
38	if err != nil {
39		return nil, Wrap(err)
40	}
41	conn, err := net.ListenPacket("udp", ":0")
42	if err != nil {
43		return nil, Wrap(err)
44	}
45	f.addr = addr
46	f.conn = conn.(*net.UDPConn)
47	return f, nil
48}
49
50type UDPHook struct {
51	Clock      clockwork.Clock
52	clientNet  string
53	clientAddr string
54	addr       *net.UDPAddr
55	conn       *net.UDPConn
56}
57
58type Frame struct {
59	Time    time.Time              `json:"time"`
60	Type    string                 `json:"type"`
61	Entry   map[string]interface{} `json:"entry"`
62	Message string                 `json:"message"`
63	Level   string                 `json:"level"`
64}
65
66// Fire fires the event to the ELK beat
67func (elk *UDPHook) Fire(e *log.Entry) error {
68	// Make a copy to safely modify
69	entry := e.WithFields(nil)
70	if cursor := findFrame(); cursor != nil {
71		t := newTraceFromFrames(*cursor, nil)
72		entry.Data[FileField] = t.String()
73		entry.Data[FunctionField] = t.Func()
74	}
75	data, err := json.Marshal(Frame{
76		Time:    elk.Clock.Now().UTC(),
77		Type:    "trace",
78		Entry:   entry.Data,
79		Message: entry.Message,
80		Level:   entry.Level.String(),
81	})
82	if err != nil {
83		return Wrap(err)
84	}
85
86	conn, err := net.ListenPacket("udp", ":0")
87	if err != nil {
88		return Wrap(err)
89	}
90	defer conn.Close()
91
92	resolvedAddr, err := net.ResolveUDPAddr("udp", "127.0.0.1:5000")
93	if err != nil {
94		return Wrap(err)
95	}
96
97	_, err = (conn.(*net.UDPConn)).WriteToUDP(data, resolvedAddr)
98	return Wrap(err)
99
100}
101
102// Levels returns logging levels supported by logrus
103func (elk *UDPHook) Levels() []log.Level {
104	return []log.Level{
105		log.PanicLevel,
106		log.FatalLevel,
107		log.ErrorLevel,
108		log.WarnLevel,
109		log.InfoLevel,
110		log.DebugLevel,
111	}
112}
113