cubefs

1
// Copyright 2022 The CubeFS Authors.
2
//
3
// Licensed under the Apache License, Version 2.0 (the "License");
4
// you may not use this file except in compliance with the License.
5
// You may obtain a copy of the License at
6
//
7
//     http://www.apache.org/licenses/LICENSE-2.0
8
//
9
// Unless required by applicable law or agreed to in writing, software
10
// distributed under the License is distributed on an "AS IS" BASIS,
11
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
12
// implied. See the License for the specific language governing
13
// permissions and limitations under the License.
14

15
package auditlog
16

17
import (
18
	"bufio"
19
	"io"
20
	"net"
21
	"net/http"
22
	"strconv"
23
	"time"
24

25
	"github.com/cubefs/cubefs/blobstore/common/rpc"
26
	"github.com/cubefs/cubefs/blobstore/common/trace"
27
)
28

29
type responseWriter struct {
30
	module     string
31
	statusCode int
32
	n          int
33
	bodyLimit  int
34
	// bodyWritten record how much bytes has been written to client
35
	bodyWritten int64
36
	// body hold some data buffer of response body, like json or form
37
	// audit log will record body buffer into log file
38
	body           []byte
39
	extra          http.Header // extra header
40
	span           trace.Span
41
	startTime      time.Time
42
	hasRecordCost  bool
43
	hasWroteHeader bool
44
	no2xxBody      bool
45

46
	http.ResponseWriter
47
}
48

49
func (w *responseWriter) Write(b []byte) (int, error) {
50
	if !w.hasWroteHeader {
51
		w.WriteHeader(http.StatusOK)
52
		w.hasWroteHeader = true
53
	}
54
	if !(w.statusCode/100 == 2 && w.no2xxBody) && w.n < w.bodyLimit {
55
		n := copy(w.body[w.n:], b)
56
		w.n += n
57
	}
58
	n, err := w.ResponseWriter.Write(b)
59
	w.bodyWritten += int64(n)
60
	return n, err
61
}
62

63
// ReadFrom implement io.ReaderFrom when io.Copy.
64
// Response with this function will not hold first body bytes in local buffer.
65
func (w *responseWriter) ReadFrom(src io.Reader) (n int64, err error) {
66
	if !w.hasWroteHeader {
67
		w.WriteHeader(http.StatusOK)
68
		w.hasWroteHeader = true
69
	}
70
	if rf, ok := w.ResponseWriter.(io.ReaderFrom); ok {
71
		n, err = rf.ReadFrom(src)
72
		w.bodyWritten += int64(n)
73
		return
74
	}
75
	n, err = io.Copy(w.ResponseWriter, src)
76
	w.bodyWritten += int64(n)
77
	return
78
}
79

80
func (w *responseWriter) WriteHeader(code int) {
81
	if w.hasWroteHeader {
82
		return
83
	}
84

85
	if !w.hasRecordCost {
86
		w.span.AppendTrackLog(w.module, w.startTime, nil)
87
		traceLog := w.span.TrackLog()
88
		for i := range traceLog {
89
			w.Header().Add(rpc.HeaderTraceLog, traceLog[i])
90
		}
91
		tags := w.span.Tags().ToSlice()
92
		for i := range tags {
93
			w.Header().Add(rpc.HeaderTraceTags, tags[i])
94
		}
95
		w.Header().Set(trace.GetTraceIDKey(), w.span.TraceID())
96
		w.hasRecordCost = true
97
	}
98
	w.statusCode = code
99
	w.ResponseWriter.WriteHeader(code)
100
	w.hasWroteHeader = true
101
}
102

103
func (w *responseWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
104
	return w.ResponseWriter.(http.Hijacker).Hijack()
105
}
106

107
func (w *responseWriter) Flush() {
108
	w.ResponseWriter.(http.Flusher).Flush()
109
}
110

111
func (w *responseWriter) ExtraHeader() http.Header {
112
	if w.extra == nil {
113
		w.extra = make(http.Header)
114
	}
115

116
	return w.extra
117
}
118

119
func (w *responseWriter) getBody() []byte {
120
	header := w.ResponseWriter.Header()
121
	length, _ := strconv.ParseInt(header.Get(rpc.HeaderContentLength), 10, 64)
122
	if (w.statusCode/100 == 2 && w.no2xxBody) || length > int64(w.n) {
123
		return nil
124
	}
125
	return w.body[:w.n]
126
}
127

128
func (w *responseWriter) getStatusCode() int {
129
	return w.statusCode
130
}
131

132
func (w *responseWriter) getHeader() M {
133
	header := w.ResponseWriter.Header()
134
	headerM := make(M)
135
	for k := range header {
136
		if k == rpc.HeaderTraceLog || k == rpc.HeaderTraceTags {
137
			headerM[k] = header[k]
138
		} else if len(header[k]) == 1 {
139
			headerM[k] = header[k][0]
140
		} else {
141
			headerM[k] = header[k]
142
		}
143
	}
144

145
	for k := range w.extra {
146
		if len(w.extra[k]) == 1 {
147
			headerM[k] = w.extra[k][0]
148
		} else {
149
			headerM[k] = w.extra[k]
150
		}
151
	}
152

153
	return headerM
154
}
155

156
func (w *responseWriter) getBodyWritten() int64 {
157
	return w.bodyWritten
158
}
159