From f1b4807eeaea50a871f5b48ca275137e6440736a Mon Sep 17 00:00:00 2001
From: Jakob Schmid <jakob.schmid@sap.com>
Date: Wed, 12 May 2021 11:42:20 +0200
Subject: [PATCH] Add Router to runtime.Params

  - Allows extending the REST API by creating a custom runtime object
  - Allows overriding existing routes since Router matches routes in order of registration

Signed-off-by: Jakob Schmid <jakob.schmid@sap.com>
---
 runtime/runtime.go    |  7 +++++++
 server/server_test.go | 22 ++++++++++++++++++++++
 2 files changed, 29 insertions(+)

diff --git a/runtime/runtime.go b/runtime/runtime.go
index 7921bc2dfb..0d273f300f 100644
--- a/runtime/runtime.go
+++ b/runtime/runtime.go
@@ -21,6 +21,7 @@ import (
 	"time"
 
 	"github.com/fsnotify/fsnotify"
+	"github.com/gorilla/mux"
 	"github.com/pkg/errors"
 	"github.com/sirupsen/logrus"
 	"go.uber.org/automaxprocs/maxprocs"
@@ -182,6 +183,11 @@ type Params struct {
 	// configured bundles and plugins to be activated/ready before listening for traffic.
 	// A value of 0 or less means no wait is exercised.
 	ReadyTimeout int
+
+	// Router is the router to which handlers for the REST API are added.
+	// Router uses a first-matching-route-wins strategy, so no existing routes are overridden
+	// If it is nil, a new mux.Router will be created
+	Router *mux.Router
 }
 
 // LoggingConfig stores the configuration for OPA's logging behaviour.
@@ -352,6 +358,7 @@ func (rt *Runtime) Serve(ctx context.Context) error {
 	defer rt.Manager.Stop(ctx)
 
 	rt.server = server.New().
+		WithRouter(rt.Params.Router).
 		WithStore(rt.Store).
 		WithManager(rt.Manager).
 		WithCompilerErrorLimit(rt.Params.ErrorLimit).
diff --git a/server/server_test.go b/server/server_test.go
index b79db50041..53a1c31ce3 100644
--- a/server/server_test.go
+++ b/server/server_test.go
@@ -18,6 +18,7 @@ import (
 	"testing"
 	"time"
 
+	"github.com/gorilla/mux"
 	"github.com/pkg/errors"
 
 	"github.com/open-policy-agent/opa/ast"
@@ -3607,6 +3608,27 @@ func TestMixedAddrTypes(t *testing.T) {
 	}
 }
 
+func TestCustomRoute(t *testing.T) {
+	router := mux.NewRouter()
+	router.HandleFunc("/customEndpoint", func(w http.ResponseWriter, r *http.Request) {
+		w.Write([]byte(`{"myCustomResponse": true}`))
+	})
+	f := newFixture(t, func(server *Server) {
+		server.WithRouter(router)
+	})
+
+	if err := f.v1(http.MethodGet, "/data", "", 200, `{"result":{}}`); err != nil {
+		t.Fatalf("Unexpected response for default server route: %v", err)
+	}
+	r, err := http.NewRequest(http.MethodGet, "/customEndpoint", nil)
+	if err != nil {
+		t.Fatalf("Unexpected error: %s", err)
+	}
+	if err := f.executeRequest(r, http.StatusOK, `{"myCustomResponse": true}`); err != nil {
+		t.Fatalf("Request to custom endpoint failed: %s", err)
+	}
+}
+
 func TestDiagnosticRoutes(t *testing.T) {
 	cases := []struct {
 		path      string