debian-mirror-gitlab/workhorse/internal/senddata/senddata.go
2021-02-22 17:27:13 +05:30

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)
}