From ec815638800aebba93ff81908cd6e94f9045e13b Mon Sep 17 00:00:00 2001 From: Chris Hager Date: Thu, 29 May 2025 20:44:42 +0200 Subject: [PATCH 1/2] rpcserver: allow custom /readyz handler --- rpcserver/jsonrpc_server.go | 25 ++++++++++++++++++++++--- 1 file changed, 22 insertions(+), 3 deletions(-) diff --git a/rpcserver/jsonrpc_server.go b/rpcserver/jsonrpc_server.go index 3c31ab4..a06b48c 100644 --- a/rpcserver/jsonrpc_server.go +++ b/rpcserver/jsonrpc_server.go @@ -96,6 +96,9 @@ type JSONRPCHandlerOpts struct { ExtractOriginFromHeader bool // GET response content GetResponseContent []byte + // Custom handler for /readyz endpoint. If not nil then it is expected to write the response to the provided ResponseWriter. + // If the custom handler returns an error, the error message is written to the ResponseWriter with a 500 status code. + ReadyHandler func(w http.ResponseWriter, r *http.Request) error } // NewJSONRPCHandler creates JSONRPC http.Handler from the map that maps method names to method functions @@ -162,9 +165,22 @@ func (h *JSONRPCHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) { stepStartAt := time.Now() - if r.Method != http.MethodPost { - // Respond with GET response content if it's set - if r.Method == http.MethodGet && len(h.GetResponseContent) > 0 { + // Some GET requests are allowed + if r.Method == http.MethodGet { + if r.URL.Path == "/livez" { + w.WriteHeader(http.StatusOK) + return + } else if r.URL.Path == "/readyz" && h.JSONRPCHandlerOpts.ReadyHandler != nil { + err := h.JSONRPCHandlerOpts.ReadyHandler(w, r) + if err == nil { + // Response was already written by the handler + return + } else { + http.Error(w, err.Error(), http.StatusInternalServerError) + incInternalErrors(h.ServerName) + return + } + } else if len(h.GetResponseContent) > 0 { w.WriteHeader(http.StatusOK) _, err := w.Write(h.GetResponseContent) if err != nil { @@ -174,7 +190,10 @@ func (h *JSONRPCHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) { } return } + } + // From here we only accept POST requests with JSON body + if r.Method != http.MethodPost { // Responsd with "only POST method is allowed" http.Error(w, errMethodNotAllowed, http.StatusMethodNotAllowed) incIncorrectRequest(h.ServerName) From 29075cefb2037a679a8c5ea446436b45bdb36dee Mon Sep 17 00:00:00 2001 From: Chris Hager Date: Fri, 30 May 2025 09:03:58 +0200 Subject: [PATCH 2/2] test --- rpcserver/jsonrpc_server.go | 18 ++++++---- rpcserver/jsonrpc_server_test.go | 61 ++++++++++++++++++++++++++++++++ 2 files changed, 72 insertions(+), 7 deletions(-) diff --git a/rpcserver/jsonrpc_server.go b/rpcserver/jsonrpc_server.go index a06b48c..8a862f5 100644 --- a/rpcserver/jsonrpc_server.go +++ b/rpcserver/jsonrpc_server.go @@ -170,17 +170,22 @@ func (h *JSONRPCHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) { if r.URL.Path == "/livez" { w.WriteHeader(http.StatusOK) return - } else if r.URL.Path == "/readyz" && h.JSONRPCHandlerOpts.ReadyHandler != nil { - err := h.JSONRPCHandlerOpts.ReadyHandler(w, r) - if err == nil { - // Response was already written by the handler + } else if r.URL.Path == "/readyz" { + if h.JSONRPCHandlerOpts.ReadyHandler == nil { + http.Error(w, "ready handler is not set", http.StatusNotFound) + incIncorrectRequest(h.ServerName) return } else { - http.Error(w, err.Error(), http.StatusInternalServerError) - incInternalErrors(h.ServerName) + // Handler is expected to write the Response + err := h.JSONRPCHandlerOpts.ReadyHandler(w, r) + if err != nil { + http.Error(w, err.Error(), http.StatusInternalServerError) + incInternalErrors(h.ServerName) + } return } } else if len(h.GetResponseContent) > 0 { + // Static response for all other GET requests w.WriteHeader(http.StatusOK) _, err := w.Write(h.GetResponseContent) if err != nil { @@ -194,7 +199,6 @@ func (h *JSONRPCHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) { // From here we only accept POST requests with JSON body if r.Method != http.MethodPost { - // Responsd with "only POST method is allowed" http.Error(w, errMethodNotAllowed, http.StatusMethodNotAllowed) incIncorrectRequest(h.ServerName) return diff --git a/rpcserver/jsonrpc_server_test.go b/rpcserver/jsonrpc_server_test.go index 7f96bc9..4bb1450 100644 --- a/rpcserver/jsonrpc_server_test.go +++ b/rpcserver/jsonrpc_server_test.go @@ -4,6 +4,7 @@ import ( "bytes" "context" "errors" + "fmt" "net/http" "net/http/httptest" "testing" @@ -120,3 +121,63 @@ func TestJSONRPCServerWithSignatureWithClient(t *testing.T) { require.NoError(t, err) require.Equal(t, 123, structResp.Field) } + +func TestJSONRPCServerDefaultLiveAndReady(t *testing.T) { + handler := testHandler(JSONRPCHandlerOpts{}) + httpServer := httptest.NewServer(handler) + defer httpServer.Close() + + // /livez (200 by default) + request, err := http.NewRequest(http.MethodGet, "/livez", nil) + require.NoError(t, err) + rr := httptest.NewRecorder() + handler.ServeHTTP(rr, request) + require.Equal(t, http.StatusOK, rr.Code) + require.Equal(t, "", rr.Body.String()) + + // /readyz (404 by default) + request, err = http.NewRequest(http.MethodGet, "/readyz", nil) + require.NoError(t, err) + rr = httptest.NewRecorder() + handler.ServeHTTP(rr, request) + require.Equal(t, http.StatusNotFound, rr.Code) +} + +func TestJSONRPCServerReadyzOK(t *testing.T) { + handler := testHandler(JSONRPCHandlerOpts{ + ReadyHandler: func(w http.ResponseWriter, r *http.Request) error { + w.WriteHeader(http.StatusOK) + _, err := w.Write([]byte("ready")) + return err + }, + }) + httpServer := httptest.NewServer(handler) + defer httpServer.Close() + + request, err := http.NewRequest(http.MethodGet, "/readyz", nil) + require.NoError(t, err) + rr := httptest.NewRecorder() + + handler.ServeHTTP(rr, request) + require.Equal(t, http.StatusOK, rr.Code) + require.Equal(t, "ready", rr.Body.String()) +} + +func TestJSONRPCServerReadyzError(t *testing.T) { + handler := testHandler(JSONRPCHandlerOpts{ + ReadyHandler: func(w http.ResponseWriter, r *http.Request) error { + return fmt.Errorf("not ready") + }, + }) + httpServer := httptest.NewServer(handler) + defer httpServer.Close() + + request, err := http.NewRequest(http.MethodGet, "/readyz", nil) + require.NoError(t, err) + rr := httptest.NewRecorder() + + handler.ServeHTTP(rr, request) + require.Equal(t, http.StatusInternalServerError, rr.Code) + fmt.Println(rr.Body.String()) + require.Equal(t, "not ready\n", rr.Body.String()) +}