summaryrefslogtreecommitdiff
path: root/workhorse/internal/senddata/senddata.go
blob: c287d2574fa72fd24fbd89c2bdab8c6116ff7b52 (plain)
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
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)
}