105 lines
2.4 KiB
Go
105 lines
2.4 KiB
Go
package senddata
|
|
|
|
import (
|
|
"net/http"
|
|
|
|
"gitlab.com/gitlab-org/gitlab/workhorse/internal/headers"
|
|
"gitlab.com/gitlab-org/gitlab/workhorse/internal/helper"
|
|
"gitlab.com/gitlab-org/gitlab/workhorse/internal/senddata/contentprocessor"
|
|
|
|
"github.com/prometheus/client_golang/prometheus"
|
|
"github.com/prometheus/client_golang/prometheus/promauto"
|
|
)
|
|
|
|
var (
|
|
sendDataResponses = promauto.NewCounterVec(
|
|
prometheus.CounterOpts{
|
|
Name: "gitlab_workhorse_senddata_responses",
|
|
Help: "How many HTTP responses have been hijacked by a workhorse senddata injecter",
|
|
},
|
|
[]string{"injecter"},
|
|
)
|
|
sendDataResponseBytes = promauto.NewCounterVec(
|
|
prometheus.CounterOpts{
|
|
Name: "gitlab_workhorse_senddata_response_bytes",
|
|
Help: "How many bytes have been written by workhorse senddata response injecters",
|
|
},
|
|
[]string{"injecter"},
|
|
)
|
|
)
|
|
|
|
type sendDataResponseWriter struct {
|
|
rw http.ResponseWriter
|
|
status int
|
|
hijacked bool
|
|
req *http.Request
|
|
injecters []Injecter
|
|
}
|
|
|
|
func SendData(h http.Handler, injecters ...Injecter) http.Handler {
|
|
return contentprocessor.SetContentHeaders(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
s := sendDataResponseWriter{
|
|
rw: w,
|
|
req: r,
|
|
injecters: injecters,
|
|
}
|
|
defer s.flush()
|
|
h.ServeHTTP(&s, r)
|
|
}))
|
|
}
|
|
|
|
func (s *sendDataResponseWriter) Header() http.Header {
|
|
return s.rw.Header()
|
|
}
|
|
|
|
func (s *sendDataResponseWriter) Write(data []byte) (int, error) {
|
|
if s.status == 0 {
|
|
s.WriteHeader(http.StatusOK)
|
|
}
|
|
if s.hijacked {
|
|
return len(data), nil
|
|
}
|
|
return s.rw.Write(data)
|
|
}
|
|
|
|
func (s *sendDataResponseWriter) WriteHeader(status int) {
|
|
if s.status != 0 {
|
|
return
|
|
}
|
|
s.status = status
|
|
|
|
if s.status == http.StatusOK && s.tryInject() {
|
|
return
|
|
}
|
|
|
|
s.rw.WriteHeader(s.status)
|
|
}
|
|
|
|
func (s *sendDataResponseWriter) tryInject() bool {
|
|
if s.hijacked {
|
|
return false
|
|
}
|
|
|
|
header := s.Header().Get(headers.GitlabWorkhorseSendDataHeader)
|
|
if header == "" {
|
|
return false
|
|
}
|
|
|
|
for _, injecter := range s.injecters {
|
|
if injecter.Match(header) {
|
|
s.hijacked = true
|
|
helper.DisableResponseBuffering(s.rw)
|
|
crw := helper.NewCountingResponseWriter(s.rw)
|
|
injecter.Inject(crw, s.req, header)
|
|
sendDataResponses.WithLabelValues(injecter.Name()).Inc()
|
|
sendDataResponseBytes.WithLabelValues(injecter.Name()).Add(float64(crw.Count()))
|
|
return true
|
|
}
|
|
}
|
|
|
|
return false
|
|
}
|
|
|
|
func (s *sendDataResponseWriter) flush() {
|
|
s.WriteHeader(http.StatusOK)
|
|
}
|