Skip to content
This repository was archived by the owner on Nov 7, 2025. It is now read-only.

Commit c1a9b05

Browse files
pdelewskipdelewski
authored
Adding middleware handling API (#1130)
This PR adds a new API call for registering middlewares into the HTTP FrontendConnector.
1 parent 7b4de20 commit c1a9b05

File tree

5 files changed

+148
-16
lines changed

5 files changed

+148
-16
lines changed

‎quesma/frontend_connectors/basic_http_frontend_connector.go‎

Lines changed: 38 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,7 @@ type BasicHTTPFrontendConnector struct {
3030
phoneHomeClient diag.PhoneHomeClient
3131
debugInfoCollector diag.DebugInfoCollector
3232
logger quesma_api.QuesmaLogger
33+
middlewares []http.Handler
3334
}
3435

3536
func (h *BasicHTTPFrontendConnector) GetChildComponents() []interface{} {
@@ -66,6 +67,7 @@ func NewBasicHTTPFrontendConnector(endpoint string, config *config.QuesmaConfigu
6667
responseMutator: func(w http.ResponseWriter) http.ResponseWriter {
6768
return w
6869
},
70+
middlewares: make([]http.Handler, 0),
6971
}
7072
}
7173

@@ -81,7 +83,39 @@ func (h *BasicHTTPFrontendConnector) GetRouter() quesma_api.Router {
8183
return h.router
8284
}
8385

86+
type ResponseWriterWithStatusCode struct {
87+
http.ResponseWriter
88+
statusCode int
89+
}
90+
91+
func (w *ResponseWriterWithStatusCode) WriteHeader(statusCode int) {
92+
w.statusCode = statusCode
93+
w.ResponseWriter.WriteHeader(statusCode)
94+
}
95+
8496
func (h *BasicHTTPFrontendConnector) ServeHTTP(w http.ResponseWriter, req *http.Request) {
97+
index := 0
98+
var runMiddleware func()
99+
100+
runMiddleware = func() {
101+
if index < len(h.middlewares) {
102+
middleware := h.middlewares[index]
103+
index++
104+
responseWriter := &ResponseWriterWithStatusCode{w, 0}
105+
middleware.ServeHTTP(responseWriter, req) // Automatically proceeds to the next middleware
106+
// Only if the middleware did not set a status code, we proceed to the next middleware
107+
if responseWriter.statusCode == 0 {
108+
runMiddleware()
109+
}
110+
111+
} else {
112+
h.finalHandler(w, req)
113+
}
114+
}
115+
runMiddleware()
116+
}
117+
118+
func (h *BasicHTTPFrontendConnector) finalHandler(w http.ResponseWriter, req *http.Request) {
85119
reqBody, err := PeekBodyV2(req)
86120
if err != nil {
87121
http.Error(w, "Error reading request body", http.StatusInternalServerError)
@@ -144,3 +178,7 @@ func ReadRequestBody(request *http.Request) ([]byte, error) {
144178
func (h *BasicHTTPFrontendConnector) GetRouterInstance() *RouterV2 {
145179
return h.routerInstance
146180
}
181+
182+
func (h *BasicHTTPFrontendConnector) AddMiddleware(middleware http.Handler) {
183+
h.middlewares = append(h.middlewares, middleware)
184+
}

‎quesma/main_test.go‎

Lines changed: 79 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -135,13 +135,13 @@ func fallbackScenario() quesma_api.QuesmaBuilder {
135135
var ingestPipeline quesma_api.PipelineBuilder = quesma_api.NewPipeline()
136136
ingestPipeline.AddFrontendConnector(ingestFrontendConnector)
137137
quesmaBuilder.AddPipeline(ingestPipeline)
138-
quesma, _ := quesmaBuilder.Build()
139-
quesma.Start()
140-
return quesma
138+
139+
return quesmaBuilder
141140
}
142141

143142
func Test_fallbackScenario(t *testing.T) {
144-
q1 := fallbackScenario()
143+
qBuilder := fallbackScenario()
144+
q1, _ := qBuilder.Build()
145145
q1.Start()
146146
stop := make(chan os.Signal, 1)
147147
emitRequests(stop)
@@ -159,3 +159,78 @@ func Test_scenario1(t *testing.T) {
159159
<-stop
160160
q1.Stop(context.Background())
161161
}
162+
163+
var middlewareCallCount int32 = 0
164+
165+
type Middleware struct {
166+
emitError bool
167+
}
168+
169+
func (m *Middleware) ServeHTTP(w http.ResponseWriter, r *http.Request) {
170+
atomic.AddInt32(&middlewareCallCount, 1)
171+
if m.emitError {
172+
http.Error(w, "middleware", http.StatusInternalServerError)
173+
}
174+
}
175+
176+
type Middleware2 struct {
177+
}
178+
179+
func (m *Middleware2) ServeHTTP(w http.ResponseWriter, r *http.Request) {
180+
atomic.AddInt32(&middlewareCallCount, 1)
181+
w.WriteHeader(200)
182+
}
183+
184+
func createMiddleWareScenario(emitError bool, cfg *config.QuesmaConfiguration) quesma_api.QuesmaBuilder {
185+
var quesmaBuilder quesma_api.QuesmaBuilder = quesma_api.NewQuesma(quesma_api.EmptyDependencies())
186+
187+
frontendConnector := frontend_connectors.NewBasicHTTPFrontendConnector(":8888", cfg)
188+
HTTPRouter := quesma_api.NewPathRouter()
189+
var fallback quesma_api.HTTPFrontendHandler = fallback
190+
HTTPRouter.AddFallbackHandler(fallback)
191+
frontendConnector.AddRouter(HTTPRouter)
192+
frontendConnector.AddMiddleware(&Middleware{emitError: emitError})
193+
frontendConnector.AddMiddleware(&Middleware2{})
194+
195+
var pipeline quesma_api.PipelineBuilder = quesma_api.NewPipeline()
196+
pipeline.AddFrontendConnector(frontendConnector)
197+
var ingestProcessor quesma_api.Processor = NewIngestProcessor()
198+
pipeline.AddProcessor(ingestProcessor)
199+
quesmaBuilder.AddPipeline(pipeline)
200+
return quesmaBuilder
201+
}
202+
203+
func Test_middleware(t *testing.T) {
204+
205+
cfg := &config.QuesmaConfiguration{
206+
DisableAuth: true,
207+
Elasticsearch: config.ElasticsearchConfiguration{
208+
Url: &config.Url{Host: "localhost:9200", Scheme: "http"},
209+
User: "",
210+
Password: "",
211+
},
212+
}
213+
{
214+
quesmaBuilder := createMiddleWareScenario(true, cfg)
215+
quesmaBuilder.Build()
216+
quesmaBuilder.Start()
217+
stop := make(chan os.Signal, 1)
218+
emitRequests(stop)
219+
<-stop
220+
quesmaBuilder.Stop(context.Background())
221+
atomic.LoadInt32(&middlewareCallCount)
222+
assert.Equal(t, int32(4), middlewareCallCount)
223+
}
224+
atomic.StoreInt32(&middlewareCallCount, 0)
225+
{
226+
quesmaBuilder := createMiddleWareScenario(false, cfg)
227+
quesmaBuilder.Build()
228+
quesmaBuilder.Start()
229+
stop := make(chan os.Signal, 1)
230+
emitRequests(stop)
231+
<-stop
232+
quesmaBuilder.Stop(context.Background())
233+
atomic.LoadInt32(&middlewareCallCount)
234+
assert.Equal(t, int32(8), middlewareCallCount)
235+
}
236+
}

‎quesma/quesma/auth_middleware.go‎

Lines changed: 15 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -24,11 +24,12 @@ type authMiddleware struct {
2424
authHeaderCache sync.Map
2525
cacheWipeInterval time.Duration
2626
esClient elasticsearch.SimpleClient
27+
v2 bool
2728
}
2829

2930
func NewAuthMiddleware(next http.Handler, esConf config.ElasticsearchConfiguration) http.Handler {
3031
esClient := elasticsearch.NewSimpleClient(&esConf)
31-
middleware := &authMiddleware{nextHttpHandler: next, esClient: *esClient, cacheWipeInterval: cacheWipeInterval}
32+
middleware := &authMiddleware{nextHttpHandler: next, esClient: *esClient, cacheWipeInterval: cacheWipeInterval, v2: false}
3233
go middleware.startCacheWipeScheduler()
3334
return middleware
3435
}
@@ -49,7 +50,9 @@ func (a *authMiddleware) ServeHTTP(w http.ResponseWriter, r *http.Request) {
4950
}
5051
if _, ok := a.authHeaderCache.Load(auth); ok {
5152
logger.Debug().Msgf("[AUTH] [%s] called by [%s] - credentials loaded from cache", r.URL, userName)
52-
a.nextHttpHandler.ServeHTTP(w, r)
53+
if !a.v2 {
54+
a.nextHttpHandler.ServeHTTP(w, r)
55+
}
5356
return
5457
}
5558

@@ -61,8 +64,9 @@ func (a *authMiddleware) ServeHTTP(w http.ResponseWriter, r *http.Request) {
6164
http.Error(w, "Unauthorized", http.StatusUnauthorized)
6265
return
6366
}
64-
65-
a.nextHttpHandler.ServeHTTP(w, r)
67+
if !a.v2 {
68+
a.nextHttpHandler.ServeHTTP(w, r)
69+
}
6670
}
6771

6872
func (a *authMiddleware) startCacheWipeScheduler() {
@@ -86,3 +90,10 @@ func (a *authMiddleware) wipeCache() {
8690
return true
8791
})
8892
}
93+
94+
func NewAuthMiddlewareV2(esConf config.ElasticsearchConfiguration) http.Handler {
95+
esClient := elasticsearch.NewSimpleClient(&esConf)
96+
middleware := &authMiddleware{esClient: *esClient, cacheWipeInterval: cacheWipeInterval, v2: true}
97+
go middleware.startCacheWipeScheduler()
98+
return middleware
99+
}

‎quesma/quesma/dual_write_proxy_v2.go‎

Lines changed: 10 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -26,14 +26,12 @@ const concurrentClientsLimitV2 = 100 // FIXME this should be configurable
2626

2727
type simultaneousClientsLimiterV2 struct {
2828
counter atomic.Int64
29-
handler http.Handler
3029
limit int64
3130
}
3231

33-
func newSimultaneousClientsLimiterV2(handler http.Handler, limit int64) *simultaneousClientsLimiterV2 {
32+
func newSimultaneousClientsLimiterV2(limit int64) *simultaneousClientsLimiterV2 {
3433
return &simultaneousClientsLimiterV2{
35-
handler: handler,
36-
limit: limit,
34+
limit: limit,
3735
}
3836
}
3937

@@ -49,7 +47,6 @@ func (c *simultaneousClientsLimiterV2) ServeHTTP(w http.ResponseWriter, r *http.
4947

5048
c.counter.Add(1)
5149
defer c.counter.Add(-1)
52-
c.handler.ServeHTTP(w, r)
5350
}
5451

5552
type dualWriteHttpProxyV2 struct {
@@ -101,12 +98,17 @@ func newDualWriteProxyV2(dependencies quesma_api.Dependencies, schemaLoader clic
10198
if err != nil {
10299
logger.Fatal().Msgf("Error building Quesma: %v", err)
103100
}
104-
105101
var limitedHandler http.Handler
106102
if config.DisableAuth {
107-
limitedHandler = newSimultaneousClientsLimiterV2(elasticHttpIngestFrontendConnector, concurrentClientsLimitV2)
103+
elasticHttpIngestFrontendConnector.AddMiddleware(newSimultaneousClientsLimiterV2(concurrentClientsLimitV2))
104+
elasticHttpQueryFrontendConnector.AddMiddleware(newSimultaneousClientsLimiterV2(concurrentClientsLimitV2))
105+
limitedHandler = elasticHttpIngestFrontendConnector
108106
} else {
109-
limitedHandler = newSimultaneousClientsLimiterV2(NewAuthMiddleware(elasticHttpIngestFrontendConnector, config.Elasticsearch), concurrentClientsLimitV2)
107+
elasticHttpQueryFrontendConnector.AddMiddleware(newSimultaneousClientsLimiterV2(concurrentClientsLimitV2))
108+
elasticHttpQueryFrontendConnector.AddMiddleware(NewAuthMiddlewareV2(config.Elasticsearch))
109+
elasticHttpIngestFrontendConnector.AddMiddleware(newSimultaneousClientsLimiterV2(concurrentClientsLimitV2))
110+
elasticHttpIngestFrontendConnector.AddMiddleware(NewAuthMiddlewareV2(config.Elasticsearch))
111+
limitedHandler = elasticHttpIngestFrontendConnector
110112
}
111113

112114
return &dualWriteHttpProxyV2{

‎quesma/v2/core/quesma_apis.go‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@ package quesma_api
55
import (
66
"context"
77
"net"
8+
"net/http"
89
)
910

1011
type InstanceNamer interface {
@@ -31,8 +32,13 @@ type FrontendConnector interface {
3132

3233
type HTTPFrontendConnector interface {
3334
FrontendConnector
35+
// AddRouter adds a router to the HTTPFrontendConnector
3436
AddRouter(router Router)
3537
GetRouter() Router
38+
// AddMiddleware adds a middleware to the HTTPFrontendConnector.
39+
// The middleware chain is executed in the order it is added
40+
// and before the router is executed.
41+
AddMiddleware(middleware http.Handler)
3642
}
3743

3844
type TCPFrontendConnector interface {

0 commit comments

Comments
 (0)