-
Notifications
You must be signed in to change notification settings - Fork 48
/
payload.go
116 lines (95 loc) · 2.56 KB
/
payload.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
package telemetry
import (
"bytes"
"compress/gzip"
"encoding/base64"
"encoding/json"
"errors"
"fmt"
"io"
"strings"
)
func parsePayload(data []byte) (metadata, uncompressedData map[string]json.RawMessage, err error) {
var arr [4]json.RawMessage
if err = json.Unmarshal(data, &arr); nil != err {
err = fmt.Errorf("unable to unmarshal payload data array: %v", err)
return
}
var dataJSON []byte
compressed := strings.Trim(string(arr[3]), `"`)
if dataJSON, err = decodeUncompress(compressed); nil != err {
err = fmt.Errorf("unable to uncompress payload: %v", err)
return
}
if err = json.Unmarshal(dataJSON, &uncompressedData); nil != err {
err = fmt.Errorf("unable to unmarshal uncompressed payload: %v", err)
return
}
if err = json.Unmarshal(arr[2], &metadata); nil != err {
err = fmt.Errorf("unable to unmarshal payload metadata: %v", err)
return
}
return
}
func decodeUncompress(input string) ([]byte, error) {
decoded, err := base64.StdEncoding.DecodeString(input)
if nil != err {
return nil, err
}
buf := bytes.NewBuffer(decoded)
gz, err := gzip.NewReader(buf)
if nil != err {
return nil, err
}
var out bytes.Buffer
io.Copy(&out, gz)
gz.Close()
return out.Bytes(), nil
}
// ExtracTraceID extracts the trace ID within a payload, if present
func ExtractTraceID(data []byte) (string, error) {
if !bytes.Contains(data, []byte("NR_LAMBDA_MONITORING")) {
return "", nil
}
_, segments, err := parsePayload(data)
if err != nil {
return "", err
}
analyticEvents, ok := segments["analytic_event_data"]
if ok {
var parsedAnalyticEvents []json.RawMessage
if err := json.Unmarshal(analyticEvents, &parsedAnalyticEvents); err != nil {
return "", err
}
if len(parsedAnalyticEvents) == 3 {
var analyticEvent [][]struct {
TraceID string `json:"traceId"`
}
if err := json.Unmarshal(parsedAnalyticEvents[2], &analyticEvent); err != nil {
return "", err
}
if len(analyticEvent) > 0 && len(analyticEvent[0]) > 0 {
return analyticEvent[0][0].TraceID, nil
}
}
}
spanEvents, ok := segments["span_event_data"]
if ok {
var parsedSpanEvents []json.RawMessage
if err := json.Unmarshal(spanEvents, &parsedSpanEvents); err != nil {
return "", err
}
if len(parsedSpanEvents) == 3 {
var spanEvent [][]struct {
TraceID string `json:"traceId"`
}
if err := json.Unmarshal(parsedSpanEvents[2], &spanEvent); err != nil {
return "", err
}
if len(spanEvent) > 0 && len(spanEvent[0]) > 0 {
return spanEvent[0][0].TraceID, nil
}
}
}
return "", errors.New("No trace ID found in payload")
}